pwact 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

pwact/active_learning/label/labeling.py

@@ -35,6 +35,7 @@ from pwact.utils.file_operation import write_to_file, copy_file, copy_dir, searc
 from pwact.utils.app_lib.common import link_pseudo_by_atom, set_input_script
 
 from pwact.data_format.configop import extract_pwdata, save_config, get_atom_type
+from pwdata import Config
 class Labeling(object):
     @staticmethod
     def kill_job(root_dir:str, itername:str):
@@ -59,9 +60,10 @@ class Labeling(object):
         self.real_explore_dir = os.path.join(self.input_param.root_dir, itername, AL_STRUCTURE.explore)
         self.md_dir = os.path.join(self.explore_dir, EXPLORE_FILE_STRUCTURE.md)
         self.select_dir = os.path.join(self.explore_dir, EXPLORE_FILE_STRUCTURE.select)
+        self.direct_dir = os.path.join(self.explore_dir, EXPLORE_FILE_STRUCTURE.direct) 
         self.real_md_dir = os.path.join(self.real_explore_dir, EXPLORE_FILE_STRUCTURE.md)
         self.real_select_dir = os.path.join(self.real_explore_dir, EXPLORE_FILE_STRUCTURE.select)
-
+        self.real_direct_dir = os.path.join(self.real_explore_dir, EXPLORE_FILE_STRUCTURE.direct)
         # labed work dir
         self.label_dir = os.path.join(self.input_param.root_dir, itername, TEMP_STRUCTURE.tmp_run_iter_dir, AL_STRUCTURE.labeling) 
         self.scf_dir = os.path.join(self.label_dir, LABEL_FILE_STRUCTURE.scf)
@@ -71,6 +73,9 @@ class Labeling(object):
         self.real_scf_dir = os.path.join(self.real_label_dir, LABEL_FILE_STRUCTURE.scf)
         self.real_result_dir = os.path.join(self.real_label_dir, LABEL_FILE_STRUCTURE.result)
 
+        self.bigmodel_dir = os.path.join(self.label_dir, LABEL_FILE_STRUCTURE.bigmodel)
+        self.real_bigmodel_dir = os.path.join(self.real_label_dir, LABEL_FILE_STRUCTURE.bigmodel)
+
     '''
     description: 
     the scf work dir file structure is as follow.
@@ -86,9 +91,8 @@ class Labeling(object):
     return {*}
     author: wuxingxing
     '''        
+
     def make_scf_work(self):
-        # read select info, and make scf
-        # ["devi_force", "file_path", "config_index"]
         candidate = pd.read_csv(os.path.join(self.select_dir, EXPLORE_FILE_STRUCTURE.candidate))
         # make scf work dir
         scf_dir_list = []
@@ -108,14 +112,51 @@ class Labeling(object):
                 atom_names = line.split()
             self.make_scf_file(scf_sub_md_sys_path, tarj_lmp, atom_names)
             scf_dir_list.append(scf_sub_md_sys_path)
-            
+        
         self.make_scf_slurm_job_files(scf_dir_list)
 
+    def make_bigmodel_work(self):
+        # copy from realdir/direct/select.xyz
+        if self.input_param.strategy.direct:
+            copy_file(os.path.join(self.real_direct_dir, EXPLORE_FILE_STRUCTURE.select_xyz), 
+                os.path.join(self.bigmodel_dir, EXPLORE_FILE_STRUCTURE.select_xyz))
+        else:
+            # copy trajs to bigmodel_dir and cvt to xyz
+            candidate = pd.read_csv(os.path.join(self.select_dir, EXPLORE_FILE_STRUCTURE.candidate))
+            # make scf work dir
+            image_list = None
+            for index, row in candidate.iterrows():
+                config_index    = int(row["config_index"])
+                sub_md_sys_path = row["file_path"]
+                atom_names = None
+                with open(os.path.join(sub_md_sys_path, LAMMPS.atom_type_file), 'r') as rf:
+                    line = rf.readline()
+                    atom_names = line.split()
+                if image_list is None:
+                    image_list = Config(data_path=os.path.join(sub_md_sys_path, EXPLORE_FILE_STRUCTURE.traj, "{}{}".format(config_index, LAMMPS.traj_postfix)), 
+                                        format=PWDATA.lammps_dump, atom_names=atom_names)
+                else:
+                    image_list.append(Config(data_path=os.path.join(sub_md_sys_path, EXPLORE_FILE_STRUCTURE.traj, "{}{}".format(config_index, LAMMPS.traj_postfix)), 
+                                        format=PWDATA.lammps_dump, atom_names=atom_names))
+            # cvt_lammps.dump to extxyz
+            image_list.to(data_path=self.bigmodel_dir, format=PWDATA.extxyz, data_name="{}".format(EXPLORE_FILE_STRUCTURE.select_xyz))
+        # copy bigmodelscript
+        copy_file(self.input_param.scf.bigmodel_script, os.path.join(self.bigmodel_dir, os.path.basename(self.input_param.scf.bigmodel_script)))
+        # make slrum file
+        self.make_bigmodel_slurm_job_files([self.bigmodel_dir])
+
     def back_label(self):
-        slurm_remain, slurm_success = get_slurm_job_run_info(self.real_scf_dir, \
-            job_patten="*-{}".format(LABEL_FILE_STRUCTURE.scf_job), \
-            tag_patten="*-{}".format(LABEL_FILE_STRUCTURE.scf_tag))
-        slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False
+        if self.input_param.scf.dft_style == DFT_STYLE.bigmodel:
+            slurm_remain, slurm_success = get_slurm_job_run_info(self.real_bigmodel_dir, \
+                job_patten="*-{}".format(LABEL_FILE_STRUCTURE.bigmodel_job), \
+                tag_patten="*-{}".format(LABEL_FILE_STRUCTURE.bigmodel_tag))
+            slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False        
+        else:
+            slurm_remain, slurm_success = get_slurm_job_run_info(self.real_scf_dir, \
+                job_patten="*-{}".format(LABEL_FILE_STRUCTURE.scf_job), \
+                tag_patten="*-{}".format(LABEL_FILE_STRUCTURE.scf_tag))
+            slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False
+        
         if slurm_done:
             # bk and do new job
             target_bk_file = add_postfix_dir(self.real_label_dir, postfix_str="bk")
@@ -147,7 +188,31 @@ class Labeling(object):
                 mission.commit_jobs()
                 mission.check_running_job()
                 mission.all_job_finished(error_type=SLURM_OUT.dft_out)
-                    
+
+    def do_bigmodel_jobs(self):
+        mission = Mission()
+        slurm_remain, slurm_success = get_slurm_job_run_info(self.bigmodel_dir, \
+            job_patten="*-{}".format(LABEL_FILE_STRUCTURE.bigmodel_job), \
+            tag_patten="*-{}".format(LABEL_FILE_STRUCTURE.bigmodel_tag))
+        slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False
+        if slurm_done is False:
+            #recover slurm jobs
+            if len(slurm_remain) > 0:
+                print("Run bigModel Job:\n")
+                print(slurm_remain)
+                for i, script_path in enumerate(slurm_remain):
+                    slurm_job = SlurmJob()
+                    tag_name = "{}-{}".format(os.path.basename(script_path).split('-')[0].strip(), LABEL_FILE_STRUCTURE.bigmodel_tag)
+                    tag = os.path.join(os.path.dirname(script_path),tag_name)
+                    slurm_job.set_tag(tag)
+                    slurm_job.set_cmd(script_path)
+                    mission.add_job(slurm_job)
+
+            if len(mission.job_list) > 0:
+                mission.commit_jobs()
+                mission.check_running_job()
+                mission.all_job_finished()
+
     def make_scf_file(self, scf_dir:str, tarj_lmp:str, atom_names:list[str]=None):
         config_index = os.path.basename(tarj_lmp).split('.')[0]
         if DFT_STYLE.vasp == self.resource.dft_style: # when do scf, the vasp input file name is 'POSCAR'
@@ -230,6 +295,42 @@ class Labeling(object):
             slurm_job_file = os.path.join(self.scf_dir, slurm_script_name)
             write_to_file(slurm_job_file, group_slurm_script, "w")
 
+
+    def make_bigmodel_slurm_job_files(self, scf_sub_list:list[str]):
+        del_file_list_by_patten(self.bigmodel_dir, "*{}".format(LABEL_FILE_STRUCTURE.scf_job))
+        group_list = split_job_for_group(1, scf_sub_list, 1)
+        
+        for group_index, group in enumerate(group_list):
+            if group[0] == "NONE":
+                continue
+            
+            jobname = "bigmodel{}".format(group_index)
+            tag_name = "{}-{}".format(group_index, LABEL_FILE_STRUCTURE.bigmodel_tag)
+            tag = os.path.join(self.bigmodel_dir, tag_name)
+            run_cmd = self.resource.dft_resource.command
+            # if self.resource.dft_resource.gpu_per_node > 0:
+            #     run_cmd = "mpirun -np {} PWmat > {}".format(self.resource.dft_resource.gpu_per_node, SLURM_OUT.md_out)
+            # else:
+            #     raise Exception("ERROR! the cpu version of pwmat not support yet!")
+            group_slurm_script = set_slurm_script_content(gpu_per_node=self.resource.dft_resource.gpu_per_node, 
+                number_node = self.resource.dft_resource.number_node, 
+                cpu_per_node = self.resource.dft_resource.cpu_per_node,
+                queue_name = self.resource.dft_resource.queue_name,
+                custom_flags = self.resource.dft_resource.custom_flags,
+                env_script = self.resource.dft_resource.env_script,
+                job_name = jobname,
+                run_cmd_template = run_cmd,
+                group = group,
+                job_tag = tag,
+                task_tag = LABEL_FILE_STRUCTURE.bigmodel_tag, 
+                task_tag_faild = LABEL_FILE_STRUCTURE.bigmodel_tag_failed,
+                parallel_num=self.resource.dft_resource.parallel_num,
+                check_type=self.resource.dft_style
+                )
+            slurm_script_name = "{}-{}".format(group_index, LABEL_FILE_STRUCTURE.bigmodel_job)
+            slurm_job_file = os.path.join(self.bigmodel_dir, slurm_script_name)
+            write_to_file(slurm_job_file, group_slurm_script, "w")
+
     '''
     description: 
     collecte OUT.MLMD to mvm-
@@ -274,18 +375,31 @@ class Labeling(object):
                 for scf_file in scf_files:
                     scf_file_path = os.path.join(scf_dir, scf_file)
                     if scf_file.lower() in DFT_STYLE.get_scf_reserve_list(self.resource.dft_style) \
-                        and scf_file.lower() not in DFT_STYLE.get_scf_del_list():# for pwmat final.config
+                        or "atom.config" in scf_file.lower() :# for the input natom.config
                         copy_file(scf_file_path, scf_file_path.replace(TEMP_STRUCTURE.tmp_run_iter_dir, ""))
 
         # scf files to pwdata format
         scf_configs = self.collect_scf_configs()
+        if len(scf_configs) > 0:
+            extract_pwdata(input_data_list=scf_configs,
+                    intput_data_format =DFT_STYLE.get_format_by_postfix(os.path.basename(scf_configs[0])),
+                    save_data_path =self.result_dir,
+                    save_data_name = INIT_BULK.get_save_format(self.input_param.data_format),
+                    save_data_format = self.input_param.data_format,
+                    data_shuffle     =self.input_param.train.data_shuffle
+            )
+            # copy to main dir
+            copy_dir(self.result_dir, self.real_result_dir)
 
-        extract_pwdata(input_data_list=scf_configs,
-                intput_data_format =DFT_STYLE.get_format_by_postfix(os.path.basename(scf_configs[0])),
-                save_data_path =self.result_dir,
-                save_data_name = INIT_BULK.get_save_format(self.input_param.data_format),
-                save_data_format = self.input_param.data_format,
-                data_shuffle     =self.input_param.train.data_shuffle
-        )
-        # copy to main dir
+    def do_post_bigmodel(self):
+        # copy the bigmodel labeled.xyz to result
+        if self.input_param.data_format == PWDATA.extxyz:
+            copy_file(os.path.join(self.bigmodel_dir, LABEL_FILE_STRUCTURE.train_xyz), os.path.join(self.result_dir, LABEL_FILE_STRUCTURE.train_xyz))
+        else:
+            images = Config(data_path=os.path.join(self.bigmodel_dir, LABEL_FILE_STRUCTURE.train_xyz), format=PWDATA.extxyz)
+            images.to(data_path=self.result_dir, format=PWDATA.pwmlff_npy)
+        # copy bigmodel dir to real dir
+        copy_dir(self.bigmodel_dir, self.real_bigmodel_dir)
         copy_dir(self.result_dir, self.real_result_dir)
+        # del slurm logs and tags
+        del_file_list_by_patten(self.real_bigmodel_dir, "slurm-*")

pwact/active_learning/train/train_model.py

@@ -33,6 +33,7 @@ class ModelTrian(object):
         self.itername = itername
         self.resource = resource
         self.input_param = input_param
+        self.train_flag = True
         self.iter = get_iter_from_iter_name(self.itername)
         # train work dir
         self.train_dir = os.path.join(self.input_param.root_dir, self.itername, TEMP_STRUCTURE.tmp_run_iter_dir, AL_STRUCTURE.train)
@@ -68,11 +69,16 @@ class ModelTrian(object):
             if not os.path.exists(model_i_dir):
                 os.makedirs(model_i_dir)
             # make train.json file
-            train_dict = self.set_train_input_dict(work_dir=model_i_dir, model_index = model_index)
+            train_dict, train_tag = self.set_train_input_dict(work_dir=model_i_dir, model_index = model_index)
             train_json_file_path = os.path.join(model_i_dir, TRAIN_FILE_STRUCTUR.train_json)
             save_json_file(train_dict, train_json_file_path)
             train_list.append(model_i_dir)
-        self.make_train_slurm_job_files(train_list)
+        if train_tag:
+            self.make_train_slurm_job_files(train_list)
+        else:
+            pre_iter_name = make_iter_name(self.iter - 1)
+            pre_iter_dir = os.path.join(self.input_param.root_dir, pre_iter_name, AL_STRUCTURE.train)
+            copy_dir(pre_iter_dir, self.train_dir)
     
     def make_train_slurm_job_files(self, train_list:list[str]):
         # make train slurm script
@@ -158,6 +164,7 @@ class ModelTrian(object):
         # search train_feature_path in iter*/label/result/*/PWdata/*
         iter_index = get_iter_from_iter_name(self.itername)
         start_iter = 0
+        train_tag = True
         while start_iter < iter_index:
             if self.input_param.data_format == PWDATA.extxyz: # result/train.xyz
                 iter_data_list = search_files(self.input_param.root_dir, 
@@ -172,6 +179,9 @@ class ModelTrian(object):
                 train_feature_path.extend(iter_data_list)
             start_iter += 1
         
+        if start_iter > 0 and len(iter_data_list) == 0:
+            train_tag = False
+
         # reset seed
         train_json[TRAIN_INPUT_PARAM.seed] = get_seed_by_time()
         train_json[TRAIN_INPUT_PARAM.raw_files] = []
@@ -181,7 +191,7 @@ class ModelTrian(object):
         train_json[TRAIN_INPUT_PARAM.format] = self.input_param.data_format
         if self.input_param.strategy.uncertainty == UNCERTAINTY.kpu:
             train_json[TRAIN_INPUT_PARAM.save_p_matrix] = True
-        return train_json
+        return train_json, train_tag
 
     def do_train_job(self):
         mission = Mission()

pwact/active_learning/user_input/init_bulk_input.py

@@ -26,13 +26,15 @@ class InitBulkParam(object):
             sys_configs = [sys_configs]
 
         # set sys_config detail
-        self.dft_style = get_required_parameter("dft_style", json_dict).lower()
+        self.dft_style = get_parameter("dft_style", json_dict, "PWMAT").lower()
         self.scf_style = get_parameter("scf_style", json_dict, None)
 
         self.sys_config:list[Stage] = []
         self.is_relax = False
         self.is_aimd = False
         self.is_scf = False
+        self.is_bigmodel=False
+        self.is_direct = False
         for index, config in enumerate(sys_configs):
             stage = Stage(config, index, sys_config_prefix, self.dft_style)
             self.sys_config.append(stage)
@@ -42,22 +44,46 @@ class InitBulkParam(object):
                 self.is_aimd = True
             if stage.scf:
                 self.is_scf = True
+            if stage.bigmodel:
+                self.is_bigmodel = True
+            if stage.direct:
+                self.is_direct = True
                 
         # for PWmat: set etot.input files and persudo files
         # for Vasp: set INCAR files and persudo files
-        self.dft_input = SCFParam(json_dict=json_dict, is_scf=self.is_scf, is_relax=self.is_relax, is_aimd=self.is_aimd, root_dir=self.root_dir, dft_style=self.dft_style, scf_style=self.scf_style)
+        self.dft_input = SCFParam(json_dict=json_dict, 
+                                    is_scf=self.is_scf, 
+                                    is_relax=self.is_relax, 
+                                    is_aimd=self.is_aimd, 
+                                    root_dir=self.root_dir, 
+                                    dft_style=self.dft_style, 
+                                    scf_style=self.scf_style,
+                                    is_bigmodel=self.is_bigmodel,
+                                    is_direct=self.is_direct)
+        
         # check and set relax etot.input file
         for config in self.sys_config:
             if self.is_relax:
                 if config.relax_input_idx >= len(self.dft_input.relax_input_list):
                     raise Exception("Error! for config '{}' 'relax_input_idx' {} not in 'relax_input'!".format(os.path.basename(config.config_file), config.relax_input_idx))
                 config.set_relax_input_file(self.dft_input.relax_input_list[config.relax_input_idx])
+            
             if self.is_scf:
                 if not os.path.exists(self.dft_input.scf_input_list[0].input_file):
                     raise Exception("Error! relabel dft input file {} not exisit!".format(self.dft_input.scf_input_list[0].input_file))
                 config.set_scf_input_file(self.dft_input.scf_input_list[0])
-        # check and set aimd etot.input file
-        for config in self.sys_config:
+            
+            if self.is_bigmodel:
+                if config.bigmodel_input_idx >= len(self.dft_input.bigmodel_input_list):
+                    raise Exception("Error! for script '{}' 'bigmodel_input_idx' {} not in 'bigmodel_input'!".format(os.path.basename(config.config_file), config.bigmodel_input_idx))
+                config.set_bigmodel_input_file(self.dft_input.bigmodel_input_list[config.bigmodel_input_idx])
+
+            if self.is_direct:
+                if config.direct_input_idx >= len(self.dft_input.direct_input_list):
+                    raise Exception("Error! for script '{}' 'direct_input_idx' {} not in 'direct_input'!".format(os.path.basename(config.config_file), config.direct_input_idx))
+                config.set_direct_input_file(self.dft_input.direct_input_list[config.direct_input_idx])
+
+            # check and set aimd etot.input file
             if self.is_aimd:
                 if config.aimd_input_idx >= len(self.dft_input.aimd_input_list):
                     raise Exception("Error! for config '{}' 'aimd_input_idx' {} not in 'aimd_input'!".format(os.path.basename(config.config_file), config.aimd_input_idx))
@@ -77,16 +103,29 @@ class Stage(object):
         self.format = get_parameter("format", json_dict, PWDATA.pwmat_config).lower()
         self.pbc = get_parameter("pbc", json_dict, [1,1,1])
         # extract config file to Config object, then use it
-        self.relax = get_parameter("relax", json_dict, True)
+        self.relax = get_parameter("relax", json_dict, False)
         self.relax_input_idx = get_parameter("relax_input_idx", json_dict, 0)
         self.relax_input_file = None
         
-        self.aimd = get_parameter("aimd", json_dict, True)
+        self.aimd = get_parameter("aimd", json_dict, False)
         self.aimd_input_idx = get_parameter("aimd_input_idx", json_dict, 0)
         self.aimd_input_file = None
         
         self.scf = get_parameter("scf", json_dict, False)
+        self.scf_input_idx = get_parameter("scf_input_idx", json_dict, 0)
+        self.scf_input_file = None
 
+        self.bigmodel = get_parameter("bigmodel", json_dict, False)
+        self.bigmodel_input_idx = get_parameter("bigmodel_input_idx", json_dict, 0)
+        self.bigmodel_script = None
+
+        self.direct = get_parameter("direct", json_dict, False)
+        self.direct_input_idx = get_parameter("direct_input_idx", json_dict, 0)
+        self.direct_script = None
+        
+        if self.bigmodel and self.aimd:
+            raise Exception("ERROR! The 'aimd' and 'bigmodel' cannot be set simultaneously!")
+        
         super_cell = get_parameter("super_cell", json_dict, [])
         super_cell = str_list_format(super_cell)
         if len(super_cell) > 0:
@@ -131,3 +170,13 @@ class Stage(object):
         self.aimd_flag_symm = input_file.flag_symm
         self.use_dftb = input_file.use_dftb
         self.use_skf = input_file.use_skf
+
+    def set_bigmodel_input_file(self, input_file:DFTInput):
+        self.bigmodel_input_file = input_file.input_file
+        self.bigmodel_kspacing = input_file.kspacing 
+        self.bigmodel_flag_symm = input_file.flag_symm
+
+    def set_direct_input_file(self, input_file:DFTInput):
+        self.direct_input_file = input_file.input_file
+        self.direct_kspacing = input_file.kspacing 
+        self.direct_flag_symm = input_file.flag_symm

pwact/active_learning/user_input/iter_input.py

@@ -105,6 +105,18 @@ class StrategyParam(object):
             if self.compress:
                 error_log = "Error! the kpu uncertainty does not support compress, please set the 'compress' in strategy dict to be false!"
                 raise Exception(error_log)
+        
+        self.direct = get_parameter("direct", json_dict, False)
+        if self.direct:
+            self.direct_script = get_parameter("direct_script", json_dict, None)
+            if self.direct_script is not None:
+                self.direct_script = os.path.abspath(self.direct_script)
+                if not os.path.exists(self.direct_script):
+                    raise Exception("ERROR! The direct script {} does not exist!".format(self.direct_script))
+            else:
+                raise Exception("ERROR! The direct script does not exist!")
+        else:
+            self.direct_script = None
 
     def to_dict(self):
         res = {}

pwact/active_learning/user_input/resource.py

@@ -20,10 +20,22 @@ class Resource(object):
             if "-in" in self.explore_resource.command:
                 self.explore_resource.command = self.explore_resource.command.split('-in')[0].strip()
             self.explore_resource.command = "{} -in {} > {}".format(self.explore_resource.command, LAMMPS.input_lammps, SLURM_OUT.md_out)
-
+        else:
+            if "explore" in json_dict.keys():
+                self.explore_resource = self.get_resource(get_required_parameter("explore", json_dict))
+            else:
+                self.explore_resource = None
         # check dft resource
-        self.dft_resource = self.get_resource(get_required_parameter("dft", json_dict))
-        
+        if "dft" in json_dict.keys():
+            self.dft_resource = self.get_resource(get_required_parameter("dft", json_dict))
+        else:
+            self.dft_resource = ResourceDetail("mpirun -np 1 PWmat", 1, 1, 1, 1, 1, None, None, None)
+
+        if "direct" in json_dict.keys():
+            self.direct_resource = self.get_resource(get_required_parameter("direct", json_dict))
+        else:
+            self.direct_resource = None
+
         if "scf" in json_dict.keys():
             self.scf_resource = self.get_resource(get_parameter("scf", json_dict, None))
         else:
@@ -33,11 +45,11 @@ class Resource(object):
         #     self.dft_resource.dftb_command  = "{} > {}".format(dftb_command, SLURM_OUT.dft_out)
         self.dft_style = dft_style
         self.scf_style = scf_style
-        if DFT_STYLE.vasp.lower() == dft_style.lower():
+        if DFT_STYLE.vasp.lower() == dft_style:
             self.dft_resource.command = "{} > {}".format(self.dft_resource.command, SLURM_OUT.dft_out)
-        elif DFT_STYLE.pwmat.lower() == dft_style.lower():
+        elif DFT_STYLE.pwmat.lower() == dft_style:
             self.dft_resource.command = "{} > {}".format(self.dft_resource.command, SLURM_OUT.dft_out)
-        elif DFT_STYLE.cp2k.lower() == dft_style.lower():
+        elif DFT_STYLE.cp2k.lower() == dft_style:
             self.dft_resource.command = "{} {} > {}".format(self.dft_resource.command, CP2K.cp2k_inp, SLURM_OUT.dft_out)
         
         if self.scf_resource is not None and scf_style is not None:
@@ -114,7 +126,7 @@ class ResourceDetail(object):
         if self.gpu_per_node is None and self.cpu_per_node is None:
             raise Exception("ERROR! Both CPU and GPU resources are not specified!")
         # check param
-        if "$SLURM_NTASKS".lower() in command.lower():
+        if "$SLURM".lower() in command.lower():
             pass
         else:
             if "mpirun -np" in command:

pwact/active_learning/user_input/scf_param.py

@@ -10,7 +10,9 @@ class SCFParam(object):
         is_scf:bool=False, 
         root_dir:str=None, 
         dft_style:str=None,
-        scf_style:str=None) -> None:# for scf relabel in init_bulk
+        scf_style:str=None,
+        is_bigmodel:bool=False,
+        is_direct:bool=False) -> None:# for scf relabel in init_bulk
         
         self.dft_style = dft_style
         self.root_dir = root_dir
@@ -24,12 +26,18 @@ class SCFParam(object):
 
         if is_scf:
             if "scf_input" in json_dict.keys(): # for init_bulk relabel
-                json_scf = get_required_parameter("scf_input", json_dict)
-                self.scf_input_list = self.set_input(json_scf, flag_symm=0)
+                if dft_style == DFT_STYLE.bigmodel:
+                    self.bigmodel_script = get_required_parameter("bigmodel_script", json_dict)
+                else:
+                    json_scf = get_required_parameter("scf_input", json_dict)
+                    self.scf_input_list = self.set_input(json_scf, flag_symm=0)
             else: # for run_iter
-                self.scf_input_list = self.set_input(json_dict, flag_symm=0)
-                if self.scf_input_list[0].use_dftb:
-                    self.use_dftb = True
+                if dft_style == DFT_STYLE.bigmodel:
+                    self.bigmodel_script = get_required_parameter("bigmodel_script", json_dict)
+                else:
+                    self.scf_input_list = self.set_input(json_dict, flag_symm=0)
+                    if self.scf_input_list[0].use_dftb:
+                        self.use_dftb = True
         if is_aimd:
             json_aimd = get_required_parameter("aimd_input", json_dict)
             self.aimd_input_list = self.set_input(json_aimd, flag_symm=0)
@@ -40,6 +48,16 @@ class SCFParam(object):
             self.relax_input_list = self.set_input(json_relax, flag_symm=3) 
             if self.relax_input_list[0].use_dftb:
                 self.use_dftb = True
+
+        if is_bigmodel: # init_bulk
+            json_bigmodel = get_required_parameter("bigmodel_input", json_dict)
+            self.bigmodel_input_list = self.set_input(json_bigmodel, flag_symm=3) 
+
+        if is_direct: # init_bulk
+            json_direct = get_required_parameter("direct_input", json_dict)
+            self.direct_input_list = self.set_input(json_direct, flag_symm=3) 
+
+        self.scf_max_num = get_parameter("scf_max_num", json_dict, None)
         # for pwmat, use 'pseudo' key
         # for vasp is INCAR file, use 'pseudo' key        
         pseudo = get_parameter("pseudo", json_dict, [])

pwact/active_learning/user_input/train_param/optimizer_param.py

@@ -6,7 +6,7 @@ class OptimizerParam(object):
 
     def set_optimizer(self, json_source:dict, nep_param:NepParam=None):
         optimizer_dict = get_parameter("optimizer", json_source, {})
-        self.opt_name = get_parameter("optimizer", optimizer_dict, "LKF")
+        self.opt_name = get_parameter("optimizer", optimizer_dict, "ADAM")
         self.batch_size = get_parameter("batch_size", optimizer_dict, 1)
         self.epochs = get_parameter("epochs", optimizer_dict, 30)
         self.print_freq = get_parameter("print_freq", optimizer_dict, 10)

pwact/main.py

@@ -5,7 +5,7 @@ import glob
 import sys
 import json
 import argparse
-from pwact.utils.constant import TEMP_STRUCTURE, UNCERTAINTY, AL_WORK, AL_STRUCTURE, LABEL_FILE_STRUCTURE, EXPLORE_FILE_STRUCTURE
+from pwact.utils.constant import TEMP_STRUCTURE, UNCERTAINTY, AL_WORK, AL_STRUCTURE, LABEL_FILE_STRUCTURE, EXPLORE_FILE_STRUCTURE, DFT_STYLE
 from pwact.utils.format_input_output import make_iter_name
 from pwact.utils.file_operation import write_to_file, del_file_list, search_files, del_dir, copy_dir
 from pwact.utils.json_operation import convert_keys_to_lowercase
@@ -79,18 +79,23 @@ def run_iter():
             if jj == 2 and not input_param.reserve_work: # delete temp_work_dir under current iteration after the labeling done
                 del_file_list([os.path.join(input_param.root_dir, iter_name, TEMP_STRUCTURE.tmp_run_iter_dir)])
 
-    print("Active learning done! \nYou could use cmd 'al_pwmlff gather_pwdata' to collect all datas sampled from iterations.")
+    print("Active learning done! \nYou could use cmd 'pwact gather_pwdata' to collect all datas sampled from iterations.")
 
 def run_fp(itername:str, resource : Resource, input_param: InputParam):
     lab = Labeling(itername, resource, input_param)
     #1. if the label work done before, back up and do new work
     lab.back_label()
     #2. make scf work
-    lab.make_scf_work()
-    #3. do scf work
-    lab.do_scf_jobs()
-    #4. collect scf configs outcar or movement, then to pwdata format
-    lab.do_post_labeling()
+    if input_param.dft_style == DFT_STYLE.bigmodel:
+        lab.make_bigmodel_work()
+        lab.do_bigmodel_jobs()
+        lab.do_post_bigmodel()
+    else:
+        lab.make_scf_work()
+        #3. do scf work
+        lab.do_scf_jobs()
+        #4. collect scf configs outcar or movement, then to pwdata format
+        lab.do_post_labeling()
     
 def do_training_work(itername:str, resource : Resource, input_param: InputParam):
     mtrain = ModelTrian(itername, resource, input_param)
@@ -129,8 +134,13 @@ def do_exploring_work(itername:str, resource : Resource, input_param: InputParam
     summary = "{}  {}\n".format(itername, summary)
     write_to_file(os.path.join(input_param.root_dir, EXPLORE_FILE_STRUCTURE.iter_select_file), summary, mode='a')
 
+    if input_param.strategy.direct:
+        md.make_drct_work()
+        md.do_drct_jobs()
+        md.post_drct()
     print("config selection done!")
-    # 5. do post process after lammps md running
+
+    # 5. do post process 
     md.post_process_md()
     print("exploring done!")