diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 0f8abb865..b319eab41 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -621,9 +621,189 @@ def model_devi_amber_args() -> list[Argument]: ] +def model_devi_gromacs_jobs_args() -> Argument: + """Return per-iteration exploration settings used by GROMACS.""" + return Argument( + "model_devi_jobs", + list, + repeat=True, + doc=( + "GROMACS exploration settings. Each entry corresponds to one " + "DP-GEN iteration." + ), + sub_fields=[ + Argument( + "sys_idx", + list[int], + optional=False, + doc="Indices of systems selected from sys_configs.", + ), + Argument( + "temps", + list[float], + optional=True, + doc="Temperatures written to ref_t/ref-t in the MDP file.", + ), + Argument( + "press", + list[float], + optional=True, + doc="Accepted for compatibility; pressure is configured by the MDP file.", + ), + Argument( + "trj_freq", + int, + optional=False, + doc="Frequency for trajectory, log, and energy output.", + ), + Argument("nsteps", int, optional=False, doc="Number of GROMACS MD steps."), + Argument( + "ensemble", + str, + optional=True, + doc="Accepted for compatibility; the MDP file controls the ensemble.", + ), + Argument( + "lambdas", + list[float], + optional=True, + doc="Deep potential scaling factors in the inclusive range [0, 1].", + ), + Argument( + "dt", + float, + optional=True, + doc="Per-iteration timestep overriding model_devi_dt.", + ), + ], + ) + + +def model_devi_gromacs_args() -> list[Argument]: + """Return model-deviation arguments supported by the GROMACS engine.""" + common_names = { + "model_devi_dt", + "model_devi_skip", + "model_devi_f_trust_lo", + "model_devi_f_trust_hi", + "model_devi_v_trust_lo", + "model_devi_v_trust_hi", + "model_devi_adapt_trust_lo", + "model_devi_numb_candi_f", + "model_devi_numb_candi_v", + "model_devi_perc_candi_f", + "model_devi_perc_candi_v", + "model_devi_f_avg_relative", + "model_devi_clean_traj", + "shuffle_poscar", + } + common_args = [ + argument for argument in model_devi_lmp_args() if argument.name in common_names + ] + settings_args = [ + Argument( + "mdp_filename", + str, + optional=False, + doc="GROMACS molecular-dynamics parameter file.", + ), + Argument( + "topol_filename", + str, + optional=False, + doc="GROMACS topology file.", + ), + Argument( + "conf_filename", + str, + optional=False, + doc="Initial GROMACS coordinate file.", + ), + Argument( + "index_filename", + str, + optional=False, + doc="DeepMD atom-index file.", + ), + Argument( + "type_filename", + str, + optional=False, + doc="DeepMD atom-type file.", + ), + Argument( + "ref_filename", + str, + optional=False, + doc="Reference TPR file used by trjconv for periodic-boundary processing.", + ), + Argument( + "ndx_filename", + str, + optional=True, + doc="Optional GROMACS index file passed to trjconv.", + ), + Argument( + "model_devi_script", + str, + optional=True, + doc="Optional user file forwarded with each GROMACS task.", + ), + Argument( + "deffnm", + str, + optional=True, + doc="GROMACS output prefix; defaults to 'deepmd' at runtime.", + ), + Argument( + "maxwarn", + int, + optional=True, + doc="Maximum warnings accepted by grompp; defaults to 1.", + ), + Argument( + "traj_filename", + str, + optional=True, + doc="Processed trajectory filename; defaults to 'deepmd_traj.gro'.", + ), + Argument( + "group_name", + str, + optional=True, + doc="Group selected twice for trjconv; defaults to 'Other'.", + ), + ] + return [ + model_devi_gromacs_jobs_args(), + *common_args, + Argument( + "model_devi_nopbc", + bool, + optional=True, + default=False, + doc=( + "Accepted for compatibility with existing inputs. GROMACS " + "trajectory processing currently applies periodic-boundary handling." + ), + ), + Argument( + "gromacs_settings", + dict, + optional=False, + sub_fields=settings_args, + doc="GROMACS input filenames and trajectory-processing settings.", + ), + ] + + def model_devi_args() -> list[Variant]: doc_model_devi_engine = "Engine for the model deviation task." doc_amber = "Amber DPRc engine. The command argument in the machine file should be path to sander." + doc_gromacs = ( + "GROMACS engine. Requires GromacsWrapper and DeePMD-kit 2 or later. " + "The machine command should invoke the gmx executable." + ) return [ Variant( "model_devi_engine", @@ -631,7 +811,7 @@ def model_devi_args() -> list[Variant]: Argument("lammps", dict, model_devi_lmp_args(), doc="LAMMPS"), Argument("amber", dict, model_devi_amber_args(), doc=doc_amber), Argument("calypso", dict, [], doc="TODO: add doc"), - Argument("gromacs", dict, [], doc="TODO: add doc"), + Argument("gromacs", dict, model_devi_gromacs_args(), doc=doc_gromacs), ], default_tag="lammps", optional=True, diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index 0ce513a96..3e04918e4 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -1806,6 +1806,20 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems): sys_counter += 1 +def _gromacs_input_files(gromacs_settings): + """Return only settings whose values name files staged for a GROMACS task.""" + non_input_settings = { + "traj_filename", + "mdp_filename", + "group_name", + "maxwarn", + "deffnm", + } + return [ + file for key, file in gromacs_settings.items() if key not in non_input_settings + ] + + def _make_model_devi_native_gromacs(iter_index, jdata, mdata, conf_systems): try: from gromacs.fileformats.mdp import MDP @@ -1868,16 +1882,10 @@ def _make_model_devi_native_gromacs(iter_index, jdata, mdata, conf_systems): task_path = os.path.join(work_path, task_name) create_path(task_path) gromacs_settings = jdata.get("gromacs_settings", "") - for key, file in gromacs_settings.items(): - if ( - key != "traj_filename" - and key != "mdp_filename" - and key != "group_name" - and key != "maxwarn" - ): - os.symlink( - os.path.join(cc, file), os.path.join(task_path, file) - ) + for file in _gromacs_input_files(gromacs_settings): + 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) diff --git a/tests/generator/test_gromacs_engine.py b/tests/generator/test_gromacs_engine.py index d74d1333d..30cfe2ee6 100644 --- a/tests/generator/test_gromacs_engine.py +++ b/tests/generator/test_gromacs_engine.py @@ -6,6 +6,8 @@ import numpy as np +from dpgen.generator.run import _gromacs_input_files + sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) __package__ = "generator" dirname = os.path.join(os.path.abspath(os.path.dirname(__file__)), "gromacs") @@ -43,6 +45,7 @@ def setUp(self): "ref_filename": "em.tpr", "model_devi_script": "model_devi.py", "traj_filename": "deepmd_traj.gro", + "deffnm": "deepmd", }, "model_devi_dt": 0.001, "model_devi_f_trust_lo": 0.05, @@ -80,7 +83,10 @@ def setUp(self): def _check_dir(self, wdir, post=True): for key in self.jdata["gromacs_settings"].keys(): - if key != "traj_filename": + if key == "deffnm": + # deffnm names generated outputs and must not be staged as input. + self.assertFalse(os.path.exists(os.path.join(wdir, "deepmd"))) + elif key != "traj_filename": self.assertTrue( os.path.exists( os.path.join(wdir, self.jdata["gromacs_settings"][key]) @@ -105,6 +111,13 @@ def _copy_outputs(self, path_1, path_2): ) shutil.copytree(os.path.join(path_1, "traj"), os.path.join(path_2, "traj")) + def test_deffnm_is_not_an_input_file(self): + """Command/output settings must not be symlinked from the source system.""" + input_files = _gromacs_input_files(self.jdata["gromacs_settings"]) + + self.assertNotIn("deepmd", input_files) + self.assertIn("processed.top", input_files) + @unittest.skipIf( importlib.util.find_spec("openbabel") is not None, "when openbabel is found, this test will be skipped. ", diff --git a/tests/test_check_examples.py b/tests/test_check_examples.py index 23520428a..fbd9eac26 100644 --- a/tests/test_check_examples.py +++ b/tests/test_check_examples.py @@ -6,12 +6,18 @@ import unittest from pathlib import Path +from dargs import Argument + from dpgen.data.arginfo import ( init_bulk_jdata_arginfo, init_reaction_jdata_arginfo, init_surf_jdata_arginfo, ) -from dpgen.generator.arginfo import run_jdata_arginfo, run_mdata_arginfo +from dpgen.generator.arginfo import ( + model_devi_args, + run_jdata_arginfo, + run_mdata_arginfo, +) from dpgen.simplify.arginfo import simplify_jdata_arginfo, simplify_mdata_arginfo from dpgen.util import normalize @@ -207,3 +213,47 @@ def test_arguments(self): with open(fn) as f: data = json.load(f) normalize(arginfo, data) + + def test_gromacs_model_devi_arguments(self): + """The GROMACS schema should accept every setting used by its runner.""" + arginfo = Argument("model_devi", dict, sub_variants=model_devi_args()) + data = { + "model_devi_engine": "gromacs", + "model_devi_jobs": [ + { + "sys_idx": [0], + "temps": [300.0], + "press": [], + "trj_freq": 10, + "nsteps": 100, + "ensemble": "nvt", + "lambdas": [0.5, 1.0], + "dt": 0.001, + } + ], + "model_devi_dt": 0.002, + "model_devi_skip": 0, + "model_devi_f_trust_lo": 0.2, + "model_devi_f_trust_hi": 0.6, + "model_devi_v_trust_lo": 1e10, + "model_devi_v_trust_hi": 1e10, + "model_devi_clean_traj": False, + "model_devi_nopbc": True, + "gromacs_settings": { + "mdp_filename": "md.mdp", + "topol_filename": "processed.top", + "conf_filename": "conf.gro", + "index_filename": "index.raw", + "type_filename": "type.raw", + "ref_filename": "em.tpr", + "ndx_filename": "index.ndx", + "model_devi_script": "model_devi.py", + "deffnm": "deepmd", + "maxwarn": 1, + "traj_filename": "deepmd_traj.gro", + "group_name": "Other", + }, + } + + normalized = arginfo.normalize_value(data) + arginfo.check_value(normalized, strict=True)