diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 2e31a1f85..b0fb94e7a 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -378,6 +378,13 @@ def model_devi_lmp_args() -> list[Argument]: doc_epsilon_v = ( "The level parameter for computing the relative virial model deviation." ) + doc_lmp_d3 = "D3 dispersion configuration for LAMMPS. When present, all sub-parameters are required." + doc_lmp_d3_enable = "Enable D3 dispersion correction. If false, D3 will be disabled even if other parameters are present." + doc_lmp_d3_damping_function = "Damping function for D3 dispersion. Options include 'original', 'zerom', 'bj', 'bjm'." + doc_lmp_d3_functional = "Exchange-correlation functional for D3 dispersion. The functional should match that used to train MLPs." + doc_lmp_d3_cutoff = "Cutoff radius for D3 dispersion (in Angstrom)." + doc_lmp_d3_cn_cutoff = "Coordination number cutoff for D3 dispersion (in Angstrom)." + doc_lmp_neigh_modify_one = "Maximum number of neighbors of one atom for 'neigh_modify one N' command. Helps with D3 compatibility." return [ model_devi_jobs_args(), @@ -491,6 +498,27 @@ def model_devi_lmp_args() -> list[Argument]: "use_relative_v", bool, optional=True, default=False, doc=doc_use_relative_v ), Argument("epsilon_v", float, optional=True, doc=doc_epsilon_v), + Argument( + "lmp_d3", + dict, + optional=True, + doc=doc_lmp_d3, + sub_fields=[ + Argument("enable", bool, optional=False, doc=doc_lmp_d3_enable), + Argument( + "damping_function", + str, + optional=False, + doc=doc_lmp_d3_damping_function, + ), + Argument("functional", str, optional=False, doc=doc_lmp_d3_functional), + Argument("cutoff", float, optional=False, doc=doc_lmp_d3_cutoff), + Argument("cn_cutoff", float, optional=False, doc=doc_lmp_d3_cn_cutoff), + ], + ), + Argument( + "lmp_neigh_modify_one", int, optional=True, doc=doc_lmp_neigh_modify_one + ), ] diff --git a/dpgen/generator/lib/lammps.py b/dpgen/generator/lib/lammps.py index 3145f331f..b4bf603d7 100644 --- a/dpgen/generator/lib/lammps.py +++ b/dpgen/generator/lib/lammps.py @@ -82,8 +82,17 @@ def make_lammps_input( ret += "atom_style atomic\n" ret += "\n" ret += "neighbor 1.0 bin\n" + + # Build neigh_modify command with applicable options + neigh_modify_one = jdata.get("lmp_neigh_modify_one") + neigh_modify_options = [] if neidelay is not None: - ret += "neigh_modify delay %d\n" % neidelay # noqa: UP031 + neigh_modify_options.append(f"delay {neidelay}") + if neigh_modify_one is not None: + neigh_modify_options.append(f"one {neigh_modify_one}") + + if neigh_modify_options: + ret += f"neigh_modify {' '.join(neigh_modify_options)}\n" ret += "\n" ret += "box tilt large\n" if nbeads is None: @@ -96,9 +105,23 @@ def make_lammps_input( graph_list = "" for ii in graphs: graph_list += ii + " " + + # Check if D3 dispersion is configured + lmp_d3 = jdata.get("lmp_d3", {}) + d3_enabled = lmp_d3.get("enable", False) if lmp_d3 else False + + if d3_enabled: + # Build D3 parameter string from validated arguments + d3_params = f"{lmp_d3['damping_function']} {lmp_d3['functional']} {lmp_d3['cutoff']} {lmp_d3['cn_cutoff']}" + if Version(deepmd_version) < Version("1"): # 0.x - ret += f"pair_style deepmd {graph_list} ${{THERMO_FREQ}} model_devi.out\n" + if d3_enabled: + ret += f"pair_style hybrid/overlay deepmd {graph_list} ${{THERMO_FREQ}} model_devi.out dispersion/d3 {d3_params}\n" + else: + ret += ( + f"pair_style deepmd {graph_list} ${{THERMO_FREQ}} model_devi.out\n" + ) else: # 1.x keywords = "" @@ -112,11 +135,29 @@ def make_lammps_input( keywords += "fparam ${ELE_TEMP}" if ele_temp_a is not None: keywords += "aparam ${ELE_TEMP}" - if nbeads is None: - ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi.out {keywords}\n" + + if d3_enabled: + # Use hybrid/overlay with D3 + if nbeads is None: + ret += f"pair_style hybrid/overlay deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi.out {keywords} dispersion/d3 {d3_params}\n" + else: + ret += f"pair_style hybrid/overlay deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi${{ibead}}.out {keywords} dispersion/d3 {d3_params}\n" else: - ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi${{ibead}}.out {keywords}\n" - ret += "pair_coeff * *\n" + # Standard deepmd only + if nbeads is None: + ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi.out {keywords}\n" + else: + ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi${{ibead}}.out {keywords}\n" + + # Add pair_coeff lines + if d3_enabled: + # D3 requires type maps (element symbols) + type_map = jdata.get("type_map", []) + type_map_str = " ".join(type_map) + ret += "pair_coeff * * deepmd\n" + ret += f"pair_coeff * * dispersion/d3 {type_map_str}\n" + else: + ret += "pair_coeff * *\n" 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" diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index 74f47f070..b5ea726b6 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -1078,11 +1078,22 @@ def revise_lmp_input_model( ): idx = find_only_one_key(lmp_lines, ["pair_style", "deepmd"]) graph_list = " ".join(task_model_list) + + # Check if D3 dispersion is configured + lmp_d3 = jdata.get("lmp_d3", {}) if jdata else {} + d3_enabled = lmp_d3.get("enable", False) if lmp_d3 else False + if Version(deepmd_version) < Version("1"): - lmp_lines[idx] = "pair_style deepmd %s %d model_devi.out\n" % ( # noqa: UP031 - graph_list, - trj_freq, - ) + if d3_enabled: + d3_params = f"{lmp_d3['damping_function']} {lmp_d3['functional']} {lmp_d3['cutoff']} {lmp_d3['cn_cutoff']}" + lmp_lines[idx] = ( + f"pair_style hybrid/overlay deepmd {graph_list} {trj_freq} model_devi.out dispersion/d3 {d3_params}\n" + ) + else: + lmp_lines[idx] = "pair_style deepmd %s %d model_devi.out\n" % ( # noqa: UP031 + graph_list, + trj_freq, + ) else: # Build keywords string like in make_lammps_input keywords = "" @@ -1097,14 +1108,106 @@ def revise_lmp_input_model( if use_ele_temp == 1: keywords += "fparam ${ELE_TEMP}" - lmp_lines[idx] = ( - "pair_style deepmd %s out_freq %d out_file model_devi.out %s\n" # noqa: UP031 - % ( - graph_list, - trj_freq, - keywords.rstrip(), + if d3_enabled: + d3_params = f"{lmp_d3['damping_function']} {lmp_d3['functional']} {lmp_d3['cutoff']} {lmp_d3['cn_cutoff']}" + lmp_lines[idx] = ( + "pair_style hybrid/overlay deepmd %s out_freq %d out_file model_devi.out %s dispersion/d3 %s\n" # noqa: UP031 + % ( + graph_list, + trj_freq, + keywords.rstrip(), + d3_params, + ) + ) + else: + lmp_lines[idx] = ( + "pair_style deepmd %s out_freq %d out_file model_devi.out %s\n" # noqa: UP031 + % ( + graph_list, + trj_freq, + keywords.rstrip(), + ) ) + return lmp_lines + + +def revise_lmp_input_pair_coeff(lmp_lines, jdata=None): + """Update pair_coeff lines for D3 support.""" + if jdata is None: + return lmp_lines + + lmp_d3 = jdata.get("lmp_d3", {}) + d3_enabled = lmp_d3.get("enable", False) if lmp_d3 else False + + if not d3_enabled: + return lmp_lines + + # D3 requires type maps (element symbols) + type_map = jdata.get("type_map", []) + type_map_str = " ".join(type_map) + + # Find pair_coeff line + pair_coeff_idx = None + for idx, line in enumerate(lmp_lines): + if line.strip().startswith("pair_coeff") and "* *" in line: + pair_coeff_idx = idx + break + + if pair_coeff_idx is None: + # If no pair_coeff found, add them after pair_style + pair_style_idx = find_only_one_key(lmp_lines, ["pair_style"]) + lmp_lines.insert(pair_style_idx + 1, "pair_coeff * * deepmd\n") + lmp_lines.insert( + pair_style_idx + 2, f"pair_coeff * * dispersion/d3 {type_map_str}\n" + ) + else: + # Replace existing pair_coeff with D3 version + lmp_lines[pair_coeff_idx] = "pair_coeff * * deepmd\n" + lmp_lines.insert( + pair_coeff_idx + 1, f"pair_coeff * * dispersion/d3 {type_map_str}\n" ) + + return lmp_lines + + +def revise_lmp_input_neigh_modify(lmp_lines, jdata=None): + """Add neigh_modify one N if requested.""" + if jdata is None: + return lmp_lines + + neigh_modify_one = jdata.get("lmp_neigh_modify_one") + if neigh_modify_one is None: + return lmp_lines + + # Find where to insert neigh_modify one N + # Look for existing neigh_modify lines or insert after neighbor command + neigh_modify_found = False + neighbor_idx = None + + for idx, line in enumerate(lmp_lines): + if line.strip().startswith("neigh_modify") and " one " in line: + neigh_modify_found = True + break + elif line.strip().startswith("neighbor"): + neighbor_idx = idx + + if not neigh_modify_found: + if neighbor_idx is not None: + lmp_lines.insert( + neighbor_idx + 1, f"neigh_modify one {neigh_modify_one}\n" + ) + else: + # Insert after units command if neighbor not found + units_idx = None + for idx, line in enumerate(lmp_lines): + if line.strip().startswith("units"): + units_idx = idx + break + if units_idx is not None: + lmp_lines.insert( + units_idx + 1, f"neigh_modify one {neigh_modify_one}\n" + ) + return lmp_lines @@ -1479,6 +1582,9 @@ def _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems): use_ele_temp=use_ele_temp, jdata=jdata, ) + # Add D3 pair_coeff and neigh_modify support for templates + lmp_lines = revise_lmp_input_pair_coeff(lmp_lines, jdata) + lmp_lines = revise_lmp_input_neigh_modify(lmp_lines, jdata) else: if len(lmp_lines[template_pair_deepmd_idx].split()) != ( len(models) @@ -1501,6 +1607,9 @@ def _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems): use_ele_temp=use_ele_temp, jdata=jdata, ) + # Add D3 pair_coeff and neigh_modify support for templates + lmp_lines = revise_lmp_input_pair_coeff(lmp_lines, jdata) + lmp_lines = revise_lmp_input_neigh_modify(lmp_lines, jdata) # use revise_lmp_input_model to raise error message if "part_style" or "deepmd" not found else: lmp_lines = revise_lmp_input_model( @@ -1512,6 +1621,10 @@ def _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems): jdata=jdata, ) + # Add D3 pair_coeff and neigh_modify support for templates + lmp_lines = revise_lmp_input_pair_coeff(lmp_lines, jdata) + lmp_lines = revise_lmp_input_neigh_modify(lmp_lines, jdata) + lmp_lines = revise_lmp_input_dump( lmp_lines, trj_freq, model_devi_merge_traj ) diff --git a/tests/generator/test_lammps.py b/tests/generator/test_lammps.py index df2d6acf0..c1b915939 100644 --- a/tests/generator/test_lammps.py +++ b/tests/generator/test_lammps.py @@ -7,6 +7,8 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) __package__ = "generator" +from dpgen.generator.lib.lammps import make_lammps_input + from .context import ( get_all_dumped_forces, get_dumped_forces, @@ -14,6 +16,211 @@ ) +class TestMakeLammpsInput(unittest.TestCase): + """Test LAMMPS input generation including D3 dispersion support.""" + + def setUp(self): + self.ensemble = "nvt" + self.conf_file = "test.lmp" + self.graphs = ["model.pb"] + self.nsteps = 1000 + self.dt = 0.001 + self.neidelay = 10 + self.trj_freq = 10 + self.mass_map = [1.0, 16.0] # H, O + self.temp = 300.0 + self.deepmd_version = "2.0" + + def test_basic_deepmd_only(self): + """Test basic LAMMPS input without D3 (backward compatibility).""" + jdata = {} + result = make_lammps_input( + self.ensemble, + self.conf_file, + self.graphs, + self.nsteps, + self.dt, + self.neidelay, + self.trj_freq, + self.mass_map, + self.temp, + jdata, + pres=1.0, + deepmd_version=self.deepmd_version, + ) + + # Should contain basic deepmd pair_style + self.assertIn("pair_style deepmd model.pb", result) + self.assertIn("pair_coeff * *", result) + # Should NOT contain hybrid/overlay or dispersion/d3 + self.assertNotIn("hybrid/overlay", result) + self.assertNotIn("dispersion/d3", result) + + def test_d3_enabled_basic(self): + """Test LAMMPS input with D3 dispersion enabled.""" + jdata = { + "lmp_d3": { + "enable": True, + "damping_function": "original", + "functional": "pbe", + "cutoff": 30.0, + "cn_cutoff": 20.0, + } + } + + result = make_lammps_input( + self.ensemble, + self.conf_file, + self.graphs, + self.nsteps, + self.dt, + self.neidelay, + self.trj_freq, + self.mass_map, + self.temp, + jdata, + pres=1.0, + deepmd_version=self.deepmd_version, + ) + + # Should contain hybrid/overlay pair_style + self.assertIn("pair_style hybrid/overlay deepmd model.pb", result) + self.assertIn("dispersion/d3 original pbe 30.0 20.0", result) + + # Should contain both pair_coeff lines + lines = result.split("\n") + deepmd_coeff_found = False + d3_coeff_found = False + for line in lines: + if "pair_coeff * * deepmd" in line: + deepmd_coeff_found = True + elif "pair_coeff * * dispersion/d3" in line: + d3_coeff_found = True + + self.assertTrue(deepmd_coeff_found, "deepmd pair_coeff not found") + self.assertTrue(d3_coeff_found, "dispersion/d3 pair_coeff not found") + + def test_d3_missing_parameters(self): + """Test that missing D3 parameters raise appropriate errors during validation.""" + # Since validation now happens in arginfo, we would need to test this differently + # For now, test that incomplete D3 config still allows basic operation + jdata = {"lmp_d3": {"enable": False}} # Incomplete but with enable=False + result = make_lammps_input( + self.ensemble, + self.conf_file, + self.graphs, + self.nsteps, + self.dt, + self.neidelay, + self.trj_freq, + self.mass_map, + self.temp, + jdata, + pres=1.0, + deepmd_version=self.deepmd_version, + ) + + # Should behave like basic deepmd since enable=False + self.assertIn("pair_style deepmd model.pb", result) + self.assertNotIn("hybrid/overlay", result) + + def test_d3_disabled(self): + """Test that D3 section with enable=False works like no D3.""" + jdata = { + "lmp_d3": { + "enable": False, + "damping_function": "original", + "functional": "pbe", + "cutoff": 30.0, + "cn_cutoff": 20.0, + } + } + + result = make_lammps_input( + self.ensemble, + self.conf_file, + self.graphs, + self.nsteps, + self.dt, + self.neidelay, + self.trj_freq, + self.mass_map, + self.temp, + jdata, + pres=1.0, + deepmd_version=self.deepmd_version, + ) + + # Should behave like basic deepmd + self.assertIn("pair_style deepmd model.pb", result) + self.assertNotIn("hybrid/overlay", result) + self.assertNotIn("dispersion/d3", result) + + def test_d3_with_neigh_modify_one(self): + """Test D3 with lmp_neigh_modify_one parameter.""" + jdata = { + "lmp_d3": { + "enable": True, + "damping_function": "original", + "functional": "pbe", + "cutoff": 30.0, + "cn_cutoff": 20.0, + }, + "lmp_neigh_modify_one": 2000, + } + + result = make_lammps_input( + self.ensemble, + self.conf_file, + self.graphs, + self.nsteps, + self.dt, + self.neidelay, + self.trj_freq, + self.mass_map, + self.temp, + jdata, + pres=1.0, + deepmd_version=self.deepmd_version, + ) + + # Should contain neigh_modify with one 2000 + self.assertIn("one 2000", result) + # Should also contain D3 configuration + self.assertIn("hybrid/overlay", result) + self.assertIn("dispersion/d3", result) + + def test_d3_different_parameters(self): + """Test D3 with different parameter values.""" + jdata = { + "lmp_d3": { + "enable": True, + "damping_function": "bj", + "functional": "pbe0", + "cutoff": 25.0, + "cn_cutoff": 15.0, + } + } + + result = make_lammps_input( + self.ensemble, + self.conf_file, + self.graphs, + self.nsteps, + self.dt, + self.neidelay, + self.trj_freq, + self.mass_map, + self.temp, + jdata, + pres=1.0, + deepmd_version=self.deepmd_version, + ) + + # Check specific parameter values + self.assertIn("dispersion/d3 bj pbe0 25.0 15.0", result) + + class TestGetDumpForce(unittest.TestCase): def setUp(self): file_content = textwrap.dedent(