support merge_traj for lammps

dpgen/generator/arginfo.py

@@ -158,6 +158,7 @@ def model_devi_lmp_args() -> List[Argument]:
     doc_model_devi_perc_candi_v = 'See model_devi_adapt_trust_lo.'
     doc_model_devi_f_avg_relative = 'Normalized the force model deviations by the RMS force magnitude along the trajectory. This key should not be used with use_relative.'
     doc_model_devi_clean_traj = 'If type of model_devi_clean_traj is bool type then it denote whether to clean traj folders in MD since they are too large. If it is Int type, then the most recent n iterations of traj folders will be retained, others will be removed.'
+    doc_model_devi_merge_traj = 'If model_devi_merge_traj is set as True, only all.lammpstrj will be generated, instead of lots of small traj files.'
     doc_model_devi_nopbc = 'Assume open boundary condition in MD simulations.'
     doc_model_devi_activation_func = 'Set activation functions for models, length of the list should be the same as numb_models, and two elements in the list of string respectively assign activation functions to the embedding and fitting nets within each model. Backward compatibility: the orginal "list of String" format is still supported, where embedding and fitting nets of one model use the same activation function, and the length of the list should be the same as numb_models.'
     doc_shuffle_poscar = 'Shuffle atoms of each frame before running simulations. The purpose is to sample the element occupation of alloys.'
@@ -193,7 +194,9 @@ def model_devi_lmp_args() -> List[Argument]:
         Argument("model_devi_f_avg_relative", bool, optional=True,
                  doc=doc_model_devi_f_avg_relative),
         Argument("model_devi_clean_traj", [
-                 bool, int], optional=False, doc=doc_model_devi_clean_traj),
+                 bool, int], optional=True, default=True , doc=doc_model_devi_clean_traj),
+        Argument("model_devi_merge_traj", [
+                 bool], optional=True, default=False , doc=doc_model_devi_merge_traj),
         Argument("model_devi_nopbc", bool, optional=True, default=False,
                  doc=doc_model_devi_nopbc),
         Argument("model_devi_activation_func", list, optional=True,

dpgen/generator/lib/lammps.py

@@ -89,7 +89,11 @@ def make_lammps_input(ensemble,
     ret+= "\n"
     ret+= "thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz\n"
     ret+= "thermo          ${THERMO_FREQ}\n"
-    ret+= "dump            1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z fx fy fz\n"
+    model_devi_merge_traj = jdata.get('model_devi_merge_traj', False)
+    if(model_devi_merge_traj is True):
+        ret+= "dump            1 all custom ${DUMP_FREQ} all.lammpstrj id type x y z fx fy fz\n"
+    else:
+        ret+= "dump            1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z fx fy fz\n"
     ret+= "restart         10000 dpgen.restart\n"
     ret+= "\n"
     if pka_e is None :
@@ -167,6 +171,42 @@ def get_dumped_forces(
     ret = np.array(ret)
     return ret
 
+def get_all_dumped_forces(
+        file_name):
+    with open(file_name) as fp:        
+        lines = fp.read().split('\n')
+
+    ret = []
+    exist_natoms = False
+    exist_atoms = False
+
+    for idx,ii in enumerate(lines):
+
+        if 'ITEM: NUMBER OF ATOMS' in ii:
+            natoms = int(lines[idx+1])
+            exist_natoms = True
+
+        if 'ITEM: ATOMS' in ii:
+            keys = ii
+            keys = keys.replace('ITEM: ATOMS', '')
+            keys = keys.split()
+            idfx = keys.index('fx')
+            idfy = keys.index('fy')
+            idfz = keys.index('fz')
+            exist_atoms = True
+
+            single_traj = []
+            for jj in range(idx+1, idx+natoms+1):
+                words = lines[jj].split()
+                single_traj.append([ float(words[jj]) for jj in [idfx, idfy, idfz] ])
+            single_traj = np.array(single_traj)
+            ret.append(single_traj)
+
+    if exist_natoms is False:
+        raise RuntimeError('wrong dump file format, cannot find number of atoms', file_name)
+    if exist_atoms is False:
+        raise RuntimeError('wrong dump file format, cannot find dump keys', file_name)
+    return ret
 
 if __name__ == '__main__':
     ret = get_dumped_forces('40.lammpstrj')

dpgen/generator/run.py

@@ -41,7 +41,7 @@
 from dpgen.generator.lib.utils import record_iter
 from dpgen.generator.lib.utils import log_task
 from dpgen.generator.lib.utils import symlink_user_forward_files
-from dpgen.generator.lib.lammps import make_lammps_input, get_dumped_forces
+from dpgen.generator.lib.lammps import make_lammps_input, get_dumped_forces, get_all_dumped_forces
 from dpgen.generator.lib.make_calypso import _make_model_devi_native_calypso,_make_model_devi_buffet
 from dpgen.generator.lib.run_calypso import gen_structures,analysis,run_calypso_model_devi
 from dpgen.generator.lib.parse_calypso import _parse_calypso_input,_parse_calypso_dis_mtx
@@ -190,9 +190,9 @@ def poscar_to_conf(poscar, conf):
     sys.to_lammps_lmp(conf)
 
 
-def dump_to_poscar(dump, poscar, type_map, fmt = "lammps/dump") :
-    sys = dpdata.System(dump, fmt = fmt, type_map = type_map)
-    sys.to_vasp_poscar(poscar)
+# def dump_to_poscar(dump, poscar, type_map, fmt = "lammps/dump") :
+#    sys = dpdata.System(dump, fmt = fmt, type_map = type_map)
+#    sys.to_vasp_poscar(poscar)
 
 def dump_to_deepmd_raw(dump, deepmd_raw, type_map, fmt='gromacs/gro', charge=None):
     system = dpdata.System(dump, fmt = fmt, type_map = type_map)
@@ -988,7 +988,9 @@ def _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems):
                 task_path = os.path.join(work_path, task_name)
                 # create task path
                 create_path(task_path)
-                create_path(os.path.join(task_path, 'traj'))
+                model_devi_merge_traj = jdata.get('model_devi_merge_traj', False)
+                if not model_devi_merge_traj :
+                    create_path(os.path.join(task_path, 'traj'))
                 # link conf
                 loc_conf_name = 'conf.lmp'
                 os.symlink(os.path.join(os.path.join('..','confs'), conf_name),
@@ -1118,7 +1120,9 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems):
                     task_path = os.path.join(work_path, task_name)
                     # dlog.info(task_path)
                     create_path(task_path)
-                    create_path(os.path.join(task_path, 'traj'))
+                    model_devi_merge_traj = jdata.get('model_devi_merge_traj', False)
+                    if not model_devi_merge_traj :
+                        create_path(os.path.join(task_path, 'traj'))
                     loc_conf_name = 'conf.lmp'
                     os.symlink(os.path.join(os.path.join('..','confs'), conf_name),
                                os.path.join(task_path, loc_conf_name) )
@@ -1430,6 +1434,7 @@ def run_md_model_devi (iter_index,
     model_devi_resources = mdata['model_devi_resources']
     use_plm = jdata.get('model_devi_plumed', False)
     use_plm_path = jdata.get('model_devi_plumed_path', False)
+    model_devi_merge_traj = jdata.get('model_devi_merge_traj', False)
 
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, model_devi_name)
@@ -1461,8 +1466,12 @@ def run_md_model_devi (iter_index,
         command = "{ if [ ! -f dpgen.restart.10000 ]; then %s -i input.lammps -v restart 0; else %s -i input.lammps -v restart 1; fi }" % (model_devi_exec, model_devi_exec)
         command = "/bin/sh -c '%s'" % command
         commands = [command]
-        forward_files = ['conf.lmp', 'input.lammps', 'traj']
-        backward_files = ['model_devi.out', 'model_devi.log', 'traj']
+
+        lmp_traj_name = 'traj'
+        if model_devi_merge_traj :
+            lmp_traj_name = 'all.lammpstrj'
+        forward_files = ['conf.lmp', 'input.lammps', lmp_traj_name]
+        backward_files = ['model_devi.out', 'model_devi.log', lmp_traj_name]
         if use_plm:
             forward_files += ['input.plumed']
            # backward_files += ['output.plumed']
@@ -1626,17 +1635,22 @@ def check_bad_box(conf_name,
             raise RuntimeError('unknow key', key)
     return is_bad
 
-
 def _read_model_devi_file(
         task_path : str,
-        model_devi_f_avg_relative : bool = False
+        model_devi_f_avg_relative : bool = False,
+        model_devi_merge_traj : bool = False
 ):
     model_devi = np.loadtxt(os.path.join(task_path, 'model_devi.out'))
     if model_devi_f_avg_relative :
-        trajs = glob.glob(os.path.join(task_path, 'traj', '*.lammpstrj'))
-        all_f = []
-        for ii in trajs:
-            all_f.append(get_dumped_forces(ii))
+        if(model_devi_merge_traj is True) : 
+            all_traj = os.path.join(task_path, 'all.lammpstrj')
+            all_f = get_all_dumped_forces(all_traj)
+        else :
+            trajs = glob.glob(os.path.join(task_path, 'traj', '*.lammpstrj'))
+            all_f = []
+            for ii in trajs:
+                all_f.append(get_dumped_forces(ii))     
+
         all_f = np.array(all_f)
         all_f = all_f.reshape([-1,3])
         avg_f = np.sqrt(np.average(np.sum(np.square(all_f), axis = 1)))
@@ -1655,6 +1669,7 @@ def _select_by_model_devi_standard(
         model_devi_engine : str,
         model_devi_skip : int = 0,
         model_devi_f_avg_relative : bool = False,
+        model_devi_merge_traj : bool = False, 
         detailed_report_make_fp : bool = True,
 ):
     if model_devi_engine == 'calypso':
@@ -1675,7 +1690,7 @@ def _select_by_model_devi_standard(
     for tt in modd_system_task :
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            all_conf = _read_model_devi_file(tt, model_devi_f_avg_relative)
+            all_conf = _read_model_devi_file(tt, model_devi_f_avg_relative, model_devi_merge_traj)
 
             if all_conf.shape == (7,):
                 all_conf = all_conf.reshape(1,all_conf.shape[0])
@@ -1739,6 +1754,7 @@ def _select_by_model_devi_adaptive_trust_low(
         perc_candi_v : float,
         model_devi_skip : int = 0,
         model_devi_f_avg_relative : bool = False,
+        model_devi_merge_traj : bool = False, 
 ):
     """
     modd_system_task    model deviation tasks belonging to one system
@@ -1769,7 +1785,7 @@ def _select_by_model_devi_adaptive_trust_low(
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
             model_devi = np.loadtxt(os.path.join(tt, 'model_devi.out'))
-            model_devi = _read_model_devi_file(tt, model_devi_f_avg_relative)
+            model_devi = _read_model_devi_file(tt, model_devi_f_avg_relative, model_devi_merge_traj)
             for ii in range(model_devi.shape[0]) :
                 if model_devi[ii][0] < model_devi_skip :
                     continue
@@ -1826,7 +1842,8 @@ def _select_by_model_devi_adaptive_trust_low(
     return accur, candi, failed, counter, f_trust_lo, v_trust_lo
 
 
-def _make_fp_vasp_inner (modd_path,
+def _make_fp_vasp_inner (iter_index, 
+                         modd_path,
                          work_path,
                          model_devi_skip,
                          v_trust_lo,
@@ -1839,6 +1856,7 @@ def _make_fp_vasp_inner (modd_path,
                          type_map,
                          jdata):
     """
+    iter_index          int             iter index
     modd_path           string          path of model devi
     work_path           string          path of fp
     fp_task_max         int             max number of tasks
@@ -1880,6 +1898,7 @@ def _make_fp_vasp_inner (modd_path,
     system_index = []
     for ii in modd_task :
         system_index.append(os.path.basename(ii).split('.')[1])
+
     set_tmp = set(system_index)
     system_index = list(set_tmp)
     system_index.sort()
@@ -1892,6 +1911,7 @@ def _make_fp_vasp_inner (modd_path,
     cluster_cutoff = jdata.get('cluster_cutoff', None)
     model_devi_adapt_trust_lo = jdata.get('model_devi_adapt_trust_lo', False)
     model_devi_f_avg_relative = jdata.get('model_devi_f_avg_relative', False)
+    model_devi_merge_traj = jdata.get('model_devi_merge_traj', False)
     # skip save *.out if detailed_report_make_fp is False, default is True
     detailed_report_make_fp = jdata.get("detailed_report_make_fp", True)
     # skip bad box criteria
@@ -1930,6 +1950,7 @@ def _trust_limitation_check(sys_idx, lim):
                         model_devi_engine,
                         model_devi_skip,
                         model_devi_f_avg_relative = model_devi_f_avg_relative,
+                        model_devi_merge_traj = model_devi_merge_traj, 
                         detailed_report_make_fp = detailed_report_make_fp,
                     )
             else:
@@ -1944,6 +1965,7 @@ def _trust_limitation_check(sys_idx, lim):
                         v_trust_hi_sys, numb_candi_v, perc_candi_v,
                         model_devi_skip = model_devi_skip,
                         model_devi_f_avg_relative = model_devi_f_avg_relative,
+                        model_devi_merge_traj = model_devi_merge_traj, 
                     )
                 dlog.info("system {0:s} {1:9s} : f_trust_lo {2:6.3f}   v_trust_lo {3:6.3f}".format(ss, 'adapted', f_trust_lo_ad, v_trust_lo_ad))
         elif model_devi_engine == "amber":
@@ -2042,17 +2064,34 @@ def _trust_limitation_check(sys_idx, lim):
         # ----------------------------------------------------------------------------
         dlog.info("system {0:s} accurate_ratio: {1:8.4f}    thresholds: {2:6.4f} and {3:6.4f}   eff. task min and max {4:4d} {5:4d}   number of fp tasks: {6:6d}".format(ss, accurate_ratio, fp_accurate_soft_threshold, fp_accurate_threshold, fp_task_min, this_fp_task_max, numb_task))
         # make fp tasks
-        model_devi_engine = jdata.get("model_devi_engine", "lammps")
+
+        # read all.lammpstrj, save in all_sys for each system_index
+        all_sys = []
+        trj_freq = None
+        if model_devi_merge_traj :
+            for ii in modd_system_task :
+                all_traj = os.path.join(ii, 'all.lammpstrj')
+                all_sys_per_task = dpdata.System(all_traj, fmt = 'lammps/dump', type_map = type_map)
+                all_sys.append(all_sys_per_task)
+            model_devi_jobs = jdata['model_devi_jobs']
+            cur_job = model_devi_jobs[iter_index]
+            trj_freq = int(_get_param_alias(cur_job, ['t_freq', 'trj_freq', 'traj_freq']))
+
         count_bad_box = 0
         count_bad_cluster = 0
         fp_candidate = sorted(fp_candidate[:numb_task])
+
         for cc in range(numb_task) :
             tt = fp_candidate[cc][0]
             ii = fp_candidate[cc][1]
             ss = os.path.basename(tt).split('.')[1]
             conf_name = os.path.join(tt, "traj")
+            conf_sys = None
             if model_devi_engine == "lammps":
-                conf_name = os.path.join(conf_name, str(ii) + '.lammpstrj')
+                if model_devi_merge_traj :
+                    conf_sys = all_sys[int(os.path.basename(tt).split('.')[-1])][int(int(ii) / trj_freq)]
+                else :
+                    conf_name = os.path.join(conf_name, str(ii) + '.lammpstrj')
                 ffmt = 'lammps/dump'
             elif model_devi_engine == "gromacs":
                 conf_name = os.path.join(conf_name, str(ii) + '.gromacstrj')
@@ -2146,7 +2185,14 @@ def _trust_limitation_check(sys_idx, lim):
         for idx, task in enumerate(fp_tasks):
             os.chdir(task)
             if model_devi_engine == "lammps":
-                dump_to_poscar('conf.dump', 'POSCAR', type_map, fmt = "lammps/dump")
+                sys = None
+                if model_devi_merge_traj:
+                    sys = conf_sys
+                else :
+                    sys = dpdata.System('conf.dump', fmt = "lammps/dump", type_map = type_map)
+                sys.to_vasp_poscar('POSCAR')
+                # dump to poscar 
+
                 if charges_map:
                     warnings.warn('"sys_charges" keyword only support for gromacs engine now.')
             elif model_devi_engine == "gromacs":
@@ -2482,7 +2528,8 @@ def _make_fp_vasp_configs(iter_index,
         f_trust_hi = jdata['model_devi_f_trust_hi']
 
     # make configs
-    fp_tasks = _make_fp_vasp_inner(modd_path, work_path,
+    fp_tasks = _make_fp_vasp_inner(iter_index, 
+                                   modd_path, work_path,
                                    model_devi_skip,
                                    v_trust_lo, v_trust_hi,
                                    f_trust_lo, f_trust_hi,

tests/dispatcher/loc/task0/dir0/test2

@@ -0,0 +1 @@
+f0c56f70-627b-445a-89de-3ad3e81f3785

tests/dispatcher/loc/task0/test0

@@ -0,0 +1 @@
+5dc17ca2-0d58-4968-af22-41536e667668

tests/dispatcher/loc/task0/test1

@@ -0,0 +1 @@
+814f28b0-38bd-493a-b400-d678c3fe1a0e

tests/dispatcher/loc/task1/dir0/test2

@@ -0,0 +1 @@
+d3c9fb33-4ffd-48c2-86bb-933fbe7fd512

tests/dispatcher/loc/task1/test0

@@ -0,0 +1 @@
+8a4f3e52-3ace-45c5-8bd1-1bbbb4d9abd0

tests/dispatcher/loc/task1/test1

@@ -0,0 +1 @@
+b8389a9b-bc75-4498-94bf-c565a4e387b6

tests/generator/context.py

@@ -5,6 +5,7 @@
 from dpgen.generator.lib.gaussian import detect_multiplicity
 from dpgen.generator.lib.ele_temp import NBandsEsti
 from dpgen.generator.lib.lammps import get_dumped_forces
+from dpgen.generator.lib.lammps import get_all_dumped_forces
 from dpgen.generator.lib.make_calypso import make_calypso_input,write_model_devi_out
 from dpgen.generator.lib.parse_calypso import _parse_calypso_input,_parse_calypso_dis_mtx
 

tests/generator/test_lammps.py

@@ -6,6 +6,7 @@
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 __package__ = 'generator'
 from .context import get_dumped_forces
+from .context import get_all_dumped_forces
 
 class TestGetDumpForce(unittest.TestCase):
     def setUp(self):
@@ -36,3 +37,44 @@ def test_read_dump(self):
         ff = ff.reshape([-1])
         for ii in range(6):
             self.assertAlmostEqual(ff[ii], self.expected_f[ii])
+
+class TestGetDumpForce(unittest.TestCase):
+    def setUp(self):
+        file_content = textwrap.dedent("""\
+ITEM: TIMESTEP
+0
+ITEM: NUMBER OF ATOMS
+2
+ITEM: BOX BOUNDS xy xz yz pp pp pp
+0.0000000000000000e+00 1.0000000000000000e+01 0.0000000000000000e+00
+0.0000000000000000e+00 1.0000000000000000e+01 0.0000000000000000e+00
+0.0000000000000000e+00 1.0000000000000000e+01 0.0000000000000000e+00
+ITEM: ATOMS id type x y z fx fy fz
+1 1 5.38154 4.06861 3.60573 0.000868817 -0.00100822 -0.000960258
+2 2 3.9454 4.80321 4.38469 0.000503458 -0.000374043 -9.15676e-05
+ITEM: TIMESTEP
+10
+ITEM: NUMBER OF ATOMS
+2
+ITEM: BOX BOUNDS xy xz yz pp pp pp
+0.0000000000000000e+00 1.0000000000000000e+01 0.0000000000000000e+00
+0.0000000000000000e+00 1.0000000000000000e+01 0.0000000000000000e+00
+0.0000000000000000e+00 1.0000000000000000e+01 0.0000000000000000e+00
+ITEM: ATOMS id type x y z fx fy fz
+1 1 5.35629 3.93297 3.70556 -0.125424 0.0481604 -0.0833015
+2 2 3.93654 4.79972 4.48179 0.134843 -0.0444238 -0.143111
+""")
+        with open('tmp.dump', 'w') as fp:
+            fp.write(file_content)
+        self.expected_f = [ 0.000868817 , -0.00100822 , -0.000960258 , 0.000503458 , -0.000374043 , -9.15676e-05 , -0.125424 , 0.0481604 , -0.0833015 , 0.134843 , -0.0444238 , -0.143111]
+    def tearDown(self):
+        if os.path.isfile('tmp.dump'):
+            os.remove('tmp.dump')
+
+    def test_read_all_dump(self):
+        ff = get_all_dumped_forces('tmp.dump')
+        ff = np.array(ff)
+        self.assertEqual(ff.shape, (2,2,3))
+        ff = ff.reshape([-1])
+        for ii in range(12):
+            self.assertAlmostEqual(ff[ii], self.expected_f[ii])