Add gromacs in model deviation (#427)

* Add Cu's init_surf examples.

* Delete personal information

* Update examples/init/cu.surf.hcp.111.json

Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

* Add gromacs as model_devi engine

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* Add gromacs in model_devi

* wyz: fix sort of atoms in gaussian input for gromacs model_devi_engine

* wyz: add modifying trj_freq for model_devi_gromacs enegine

* Add gromacs in model_devi

* wyz: add model_devi_dt in gmx enengine

* add unittest for gmx engine

* wyz: delete extra vacant lines

* Delete copy files

* Add GromacsWrapper Dependency

* Fix unittest in Gromacs model_devi

* Fix unittest of DP-GEN + gromacs

* Fix unittest in gromacs model_devi

* Fix unittest in gromacs model_devi

* Add dpgen-gromacs-gaussian example

* Add model_devi script

* Delete model_devi scripts. This hass been in tests

* Fix unittest in model_devi

* Fix unittest in Gromacs model_devi

Co-authored-by: Han Wang <amcadmus@gmail.com>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: zjgemi <zjgemi@gmail.com>

Author: 4 people

dpgen/generator/run.py

@@ -4,7 +4,7 @@
 init: data
 iter:
         00.train
-        01.mode_devi
+        01.model_devi
         02.vasp
         03.data
 """
@@ -62,6 +62,11 @@
 from dpgen import ROOT_PATH
 from pymatgen.io.vasp import Incar,Kpoints,Potcar
 from dpgen.auto_test.lib.vasp import make_kspacing_kpoints
+try:
+    from gromacs.fileformats.mdp import MDP
+except ImportError:
+    dlog.info("GromacsWrapper>=0.8.0 is needed for DP-GEN + Gromacs.")
+    pass
 
 template_name = 'template'
 train_name = '00.train'
@@ -176,10 +181,14 @@ def poscar_to_conf(poscar, conf):
     sys.to_lammps_lmp(conf)
 
 
-def dump_to_poscar(dump, poscar, type_map) :
-    sys = dpdata.System(dump, fmt = 'lammps/dump', type_map = type_map)
+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'):
+    system = dpdata.System(dump, fmt = fmt, type_map = type_map)
+    system.to_deepmd_raw(deepmd_raw)
+
 
 def make_train (iter_index,
                 jdata,
@@ -684,6 +693,9 @@ def revise_by_keys(lmp_lines, keys, values):
 def make_model_devi (iter_index,
                      jdata,
                      mdata) :
+    # The MD engine to perform model deviation
+    # Default is lammps
+    model_devi_engine = jdata.get('model_devi_engine', "lammps")
     model_devi_jobs = jdata['model_devi_jobs']
     if (iter_index >= len(model_devi_jobs)) :
         return False
@@ -729,24 +741,27 @@ def make_model_devi (iter_index,
     for ss in conf_systems:
         conf_counter = 0
         for cc in ss :
-            conf_name = make_model_devi_conf_name(sys_idx[sys_counter], conf_counter)
-            orig_poscar_name = conf_name + '.orig.poscar'
-            poscar_name = conf_name + '.poscar'
-            lmp_name = conf_name + '.lmp'
-            if shuffle_poscar :
-                os.symlink(cc, os.path.join(conf_path, orig_poscar_name))
-                poscar_shuffle(os.path.join(conf_path, orig_poscar_name),
-                               os.path.join(conf_path, poscar_name))
-            else :
-                os.symlink(cc, os.path.join(conf_path, poscar_name))
-            if 'sys_format' in jdata:
-                fmt = jdata['sys_format']
-            else:
-                fmt = 'vasp/poscar'
-            system = dpdata.System(os.path.join(conf_path, poscar_name), fmt = fmt, type_map = jdata['type_map'])
-            if jdata.get('model_devi_nopbc', False):
-                system.remove_pbc()
-            system.to_lammps_lmp(os.path.join(conf_path, lmp_name))
+            if model_devi_engine == "lammps":
+                conf_name = make_model_devi_conf_name(sys_idx[sys_counter], conf_counter)
+                orig_poscar_name = conf_name + '.orig.poscar'
+                poscar_name = conf_name + '.poscar'
+                lmp_name = conf_name + '.lmp'
+                if shuffle_poscar :
+                    os.symlink(cc, os.path.join(conf_path, orig_poscar_name))
+                    poscar_shuffle(os.path.join(conf_path, orig_poscar_name),
+                                   os.path.join(conf_path, poscar_name))
+                else :
+                    os.symlink(cc, os.path.join(conf_path, poscar_name))
+                if 'sys_format' in jdata:
+                    fmt = jdata['sys_format']
+                else:
+                    fmt = 'vasp/poscar'
+                system = dpdata.System(os.path.join(conf_path, poscar_name), fmt = fmt, type_map = jdata['type_map'])
+                if jdata.get('model_devi_nopbc', False):
+                    system.remove_pbc()
+                system.to_lammps_lmp(os.path.join(conf_path, lmp_name))
+            elif model_devi_engine == "gromacs":
+                pass
             conf_counter += 1
         sys_counter += 1
 
@@ -756,7 +771,12 @@ def make_model_devi (iter_index,
     use_plm = jdata.get('model_devi_plumed', False)
     use_plm_path = jdata.get('model_devi_plumed_path', False)
     if input_mode == "native":
-        _make_model_devi_native(iter_index, jdata, mdata, conf_systems)
+        if model_devi_engine == "lammps":
+            _make_model_devi_native(iter_index, jdata, mdata, conf_systems)
+        elif model_devi_engine == "gromacs":
+            _make_model_devi_native_gromacs(iter_index, jdata, mdata, conf_systems)
+        else:
+            raise RuntimeError("unknown model_devi engine", model_devi_engine)
     elif input_mode == "revise_template":
         _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems)
     else:
@@ -997,12 +1017,98 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems):
             conf_counter += 1
         sys_counter += 1
 
+def _make_model_devi_native_gromacs(iter_index, jdata, mdata, conf_systems):
+    model_devi_jobs = jdata['model_devi_jobs']
+    if (iter_index >= len(model_devi_jobs)) :
+        return False
+    cur_job = model_devi_jobs[iter_index]
+    dt = cur_job.get("dt", None)
+    if dt is not None:
+        model_devi_dt = dt
+    else:
+        model_devi_dt = jdata['model_devi_dt']
+    nsteps = cur_job.get("nsteps", None)
+    if nsteps is None:
+        raise RuntimeError("nsteps is None, you should set nsteps in model_devi_jobs!")
+    # Currently Gromacs engine is not supported for different temperatures!
+    # If you want to change temperatures, you should change it in mdp files.
+  
+    sys_idx = expand_idx(cur_job['sys_idx'])
+    if (len(sys_idx) != len(list(set(sys_idx)))) :
+        raise RuntimeError("system index should be uniq")
+
+    mass_map = jdata['mass_map']
+
+    iter_name = make_iter_name(iter_index)
+    train_path = os.path.join(iter_name, train_name)
+    train_path = os.path.abspath(train_path)
+    models = glob.glob(os.path.join(train_path, "graph*pb"))
+    task_model_list = []
+    for ii in models:
+        task_model_list.append(os.path.join('..', os.path.basename(ii)))
+    work_path = os.path.join(iter_name, model_devi_name)
+
+    sys_counter = 0
+    for ss in conf_systems:
+        conf_counter = 0
+        task_counter = 0
+        for cc in ss :
+            task_name = make_model_devi_task_name(sys_idx[sys_counter], task_counter)
+            #conf_name = make_model_devi_conf_name(sys_idx[sys_counter], conf_counter) + '.lmp'
+            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'))
+            #loc_conf_name = 'conf.lmp'
+            gromacs_settings = jdata.get("gromacs_settings" , "")
+            for key,file in gromacs_settings.items():
+                if key != "traj_filename" and key != "mdp_filename":
+                    os.symlink(os.path.join(cc,file), os.path.join(task_path, file))
+            
+            # input.json for DP-Gromacs
+            with open(os.path.join(cc, "input.json")) as f:
+                input_json = json.load(f)
+            input_json["graph_file"] = models[0]
+            with open(os.path.join(task_path,'input.json'), 'w') as _outfile:
+                json.dump(input_json, _outfile, indent = 4)
+
+            # trj_freq
+            trj_freq = cur_job.get("trj_freq", 10)
+            mdp = MDP()
+            mdp.read(os.path.join(cc, gromacs_settings['mdp_filename']))
+            mdp['nstcomm'] = trj_freq
+            mdp['nstxout'] = trj_freq
+            mdp['nstlog'] = trj_freq
+            mdp['nstenergy'] = trj_freq
+            # dt
+            mdp['dt'] = dt
+            mdp.write(os.path.join(task_path, gromacs_settings['mdp_filename']))
+
+            cwd_ = os.getcwd()
+            os.chdir(task_path)
+            job = {}
+            
+            job["model_devi_dt"] =  model_devi_dt
+            job["nsteps"] = nsteps
+            with open('job.json', 'w') as _outfile:
+                json.dump(job, _outfile, indent = 4)
+            os.chdir(cwd_)
+            
+            task_counter += 1
+            conf_counter += 1
+        sys_counter += 1
+
+
 
 def run_model_devi (iter_index,
                     jdata,
                     mdata) :
     #rmdlog.info("This module has been run !")
     lmp_exec = mdata['lmp_command']
+    # Angus: lmp_exec name should be changed to model_devi_exec.
+    # We should also change make_dispatcher
+    # For now, I will use this name for gromacs command
+
     model_devi_group_size = mdata['model_devi_group_size']
     model_devi_resources = mdata['model_devi_resources']
     use_plm = jdata.get('model_devi_plumed', False)
@@ -1014,10 +1120,6 @@ def run_model_devi (iter_index,
 
     all_task = glob.glob(os.path.join(work_path, "task.*"))
     all_task.sort()
-    command = "{ if [ ! -f dpgen.restart.10000 ]; then %s -i input.lammps -v restart 0; else %s -i input.lammps -v restart 1; fi }" % (lmp_exec, lmp_exec)
-    command = "/bin/sh -c '%s'" % command
-    commands = [command]
-
     fp = open (os.path.join(work_path, 'cur_job.json'), 'r')
     cur_job = json.load (fp)
 
@@ -1036,14 +1138,37 @@ def run_model_devi (iter_index,
     #dlog.info("run_tasks in run_model_deviation",run_tasks_)
     all_models = glob.glob(os.path.join(work_path, 'graph*pb'))
     model_names = [os.path.basename(ii) for ii in all_models]
-    forward_files = ['conf.lmp', 'input.lammps', 'traj']
-    backward_files = ['model_devi.out', 'model_devi.log', 'traj']
-    if use_plm:
-        forward_files += ['input.plumed']
-       # backward_files += ['output.plumed']
-        backward_files += ['output.plumed','COLVAR','dump.0.xyz']
-        if use_plm_path:
-            forward_files += ['plmpath.pdb']
+
+    model_devi_engine = jdata.get("model_devi_engine", "lammps")
+    if model_devi_engine == "lammps":
+        command = "{ if [ ! -f dpgen.restart.10000 ]; then %s -i input.lammps -v restart 0; else %s -i input.lammps -v restart 1; fi }" % (lmp_exec, lmp_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']
+        if use_plm:
+            forward_files += ['input.plumed']
+           # backward_files += ['output.plumed']
+            backward_files += ['output.plumed','COLVAR','dump.0.xyz']
+            if use_plm_path:
+                forward_files += ['plmpath.pdb']
+    elif model_devi_engine == "gromacs":
+        gromacs_settings = jdata.get("gromacs_settings", {})
+        mdp_filename = gromacs_settings.get("mdp_filename", "md.mdp")
+        topol_filename = gromacs_settings.get("topol_filename", "processed.top")
+        conf_filename = gromacs_settings.get("conf_filename", "conf.gro")
+        index_filename = gromacs_settings.get("index_filename", "index.raw")
+        deffnm = gromacs_settings.get("deffnm", "deepmd")
+        maxwarn = gromacs_settings.get("maxwarn", 1)
+        nsteps = cur_job["nsteps"]
+
+        command = "%s grompp -f %s -p %s -c %s -o %s -maxwarn %d" % (lmp_exec, mdp_filename, topol_filename, conf_filename, deffnm, maxwarn)
+        command += "&& %s mdrun -deffnm %s -nsteps %d" %(lmp_exec, deffnm, nsteps) 
+        commands = [command]
+        
+        forward_files = [mdp_filename, topol_filename, conf_filename, index_filename,  "input.json" ]
+        backward_files = ["%s.tpr" % deffnm, "%s.log" %deffnm , 'model_devi.out', 'model_devi.log']
+
 
     cwd = os.getcwd()
     dispatcher = make_dispatcher(mdata['model_devi_machine'], mdata['model_devi_resources'], work_path, run_tasks, model_devi_group_size)
@@ -1240,14 +1365,20 @@ def _make_fp_vasp_inner (modd_path,
             numb_task = 0
         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")
         count_bad_box = 0
         count_bad_cluster = 0
         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_name = os.path.join(conf_name, str(ii) + '.lammpstrj')
+            if model_devi_engine == "lammps":
+                conf_name = os.path.join(conf_name, str(ii) + '.lammpstrj')
+            elif model_devi_engine == "gromacs":
+                conf_name = os.path.join(conf_name, str(ii) + '.gromacstrj')
+            else:
+                raise RuntimeError("unknown model_devi engine", model_devi_engine)
             conf_name = os.path.abspath(conf_name)
             if skip_bad_box is not None:
                 skip = check_bad_box(conf_name, skip_bad_box)
@@ -1295,7 +1426,13 @@ def _make_fp_vasp_inner (modd_path,
         cwd = os.getcwd()
         for ii in fp_tasks:
             os.chdir(ii)
-            dump_to_poscar('conf.dump', 'POSCAR', type_map)
+            if model_devi_engine == "lammps":
+                dump_to_poscar('conf.dump', 'POSCAR', type_map, fmt = "lammps/dump")
+            elif model_devi_engine == "gromacs":
+                # dump_to_poscar('conf.dump', 'POSCAR', type_map, fmt = "gromacs/gro")
+                dump_to_deepmd_raw('conf.dump', 'deepmd.raw', type_map, fmt = 'gromacs/gro')
+            else:
+                raise RuntimeError("unknown model_devi engine", model_devi_engine)
             os.chdir(cwd)
     return fp_tasks
 
@@ -1676,9 +1813,14 @@ def make_fp_gaussian(iter_index,
     else:
         fp_params = jdata['fp_params']
     cwd = os.getcwd()
+
+    model_devi_engine = jdata.get('model_devi_engine', 'lammps')
     for ii in fp_tasks:
         os.chdir(ii)
-        sys_data = dpdata.System('POSCAR').data
+        if model_devi_engine == "lammps":
+            sys_data = dpdata.System('POSCAR').data
+        elif model_devi_engine == "gromacs":
+            sys_data = dpdata.System("deepmd.raw", fmt='deepmd/raw').data
         ret = make_gaussian_input(sys_data, fp_params)
         with open('input', 'w') as fp:
             fp.write(ret)
@@ -2285,10 +2427,6 @@ def set_version(mdata):
     mdata['deepmd_version'] = deepmd_version
     return mdata
 
-
-
-
-
 def run_iter (param_file, machine_file) :
     try:
        import ruamel