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

pwact/active_learning/init_bulk/relabel.py

@@ -16,18 +16,27 @@
     
 """
 import os
-
+import glob
+import json
+import bisect
 from pwact.active_learning.user_input.resource import Resource
 from pwact.active_learning.user_input.init_bulk_input import InitBulkParam
 from pwact.active_learning.init_bulk.duplicate_scale import get_config_files_with_order
 
-from pwact.utils.constant import PWMAT, INIT_BULK, TEMP_STRUCTURE, SLURM_OUT, DFT_STYLE
+from pwact.utils.constant import PWMAT, INIT_BULK, TEMP_STRUCTURE, SLURM_OUT, DFT_STYLE, PWDATA, VASP
 from pwact.active_learning.slurm.slurm import SlurmJob, Mission
 from pwact.utils.slurm_script import get_slurm_job_run_info, split_job_for_group, set_slurm_script_content
     
-from pwact.utils.file_operation import write_to_file, link_file, search_files, del_file_list_by_patten
+from pwact.utils.file_operation import write_to_file, link_file, del_dir, del_file_list_by_patten, get_random_nums
 from pwact.utils.app_lib.common import link_pseudo_by_atom, set_input_script
-from pwact.data_format.configop import save_config, get_atom_type, load_config
+from pwact.data_format.configop import extract_pwdata, save_config, get_atom_type, load_config
+
+import pandas as pd
+from pwdata import Config
+
+# from pwact.utils.constant import DFT_TYPE, VASP, PWDATA, AL_STRUCTURE, TEMP_STRUCTURE,\
+#     LABEL_FILE_STRUCTURE, EXPLORE_FILE_STRUCTURE, LAMMPS, SLURM_OUT, DFT_STYLE, PWMAT, INIT_BULK
+# from pwact.utils.file_operation import write_to_file, copy_file, copy_dir, search_files, mv_file, add_postfix_dir, del_dir, del_file_list_by_patten, link_file
 
 class Relabel(object):
     def __init__(self, resource: Resource, input_param:InitBulkParam):
@@ -42,50 +51,104 @@ class Relabel(object):
         
         self.scf_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.scf)
         self.real_scf_dir = os.path.join(self.input_param.root_dir, INIT_BULK.scf)
-       
-    def make_scf_work(self):
-        scf_paths = []
-        use_dftb = False
-        for init_config in self.init_configs:
-            if init_config.scf is False:
-                continue
-            init_config_name = "init_config_{}".format(init_config.config_index)
-            #1. read construtures from aimd dir
-
-            #2. set relabel dir 
-            #       read trajs from ./aimd/init_config_0/relax/0_aimd/
-            #       make scf dir ./relabel/init_config_0/relax/0_aimd/10-scf/files
-            traj_list = search_files(os.path.join(self.aimd_dir, init_config_name), "*/*aimd")
-            for traj_dir in traj_list:
-                scf_dir = os.path.join(self.scf_dir, init_config_name, \
-                    os.path.basename(os.path.dirname(traj_dir)),\
-                    os.path.basename(traj_dir))
 
-                traj_file_name = DFT_STYLE.get_aimd_config(self.resource.dft_style)
-
-                scf_lsit = self.make_scf_file(
-                    scf_dir      =scf_dir, 
-                    traj_file    =os.path.join(traj_dir, traj_file_name), 
-                    traj_format  =DFT_STYLE.get_format_by_postfix(traj_file_name), 
-                    interval     = self.input_param.interval,
-                    target_format=DFT_STYLE.get_pwdata_format(self.input_param.scf_style, is_cp2k_coord=True),
-                    input_file   =init_config.scf_input_file,
-                    kspacing     =init_config.scf_kspacing, 
-                    flag_symm    =init_config.scf_flag_symm,
-                    is_dftb      = False,
-                    in_skf       =None)
+        self.bigmodel_dir = os.path.join(self.input_param.root_dir, TEMP_STRUCTURE.tmp_init_bulk_dir, INIT_BULK.bigmodel)
+        self.real_bigmodel_dir = os.path.join(self.input_param.root_dir, INIT_BULK.bigmodel)
+        
+        self.direct_dir = os.path.join(self.bigmodel_dir, INIT_BULK.direct)
+        self.real_direct_dir = os.path.join(self.real_bigmodel_dir, INIT_BULK.direct)
 
-                scf_paths.extend(scf_lsit)
-        # make slurm script and slurm job
-        self.make_scf_slurm_job_files(scf_paths, use_dftb)
-    
     def check_work_done(self):
         slurm_remain, slurm_success = get_slurm_job_run_info(self.scf_dir, \
             job_patten="*-{}".format(INIT_BULK.scf_job), \
             tag_patten="*-{}".format(INIT_BULK.scf_tag))
         slurm_done = True if len(slurm_remain) == 0 and len(slurm_success) > 0 else False # len(slurm_remain) > 0 exist slurm jobs need to do
         return slurm_done
-            
+
+    def make_scf_work(self):
+        def find_position_binary(prefix_sum, N):
+            idx = bisect.bisect_right(prefix_sum, N)
+            if idx == 0:
+                return 0
+            elif idx >= len(prefix_sum):
+                return (len(prefix_sum)-1)
+            else:
+                return idx
+
+        def compute_prefix_sum(arr):
+            prefix_sum = []
+            current_sum = 0
+            for num in arr:
+                current_sum += num
+                prefix_sum.append(current_sum)
+            return prefix_sum
+                            
+        candidate = Config(data_path=os.path.join(self.direct_dir, INIT_BULK.direct_traj), format=PWDATA.extxyz)
+        # from idx get config idx
+        candidate_idx = json.load(open(os.path.join(self.direct_dir, INIT_BULK.candidate_idx)))
+        candidate_idx_sum = compute_prefix_sum([candidate_idx[_]['num'] for _ in candidate_idx.keys()])
+        _tmp = Config(data_path=os.path.join(self.direct_dir, INIT_BULK.direct_traj), format=PWDATA.extxyz)
+        scf_dir_list = []
+        if self.input_param.dft_input.scf_max_num is not None:
+            random_list = get_random_nums(0, len(candidate.images), self.input_param.dft_input.scf_max_num, seed=2024)
+        else:
+            random_list = None
+        for index, image in enumerate(candidate.images):
+            if random_list is not None and index not in random_list:
+                continue
+            _idx = find_position_binary(candidate_idx_sum, index)
+            config_idx = candidate_idx["{}".format(_idx)]['idx']
+            scf_dir = os.path.join(self.scf_dir, "{}".format(index))
+            if not os.path.exists(scf_dir):
+                os.makedirs(scf_dir)
+                
+            _tmp.images = [image]
+            _tmp.to(data_path=scf_dir, data_name=PWMAT.atom_config,
+                         format=PWDATA.pwmat_config)
+            self.make_scf_file(
+                    scf_dir      =scf_dir, 
+                    traj_file    =os.path.join(scf_dir, PWMAT.atom_config), 
+                    traj_format  =PWDATA.pwmat_config, 
+                    target_format=DFT_STYLE.get_pwdata_format(self.input_param.dft_style, is_cp2k_coord=True),
+                    input_file   =self.init_configs[config_idx].scf_input_file,
+                    kspacing     =self.init_configs[config_idx].scf_kspacing, 
+                    flag_symm    =self.init_configs[config_idx].scf_flag_symm,
+                    is_dftb      = False,
+                    in_skf       =None)
+
+            scf_dir_list.append(scf_dir)
+        
+        self.make_scf_slurm_job_files(scf_dir_list)
+
+    def make_scf_slurm_job_files(self, scf_dir_list:list[str]):
+        del_file_list_by_patten(self.scf_dir, "*{}".format(INIT_BULK.scf_job))
+        group_list = split_job_for_group(self.resource.dft_resource.group_size, scf_dir_list, self.resource.dft_resource.parallel_num)
+        for group_index, group in enumerate(group_list):
+            if group[0] == "NONE":
+                continue
+            jobname = "scf{}".format(group_index)
+            tag_name = "{}-{}".format(group_index, INIT_BULK.scf_tag)
+            tag = os.path.join(self.scf_dir, tag_name)
+            run_cmd = self.resource.dft_resource.command
+            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 = INIT_BULK.scf_tag, 
+                task_tag_faild = INIT_BULK.scf_tag_failed,
+                parallel_num=self.resource.dft_resource.parallel_num,
+                check_type=None
+                )
+            slurm_script_name = "{}-{}".format(group_index, INIT_BULK.scf_job)
+            slurm_job_file =  os.path.join(self.scf_dir, slurm_script_name)
+            write_to_file(slurm_job_file, group_slurm_script, "w")
+
     def do_scf_jobs(self):
         mission = Mission()
         slurm_remain, slurm_success = get_slurm_job_run_info(self.scf_dir, \
@@ -111,89 +174,59 @@ class Relabel(object):
                 mission.all_job_finished(error_type=SLURM_OUT.dft_out)
                 # mission.move_slurm_log_to_slurm_work_dir()
 
-    '''
-    description: 
-        input_file is scf input control file, for vasp is incar, for pwmat is etot.input
-    return {*}
-    author: wuxingxing
-    '''    
-    def make_scf_file(self, scf_dir:str, traj_file:str, traj_format:str, interval:int, target_format:str, \
-                input_file:str, kspacing:float=None, flag_symm:int=None, is_dftb:bool=False, in_skf:str=None):
-        config = load_config(format=traj_format, config=traj_file)
-        index_list = list(range(0, len(config), interval))
-        scf_lsit = []
-        for index in index_list:
-            save_dir = os.path.join(scf_dir, "{}-{}".format(index, INIT_BULK.scf))
-            if not os.path.exists(save_dir):
-                os.makedirs(save_dir)
-            #1. set config file
-            target_config = save_config(config=config[index], 
-                                        input_format=traj_format,# or None, the same
+    def make_scf_file(self, 
+                    scf_dir, 
+                    traj_file    , 
+                    traj_format  , # the input is pwmat/config
+                    target_format,
+                    input_file   ,
+                    kspacing     =None, 
+                    flag_symm    =None,
+                    is_dftb      =None,
+                    in_skf       =None,
+                    atom_names:list[str]=None):
+        if DFT_STYLE.pwmat == self.resource.dft_style:
+            target_config = traj_file
+            pass
+        else:
+            if DFT_STYLE.vasp == self.resource.dft_style: # when do scf, the vasp input file name is 'POSCAR'
+                save_name = VASP.poscar
+            else:
+                save_name="{}".format(DFT_STYLE.get_normal_config(self.resource.dft_style))# for cp2k this param will be set as coord.xzy
+            target_config = save_config(config=traj_file,
+                                        input_format=traj_format,
                                         wrap = False, 
                                         direct = True, 
                                         sort = True, 
-                                        save_format=target_format, 
-                                        save_path=save_dir, 
-                                        save_name=DFT_STYLE.get_normal_config(self.input_param.scf_style))
+                                        save_name = save_name,
+                                        save_format=DFT_STYLE.get_pwdata_format(dft_style=self.resource.dft_style, is_cp2k_coord=True),
+                                        save_path=scf_dir, 
+                                        atom_names=atom_names)
 
-            atom_type_list, _ = get_atom_type(config[index])
-            #2. set pseudo files
-            # if not is_dftb:
-            pseudo_names = link_pseudo_by_atom(
-                pseudo_list     = self.input_param.dft_input.scf_pseudo, 
-                target_dir      = save_dir, 
-                atom_order      = atom_type_list, 
-                dft_style       = self.resource.scf_style,
+        #2.
+        atomic_name_list, atomic_number_list = get_atom_type(traj_file, traj_format)
+        #1. set pseudo files
+        pseudo_names = link_pseudo_by_atom(
+                pseudo_list = self.input_param.dft_input.pseudo, 
+                target_dir = scf_dir, 
+                atom_order = atomic_name_list, 
+                dft_style = self.resource.dft_style,
                 basis_set_file  =self.input_param.dft_input.basis_set_file,
                 potential_file  =self.input_param.dft_input.potential_file
                 )
-            # else:
-            #     # link in.skf path to aimd dir
-            #     pseudo_names = []
-            #     target_dir = os.path.join(aimd_dir, PWMAT.in_skf)
-            #     link_file(in_skf, target_dir)
-            #3. make dft input file
-            set_input_script(
-                input_file=input_file,
-                config=target_config,
-                dft_style=self.resource.scf_style,
-                kspacing=kspacing, 
-                flag_symm=flag_symm, 
-                save_dir = save_dir,
-                pseudo_names=pseudo_names,
-                gaussian_base_param=self.input_param.dft_input.gaussian_base_param,# these for cp2k
-            )
-            scf_lsit.append(save_dir)
-        return scf_lsit
-
-    def make_scf_slurm_job_files(self, scf_dir_list:list[str],use_dftb: bool=False):
-        del_file_list_by_patten(self.scf_dir, "*{}".format(INIT_BULK.scf_job))
-        group_list = split_job_for_group(self.resource.scf_resource.group_size, scf_dir_list, self.resource.scf_resource.parallel_num)
-        for group_index, group in enumerate(group_list):
-            if group[0] == "NONE":
-                continue
-            jobname = "scf{}".format(group_index)
-            tag_name = "{}-{}".format(group_index, INIT_BULK.scf_tag)
-            tag = os.path.join(self.scf_dir, tag_name)
-            run_cmd = self.resource.scf_resource.command
-            group_slurm_script = set_slurm_script_content(gpu_per_node=self.resource.scf_resource.gpu_per_node, 
-                number_node = self.resource.scf_resource.number_node, 
-                cpu_per_node = self.resource.scf_resource.cpu_per_node,
-                queue_name = self.resource.scf_resource.queue_name,
-                custom_flags = self.resource.scf_resource.custom_flags,
-                env_script = self.resource.scf_resource.env_script,
-                job_name = jobname,
-                run_cmd_template = run_cmd,
-                group = group,
-                job_tag = tag,
-                task_tag = INIT_BULK.scf_tag, 
-                task_tag_faild = INIT_BULK.scf_tag_failed,
-                parallel_num=self.resource.scf_resource.parallel_num,
-                check_type=self.resource.scf_style
-                )
-            slurm_script_name = "{}-{}".format(group_index, INIT_BULK.scf_job)
-            slurm_job_file =  os.path.join(self.scf_dir, slurm_script_name)
-            write_to_file(slurm_job_file, group_slurm_script, "w")
+    
+        #2. make etot.input file
+        set_input_script(
+            input_file=input_file,
+            config=target_config,
+            dft_style=self.resource.dft_style,
+            kspacing=kspacing, 
+            flag_symm=flag_symm, 
+            save_dir = scf_dir,
+            pseudo_names=pseudo_names,
+            gaussian_base_param=self.input_param.dft_input.gaussian_base_param,# these for cp2k
+            is_scf = True
+        )
 
     def do_post_process(self):
         if os.path.exists(self.scf_dir):

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()