From 64adbc41995bd6d0eaaf37e2c4984e47c2574e1e Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 8 Jan 2026 11:16:39 +0800 Subject: [PATCH 01/25] add dplr(torch) --- deepmd/pt/modifier/__init__.py | 4 + deepmd/pt/modifier/dipole_charge.py | 347 ++++++++++++++++++ pyproject.toml | 1 + source/tests/pt/modifier/__init__.py | 1 + .../pt/{ => modifier}/test_data_modifier.py | 2 +- .../tests/pt/modifier/test_dipole_charge.py | 220 +++++++++++ source/tests/pt/modifier/water | 1 + source/tests/pt/modifier/water_tensor | 1 + 8 files changed, 576 insertions(+), 1 deletion(-) create mode 100644 deepmd/pt/modifier/dipole_charge.py create mode 100644 source/tests/pt/modifier/__init__.py rename source/tests/pt/{ => modifier}/test_data_modifier.py (99%) create mode 100644 source/tests/pt/modifier/test_dipole_charge.py create mode 120000 source/tests/pt/modifier/water create mode 120000 source/tests/pt/modifier/water_tensor diff --git a/deepmd/pt/modifier/__init__.py b/deepmd/pt/modifier/__init__.py index 71d847bcbc..f196ebf000 100644 --- a/deepmd/pt/modifier/__init__.py +++ b/deepmd/pt/modifier/__init__.py @@ -7,9 +7,13 @@ from .base_modifier import ( BaseModifier, ) +from .dipole_charge import ( + DipoleChargeModifier, +) __all__ = [ "BaseModifier", + "DipoleChargeModifier", "get_data_modifier", ] diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py new file mode 100644 index 0000000000..3ab1941c97 --- /dev/null +++ b/deepmd/pt/modifier/dipole_charge.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import torch +from torch_admp.pme import ( + CoulombForceModule, +) +from torch_admp.utils import ( + calc_grads, +) + +from deepmd.pt.modifier.base_modifier import ( + BaseModifier, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.utils import ( + to_torch_tensor, +) + + +@BaseModifier.register("dipole_charge") +class DipoleChargeModifier(BaseModifier): + """Parameters + ---------- + model_name + The model file for the DeepDipole model + model_charge_map + Gives the amount of charge for the wfcc + sys_charge_map + Gives the amount of charge for the real atoms + ewald_h + Grid spacing of the reciprocal part of Ewald sum. Unit: A + ewald_beta + Splitting parameter of the Ewald sum. Unit: A^{-1} + """ + + def __new__( + cls, *args: tuple, model_name: str | None = None, **kwargs: dict + ) -> "DipoleChargeModifier": + return super().__new__(cls, model_name) + + def __init__( + self, + model_name: str, + model_charge_map: list[float], + sys_charge_map: list[float], + ewald_h: float = 1.0, + ewald_beta: float = 1.0, + ) -> None: + """Constructor.""" + super().__init__() + self.modifier_type = "dipole_charge" + self.model_name = model_name + + self.model = torch.jit.load(model_name, map_location=env.DEVICE) + self.rcut = self.model.get_rcut() + self.type_map = self.model.get_type_map() + sel_type = self.model.get_sel_type() + self.sel_type = to_torch_tensor(np.array(sel_type)) + self.model_charge_map = to_torch_tensor(np.array(model_charge_map)) + self.sys_charge_map = to_torch_tensor(np.array(sys_charge_map)) + self._model_charge_map = model_charge_map + self._sys_charge_map = sys_charge_map + + # init ewald recp + self.ewald_h = ewald_h + self.ewald_beta = ewald_beta + self.er = CoulombForceModule( + rcut=self.rcut, + rspace=False, + kappa=ewald_beta, + spacing=ewald_h, + ) + self.placeholder_pairs = torch.ones((1, 2), device=env.DEVICE, dtype=torch.long) + self.placeholder_ds = torch.ones((1), device=env.DEVICE, dtype=torch.float64) + self.placeholder_buffer_scales = torch.zeros( + (1), device=env.DEVICE, dtype=torch.float64 + ) + + def serialize(self) -> dict: + """Serialize the modifier. + + Returns + ------- + dict + The serialized data + """ + data = { + "@class": "Modifier", + "type": self.modifier_type, + "@version": 3, + "model_name": self.model_name, + "model_charge_map": self._model_charge_map, + "sys_charge_map": self._sys_charge_map, + "ewald_h": self.ewald_h, + "ewald_beta": self.ewald_beta, + } + return data + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + """Compute energy, force, and virial corrections for dipole-charge systems. + + This method extends the system with Wannier Function Charge Centers (WFCC) + by adding dipole vectors to atomic coordinates for selected atom types. + It then calculates the electrostatic interactions using Ewald reciprocal + summation to obtain energy, force, and virial corrections. + + Parameters + ---------- + coord : torch.Tensor + The coordinates of atoms with shape (nframes, natoms, 3) + atype : torch.Tensor + The atom types with shape (nframes, natoms) + box : torch.Tensor | None, optional + The simulation box with shape (nframes, 3, 3), by default None + Note: This modifier can only be applied for periodic systems + fparam : torch.Tensor | None, optional + Frame parameters with shape (nframes, nfp), by default None + aparam : torch.Tensor | None, optional + Atom parameters with shape (nframes, natoms, nap), by default None + do_atomic_virial : bool, optional + Whether to compute atomic virial, by default False + + Returns + ------- + dict[str, torch.Tensor] + Dictionary containing the correction terms: + - energy: Energy correction tensor with shape (nframes, 1) + - force: Force correction tensor with shape (nframes, natoms+nsel, 3) + - virial: Virial correction tensor with shape (nframes, 3, 3) + """ + if box is None: + raise RuntimeWarning( + "dipole_charge data modifier can only be applied for periodic systems." + ) + else: + modifier_pred = {} + nframes = coord.shape[0] + natoms = coord.shape[1] + + input_box = box.reshape(nframes, 9) + input_box.requires_grad_(True) + + detached_box = input_box.detach() + sfactor = torch.matmul( + torch.linalg.inv(detached_box.reshape(nframes, 3, 3)), + input_box.reshape(nframes, 3, 3), + ) + input_coord = torch.matmul(coord, sfactor).reshape(nframes, -1) + + extended_coord, extended_charge = self.extend_system( + input_coord, + atype, + input_box, + fparam, + aparam, + ) + + tot_e = [] + # add Ewald reciprocal correction + for ii in range(nframes): + self.er( + extended_coord[ii].reshape((-1, 3)), + input_box[ii].reshape((3, 3)), + self.placeholder_pairs, + self.placeholder_ds, + self.placeholder_buffer_scales, + {"charge": extended_charge[ii].reshape((-1,))}, + ) + tot_e.append(self.er.reciprocal_energy.unsqueeze(0)) + # nframe, + tot_e = torch.concat(tot_e, dim=0) + # nframe, nat * 3 + tot_f = -calc_grads(tot_e, input_coord) + # nframe, nat, 3 + tot_f = torch.reshape(tot_f, (nframes, natoms, 3)) + # nframe, 9 + tot_v = calc_grads(tot_e, input_box) + tot_v = torch.reshape(tot_v, (nframes, 3, 3)) + # nframe, 3, 3 + tot_v = -torch.matmul( + tot_v.transpose(2, 1), input_box.reshape(nframes, 3, 3) + ) + + modifier_pred["energy"] = tot_e + modifier_pred["force"] = tot_f + modifier_pred["virial"] = tot_v + return modifier_pred + + def extend_system( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Extend the system with WFCC (Wannier Function Charge Centers). + + Parameters + ---------- + coord : torch.Tensor + The coordinates of atoms with shape (nframes, natoms * 3) + atype : torch.Tensor + The atom types with shape (nframes, natoms) + box : torch.Tensor + The simulation box with shape (nframes, 9) + fparam : torch.Tensor | None, optional + Frame parameters with shape (nframes, nfp), by default None + aparam : torch.Tensor | None, optional + Atom parameters with shape (nframes, natoms, nap), by default None + + Returns + ------- + tuple + (extended_coord, extended_charge) + extended_coord : torch.Tensor + Extended coordinates with shape (nframes, (natoms + nsel) * 3) + extended_charge : torch.Tensor + Extended charges with shape (nframes, natoms + nsel) + """ + nframes = coord.shape[0] + mask = make_mask(self.sel_type, atype) + + extended_coord = self.extend_system_coord( + coord, + atype, + box, + fparam, + aparam, + ) + # Get ion charges based on atom types + # nframe x nat + ion_charge = self.sys_charge_map[atype] + # Initialize wfcc charges + wc_charge = torch.zeros_like(ion_charge) + # Assign charges to selected atom types + for ii, charge in enumerate(self.model_charge_map): + wc_charge[atype == self.sel_type[ii]] = charge + # Get the charges for selected atoms only + wc_charge_selected = wc_charge[mask].reshape(nframes, -1) + # Concatenate ion charges and wfcc charges + extended_charge = torch.cat([ion_charge, wc_charge_selected], dim=1) + return extended_coord, extended_charge + + def extend_system_coord( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> torch.Tensor: + """Extend the system with WFCC (Wannier Function Charge Centers). + + This function calculates Wannier Function Charge Centers (WFCC) by adding dipole + vectors to atomic coordinates for selected atom types, then concatenates these + WFCC coordinates with the original atomic coordinates. + + Parameters + ---------- + coord : torch.Tensor + The coordinates of atoms with shape (nframes, natoms * 3) + atype : torch.Tensor + The atom types with shape (nframes, natoms) + box : torch.Tensor + The simulation box with shape (nframes, 9) + fparam : torch.Tensor | None, optional + Frame parameters with shape (nframes, nfp), by default None + aparam : torch.Tensor | None, optional + Atom parameters with shape (nframes, natoms, nap), by default None + + Returns + ------- + all_coord : torch.Tensor + Extended coordinates with shape (nframes, (natoms + nsel) * 3) + where nsel is the number of selected atoms + """ + mask = make_mask(self.sel_type, atype) + + nframes = coord.shape[0] + natoms = coord.shape[1] // 3 + + all_dipole = [] + for ii in range(nframes): + dipole_batch = self.model( + coord=coord[ii].reshape(1, -1), + atype=atype[ii].reshape(1, -1), + box=box[ii].reshape(1, -1), + do_atomic_virial=False, + fparam=fparam[ii].reshape(1, -1) if fparam is not None else None, + aparam=aparam[ii].reshape(1, -1) if aparam is not None else None, + ) + # Extract dipole from the output dictionary + all_dipole.append(dipole_batch["dipole"]) + + # nframe x natoms x 3 + dipole = torch.cat(all_dipole, dim=0) + assert dipole.shape[0] == nframes + + dipole_reshaped = dipole.reshape(nframes, natoms, 3) + coord_reshaped = coord.reshape(nframes, natoms, 3) + _wfcc_coord = coord_reshaped + dipole_reshaped + # Apply mask and reshape + wfcc_coord = _wfcc_coord[mask.unsqueeze(-1).expand_as(_wfcc_coord)] + wfcc_coord = wfcc_coord.reshape(nframes, -1) + all_coord = torch.cat((coord, wfcc_coord), dim=1) + return all_coord + + +@torch.jit.export +def make_mask( + sel_type: torch.Tensor, + atype: torch.Tensor, +) -> torch.Tensor: + """Create a boolean mask for selected atom types. + + Parameters + ---------- + sel_type : torch.Tensor + The selected atom types to create a mask for + atype : torch.Tensor + The atom types in the system + + Returns + ------- + mask : torch.Tensor + Boolean mask where True indicates atoms of selected types + """ + # Ensure tensors are of the right type + sel_type = sel_type.to(torch.long) + atype = atype.to(torch.long) + + # Create mask using broadcasting + mask = torch.zeros_like(atype, dtype=torch.bool) + for t in sel_type: + mask = mask | (atype == t) + return mask diff --git a/pyproject.toml b/pyproject.toml index bd403dfaf2..1d65e8e9a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -165,6 +165,7 @@ pin_pytorch_cpu = [ # macos x86 has been deprecated "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'", "torch; platform_machine=='x86_64' and platform_system == 'Darwin'", + "torch_admp @ git+https://github.com/chiahsinchu/torch-admp.git@v1.1.0a", ] pin_pytorch_gpu = [ "torch==2.10.0", diff --git a/source/tests/pt/modifier/__init__.py b/source/tests/pt/modifier/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/pt/modifier/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/pt/test_data_modifier.py b/source/tests/pt/modifier/test_data_modifier.py similarity index 99% rename from source/tests/pt/test_data_modifier.py rename to source/tests/pt/modifier/test_data_modifier.py index 18d66ef2ff..18c909af65 100644 --- a/source/tests/pt/test_data_modifier.py +++ b/source/tests/pt/modifier/test_data_modifier.py @@ -52,7 +52,7 @@ DeepmdData, ) -from ..consistent.common import ( +from ...consistent.common import ( parameterized, ) diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py new file mode 100644 index 0000000000..263453e72c --- /dev/null +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import unittest +from pathlib import ( + Path, +) + +import numpy as np +import torch + +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) +from deepmd.pt.entrypoints.main import ( + freeze, + get_trainer, +) +from deepmd.pt.modifier import DipoleChargeModifier as PTDipoleChargeModifier +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) +from deepmd.tf.modifier import DipoleChargeModifier as TFDipoleChargeModifier + +SEED = 1 +DTYPE = torch.float64 + + +def ref_data(): + all_box = np.load(str(Path(__file__).parent / "water/data/data_0/set.000/box.npy")) + all_coord = np.load( + str(Path(__file__).parent / "water/data/data_0/set.000/coord.npy") + ) + nframe = len(all_box) + rng = np.random.default_rng(SEED) + selected_id = rng.integers(nframe) + + coord = all_coord[selected_id].reshape(1, -1) + box = all_box[selected_id].reshape(1, -1) + atype = np.loadtxt( + str(Path(__file__).parent / "water/data/data_0/type.raw"), + dtype=int, + ).reshape(1, -1) + return coord, box, atype + + +class TestDipoleChargeModifier(unittest.TestCase): + def setUp(self) -> None: + # setup parameter + # numerical consistency can only be achieved with high prec + self.ewald_h = 0.1 + self.ewald_beta = 0.5 + self.model_charge_map = [-8.0] + self.sys_charge_map = [6.0, 1.0] + self.descriptor_dict = { + "type": "se_e2_a", + "sel": [46, 92], + "rcut_smth": 0.5, + "rcut": 4.00, + "neuron": [ + 25, + 50, + ], + } + + # Train DW model + input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") + with open(input_json, encoding="utf-8") as f: + config = json.load(f) + config["model"]["descriptor"].update(self.descriptor_dict) + config["training"]["numb_steps"] = 1 + config["training"]["save_freq"] = 1 + config["learning_rate"]["start_lr"] = 1.0 + config["training"]["training_data"]["systems"] = [ + str(Path(__file__).parent / "water_tensor/dipole/O78H156"), + ] + config["training"]["validation_data"]["systems"] = [ + str(Path(__file__).parent / "water_tensor/dipole/O96H192") + ] + + trainer = get_trainer(config) + trainer.run() + freeze( + model="model.ckpt.pt", + output="dw_model.pth", + head=None, + ) + # Convert pb model to pth model + convert_backend(INPUT="dw_model.pth", OUTPUT="dw_model.pb") + + self.dm_pt = PTDipoleChargeModifier( + "dw_model.pth", + self.model_charge_map, + self.sys_charge_map, + self.ewald_h, + self.ewald_beta, + ) + self.dm_tf = TFDipoleChargeModifier( + "dw_model.pb", + self.model_charge_map, + self.sys_charge_map, + self.ewald_h, + self.ewald_beta, + ) + + def test_jit(self): + torch.jit.script(self.dm_pt) + + def test_consistency(self): + coord, box, atype = ref_data() + # consistent with the input shape from BaseModifier.modify_data + t_coord = to_torch_tensor(coord).to(DTYPE).reshape(1, -1, 3) + t_box = to_torch_tensor(box).to(DTYPE).reshape(1, 3, 3) + t_atype = to_torch_tensor(atype).to(torch.long) + + dm_pred = self.dm_pt( + coord=t_coord, + atype=t_atype, + box=t_box, + ) + e, f, v = self.dm_tf.eval( + coord=coord, + box=box, + atype=atype.reshape(-1), + ) + + np.testing.assert_allclose( + to_numpy_array(dm_pred["energy"]).reshape(-1), e.reshape(-1), rtol=1e-4 + ) + np.testing.assert_allclose( + to_numpy_array(dm_pred["force"]).reshape(-1), f.reshape(-1), rtol=1e-4 + ) + np.testing.assert_allclose( + to_numpy_array(dm_pred["virial"]).reshape(-1), v.reshape(-1), rtol=1e-4 + ) + + def test_serialize(self): + """Test the serialize method of DipoleChargeModifier.""" + coord, box, atype = ref_data() + # consistent with the input shape from BaseModifier.modify_data + t_coord = to_torch_tensor(coord).to(DTYPE).reshape(1, -1, 3) + t_box = to_torch_tensor(box).to(DTYPE).reshape(1, 3, 3) + t_atype = to_torch_tensor(atype).to(torch.long) + + dm0 = self.dm_pt.to(env.DEVICE) + dm1 = PTDipoleChargeModifier.deserialize(dm0.serialize()).to(env.DEVICE) + + ret0 = dm0( + coord=t_coord, + atype=t_atype, + box=t_box, + ) + ret1 = dm1( + coord=t_coord, + atype=t_atype, + box=t_box, + ) + + np.testing.assert_allclose( + to_numpy_array(ret0["energy"]), to_numpy_array(ret1["energy"]) + ) + np.testing.assert_allclose( + to_numpy_array(ret0["force"]), to_numpy_array(ret1["force"]) + ) + np.testing.assert_allclose( + to_numpy_array(ret0["virial"]), to_numpy_array(ret1["virial"]) + ) + + def test_box_none_warning(self): + """Test that a RuntimeWarning is raised when box is None.""" + coord, box, atype = ref_data() + # consistent with the input shape from BaseModifier.modify_data + t_coord = to_torch_tensor(coord).to(DTYPE).reshape(1, -1, 3) + t_atype = to_torch_tensor(atype).to(torch.long) + + with self.assertRaises(RuntimeWarning) as context: + self.dm_pt( + coord=t_coord, + atype=t_atype, + box=None, # Pass None to trigger the warning + ) + + self.assertIn( + "dipole_charge data modifier can only be applied for periodic systems", + str(context.exception), + ) + + def test_train(self): + input_json = str(Path(__file__).parent / "water/se_e2_a.json") + with open(input_json, encoding="utf-8") as f: + config = json.load(f) + config["model"]["descriptor"].update(self.descriptor_dict) + config["training"]["save_freq"] = 1 + config["learning_rate"]["start_lr"] = 1.0 + config["training"]["training_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/data_0"), + str(Path(__file__).parent / "water/data/data_1"), + ] + config["training"]["validation_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single"), + ] + config["training"]["numb_steps"] = 1 + + trainer = get_trainer(config) + trainer.run() + + def tearDown(self) -> None: + for f in os.listdir("."): + if f.startswith("frozen_model") and f.endswith(".pth"): + os.remove(f) + if f.startswith("dw_model") and (f.endswith(".pth") or f.endswith(".pb")): + os.remove(f) + if f.startswith("model.ckpt") and f.endswith(".pt"): + os.remove(f) + if f in ["lcurve.out", "checkpoint"]: + os.remove(f) diff --git a/source/tests/pt/modifier/water b/source/tests/pt/modifier/water new file mode 120000 index 0000000000..b4ce4e224a --- /dev/null +++ b/source/tests/pt/modifier/water @@ -0,0 +1 @@ +../model/water/ \ No newline at end of file diff --git a/source/tests/pt/modifier/water_tensor b/source/tests/pt/modifier/water_tensor new file mode 120000 index 0000000000..a8c63dbb30 --- /dev/null +++ b/source/tests/pt/modifier/water_tensor @@ -0,0 +1 @@ +../water_tensor/ \ No newline at end of file From bc98d0e4ec9ad6da2b92eb5a02960a1b62dc279d Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 8 Jan 2026 12:54:24 +0800 Subject: [PATCH 02/25] update torch_admp version; minor update in dipole_charge non-pbc assertion --- deepmd/pt/modifier/dipole_charge.py | 2 +- pyproject.toml | 2 +- source/tests/pt/modifier/test_dipole_charge.py | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index 3ab1941c97..587161e247 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -139,7 +139,7 @@ def forward( - virial: Virial correction tensor with shape (nframes, 3, 3) """ if box is None: - raise RuntimeWarning( + raise RuntimeError( "dipole_charge data modifier can only be applied for periodic systems." ) else: diff --git a/pyproject.toml b/pyproject.toml index 1d65e8e9a9..63b648fdbb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -165,7 +165,7 @@ pin_pytorch_cpu = [ # macos x86 has been deprecated "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'", "torch; platform_machine=='x86_64' and platform_system == 'Darwin'", - "torch_admp @ git+https://github.com/chiahsinchu/torch-admp.git@v1.1.0a", + "torch_admp @ git+https://github.com/chiahsinchu/torch-admp.git@v1.1.1", ] pin_pytorch_gpu = [ "torch==2.10.0", diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index 263453e72c..4a26d3523f 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -171,17 +171,17 @@ def test_serialize(self): ) def test_box_none_warning(self): - """Test that a RuntimeWarning is raised when box is None.""" - coord, box, atype = ref_data() + """Test that a RuntimeError is raised when box is None.""" + coord, _b, atype = ref_data() # consistent with the input shape from BaseModifier.modify_data t_coord = to_torch_tensor(coord).to(DTYPE).reshape(1, -1, 3) t_atype = to_torch_tensor(atype).to(torch.long) - with self.assertRaises(RuntimeWarning) as context: + with self.assertRaises(RuntimeError) as context: self.dm_pt( coord=t_coord, atype=t_atype, - box=None, # Pass None to trigger the warning + box=None, # Pass None to trigger the error ) self.assertIn( From 284ecf25e15d250b46ca9ea51ec1aeb78e39c24d Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 8 Jan 2026 13:57:00 +0800 Subject: [PATCH 03/25] minor code improvement based on @coderabbitai --- deepmd/pt/modifier/dipole_charge.py | 8 ++++++-- source/tests/pt/modifier/test_dipole_charge.py | 17 +++++++---------- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index 587161e247..f412565228 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -129,13 +129,14 @@ def forward( Atom parameters with shape (nframes, natoms, nap), by default None do_atomic_virial : bool, optional Whether to compute atomic virial, by default False + Note: This parameter is currently not implemented and is ignored Returns ------- dict[str, torch.Tensor] Dictionary containing the correction terms: - energy: Energy correction tensor with shape (nframes, 1) - - force: Force correction tensor with shape (nframes, natoms+nsel, 3) + - force: Force correction tensor with shape (nframes, natoms, 3) - virial: Virial correction tensor with shape (nframes, 3, 3) """ if box is None: @@ -305,7 +306,10 @@ def extend_system_coord( # nframe x natoms x 3 dipole = torch.cat(all_dipole, dim=0) - assert dipole.shape[0] == nframes + if dipole.shape[0] != nframes: + raise RuntimeError( + f"Dipole shape mismatch: expected {nframes} frames, got {dipole.shape[0]}" + ) dipole_reshaped = dipole.reshape(nframes, natoms, 3) coord_reshaped = coord.reshape(nframes, natoms, 3) diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index 4a26d3523f..07da63b9bd 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json import os +import tempfile import unittest from pathlib import ( Path, @@ -50,6 +51,9 @@ def ref_data(): class TestDipoleChargeModifier(unittest.TestCase): def setUp(self) -> None: + self.test_dir = tempfile.TemporaryDirectory() + self.orig_dir = os.getcwd() + os.chdir(self.test_dir.name) # setup parameter # numerical consistency can only be achieved with high prec self.ewald_h = 0.1 @@ -170,7 +174,7 @@ def test_serialize(self): to_numpy_array(ret0["virial"]), to_numpy_array(ret1["virial"]) ) - def test_box_none_warning(self): + def test_box_none_error(self): """Test that a RuntimeError is raised when box is None.""" coord, _b, atype = ref_data() # consistent with the input shape from BaseModifier.modify_data @@ -209,12 +213,5 @@ def test_train(self): trainer.run() def tearDown(self) -> None: - for f in os.listdir("."): - if f.startswith("frozen_model") and f.endswith(".pth"): - os.remove(f) - if f.startswith("dw_model") and (f.endswith(".pth") or f.endswith(".pb")): - os.remove(f) - if f.startswith("model.ckpt") and f.endswith(".pt"): - os.remove(f) - if f in ["lcurve.out", "checkpoint"]: - os.remove(f) + os.chdir(self.orig_dir) + self.test_dir.cleanup() From 1f53cfc7a4665e73c90abaa4ceb0ac0de0b7dc5b Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 8 Jan 2026 18:27:43 +0800 Subject: [PATCH 04/25] change source of torch_admp from github to pypi in pyproject.toml; minor code improvement --- deepmd/pt/modifier/dipole_charge.py | 13 +++++++++---- pyproject.toml | 2 +- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index f412565228..ac4552e1c7 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -63,6 +63,13 @@ def __init__( self._model_charge_map = model_charge_map self._sys_charge_map = sys_charge_map + # Validate that model_charge_map and sel_type have matching lengths + if len(model_charge_map) != len(sel_type): + raise ValueError( + f"model_charge_map length ({len(model_charge_map)}) must match " + f"sel_type length ({len(sel_type)})" + ) + # init ewald recp self.ewald_h = ewald_h self.ewald_beta = ewald_beta @@ -344,8 +351,6 @@ def make_mask( sel_type = sel_type.to(torch.long) atype = atype.to(torch.long) - # Create mask using broadcasting - mask = torch.zeros_like(atype, dtype=torch.bool) - for t in sel_type: - mask = mask | (atype == t) + # Create mask using broadcasting for JIT compatibility + mask = (atype.unsqueeze(-1) == sel_type).any(dim=-1) return mask diff --git a/pyproject.toml b/pyproject.toml index 63b648fdbb..18b082771d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -165,7 +165,7 @@ pin_pytorch_cpu = [ # macos x86 has been deprecated "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'", "torch; platform_machine=='x86_64' and platform_system == 'Darwin'", - "torch_admp @ git+https://github.com/chiahsinchu/torch-admp.git@v1.1.1", + "torch-admp==1.1.1", ] pin_pytorch_gpu = [ "torch==2.10.0", From 9be128efc33846e013862d911bff62755b83a56a Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 10 Jan 2026 14:38:45 +0800 Subject: [PATCH 05/25] add torch-admp to dependencies --- backend/dynamic_metadata.py | 5 ++++- backend/find_pytorch.py | 10 +++++++++- pyproject.toml | 3 +++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index e7763cac84..ab6acc5d78 100644 --- a/backend/dynamic_metadata.py +++ b/backend/dynamic_metadata.py @@ -48,8 +48,11 @@ def dynamic_metadata( ] optional_dependencies["lmp"].extend(find_libpython_requires) optional_dependencies["ipi"].extend(find_libpython_requires) + torch_static_requirement = optional_dependencies.pop("torch", None) return { **optional_dependencies, **get_tf_requirement(tf_version), - **get_pt_requirement(pt_version), + **get_pt_requirement( + pt_version, static_requirement=torch_static_requirement + ), } diff --git a/backend/find_pytorch.py b/backend/find_pytorch.py index d50f57bf5e..4140301f34 100644 --- a/backend/find_pytorch.py +++ b/backend/find_pytorch.py @@ -90,7 +90,10 @@ def find_pytorch() -> tuple[str | None, list[str]]: @lru_cache -def get_pt_requirement(pt_version: str = "") -> dict: +def get_pt_requirement( + pt_version: str = "", + static_requirement: list[str] | None = None, +) -> dict: """Get PyTorch requirement when PT is not installed. If pt_version is not given and the environment variable `PYTORCH_VERSION` is set, use it as the requirement. @@ -99,6 +102,8 @@ def get_pt_requirement(pt_version: str = "") -> dict: ---------- pt_version : str, optional PT version + static_requirement : list[str] or None, optional + Static requirements Returns ------- @@ -125,6 +130,8 @@ def get_pt_requirement(pt_version: str = "") -> dict: mpi_requirement = ["mpich"] else: mpi_requirement = [] + if static_requirement is None: + static_requirement = [] return { "torch": [ @@ -138,6 +145,7 @@ def get_pt_requirement(pt_version: str = "") -> dict: else "torch>=2.1.0", *mpi_requirement, *cibw_requirement, + *static_requirement, ], } diff --git a/pyproject.toml b/pyproject.toml index 18b082771d..06c5539be0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,6 +140,9 @@ jax = [ # The pinning of ml_dtypes may conflict with TF # 'jax-ai-stack;python_version>="3.10"', ] +torch = [ + "torch-admp", +] [tool.deepmd_build_backend.scripts] dp = "deepmd.main:main" From b97aa84d40889dd151bb56312c6f2897b5c018af Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 10 Jan 2026 15:07:50 +0800 Subject: [PATCH 06/25] pass tuple --- backend/dynamic_metadata.py | 2 +- backend/find_pytorch.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index ab6acc5d78..8ba532d637 100644 --- a/backend/dynamic_metadata.py +++ b/backend/dynamic_metadata.py @@ -53,6 +53,6 @@ def dynamic_metadata( **optional_dependencies, **get_tf_requirement(tf_version), **get_pt_requirement( - pt_version, static_requirement=torch_static_requirement + pt_version, static_requirement=tuple(torch_static_requirement) ), } diff --git a/backend/find_pytorch.py b/backend/find_pytorch.py index 4140301f34..9fedf000b5 100644 --- a/backend/find_pytorch.py +++ b/backend/find_pytorch.py @@ -92,7 +92,7 @@ def find_pytorch() -> tuple[str | None, list[str]]: @lru_cache def get_pt_requirement( pt_version: str = "", - static_requirement: list[str] | None = None, + static_requirement: tuple[str] | None = None, ) -> dict: """Get PyTorch requirement when PT is not installed. @@ -102,7 +102,7 @@ def get_pt_requirement( ---------- pt_version : str, optional PT version - static_requirement : list[str] or None, optional + static_requirement : tuple[str] or None, optional Static requirements Returns @@ -131,7 +131,7 @@ def get_pt_requirement( else: mpi_requirement = [] if static_requirement is None: - static_requirement = [] + static_requirement = () return { "torch": [ From f97f298fd76f0697b9fd7409fe630e851043d969 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 10 Jan 2026 15:58:40 +0800 Subject: [PATCH 07/25] Fix pop method for torch_static_requirement Signed-off-by: Jinzhe Zeng --- backend/dynamic_metadata.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index 8ba532d637..d5f7b370b5 100644 --- a/backend/dynamic_metadata.py +++ b/backend/dynamic_metadata.py @@ -48,7 +48,7 @@ def dynamic_metadata( ] optional_dependencies["lmp"].extend(find_libpython_requires) optional_dependencies["ipi"].extend(find_libpython_requires) - torch_static_requirement = optional_dependencies.pop("torch", None) + torch_static_requirement = optional_dependencies.pop("torch", ()) return { **optional_dependencies, **get_tf_requirement(tf_version), From 860d2bb6814e0c98d1cf0899aa6505de3067602c Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sat, 10 Jan 2026 17:17:18 +0800 Subject: [PATCH 08/25] docs: add PyTorch backend documentation to DPLR model --- deepmd/utils/argcheck.py | 2 +- doc/model/dplr.md | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 8c20bb8bf4..c82b24d52a 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -2310,7 +2310,7 @@ def model_args(exclude_hybrid: bool = False) -> list[Argument]: [], [modifier_variant_type_args()], optional=True, - doc=doc_only_tf_supported + doc_modifier, + doc=doc_modifier, ), Argument( "compress", diff --git a/doc/model/dplr.md b/doc/model/dplr.md index 61327bb55e..2e406a0167 100644 --- a/doc/model/dplr.md +++ b/doc/model/dplr.md @@ -1,10 +1,10 @@ # Deep potential long-range (DPLR) {{ tensorflow_icon }} :::{note} -**Supported backends**: TensorFlow {{ tensorflow_icon }} +**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} ::: -Notice: **The interfaces of DPLR are not stable and subject to change** +Notice: **The interfaces of DPLR are not stable and subject to change. In addition, DP/LAMMPS interface does not yet support PyTorch DPLR models. Use `dp convert` to convert your model to TensorFlow backend for LAMMPS simulations.** The method of DPLR is described in [this paper][1]. One is recommended to read the paper before using the DPLR. From 867ce69f2577b2ebc7e5c991dcf0f0d7da8a4220 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 15 Jan 2026 18:34:04 +0800 Subject: [PATCH 09/25] update freeze/eval with model (de-)serialization --- deepmd/dpmodel/modifier/base_modifier.py | 1 - deepmd/pt/entrypoints/main.py | 18 ++--- deepmd/pt/infer/deep_eval.py | 18 +++-- deepmd/pt/modifier/base_modifier.py | 1 + .../tests/pt/modifier/test_data_modifier.py | 68 +++++++++++++++++-- 5 files changed, 80 insertions(+), 26 deletions(-) diff --git a/deepmd/dpmodel/modifier/base_modifier.py b/deepmd/dpmodel/modifier/base_modifier.py index b3e3c2ca7b..641e739ae4 100644 --- a/deepmd/dpmodel/modifier/base_modifier.py +++ b/deepmd/dpmodel/modifier/base_modifier.py @@ -36,7 +36,6 @@ def serialize(self) -> dict: dict The serialized data """ - pass @classmethod def deserialize(cls, data: dict) -> "BaseModifier": diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 46ad8a6cd0..0f68052287 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import argparse import copy -import io import json import logging import os +import pickle from pathlib import ( Path, ) @@ -401,17 +401,11 @@ def freeze( model.eval() model = torch.jit.script(model) - dm_output = "data_modifier.pth" - extra_files = {dm_output: ""} - if tester.modifier is not None: - dm = tester.modifier - dm.eval() - buffer = io.BytesIO() - torch.jit.save( - torch.jit.script(dm), - buffer, - ) - extra_files = {dm_output: buffer.getvalue()} + extra_files = {"modifier_data": ""} + dm = tester.modifier + if dm is not None: + bytes_data = pickle.dumps(dm.serialize()) + extra_files = {"modifier_data": bytes_data} torch.jit.save( model, output, diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 6e63ecb2fc..b909f5416d 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import io import json import logging +import pickle from collections.abc import ( Callable, ) @@ -49,6 +49,9 @@ from deepmd.pt.model.network.network import ( TypeEmbedNetConsistent, ) +from deepmd.pt.modifier import ( + BaseModifier, +) from deepmd.pt.train.wrapper import ( ModelWrapper, ) @@ -172,19 +175,20 @@ def __init__( self.dp = ModelWrapper(model) self.dp.load_state_dict(state_dict) elif str(self.model_path).endswith(".pth"): - extra_files = {"data_modifier.pth": ""} + extra_files = {"modifier_data": ""} model = torch.jit.load( model_file, map_location=env.DEVICE, _extra_files=extra_files ) modifier = None # Load modifier if it exists in extra_files - if len(extra_files["data_modifier.pth"]) > 0: - # Create a file-like object from the in-memory data - modifier_data = extra_files["data_modifier.pth"] + if len(extra_files["modifier_data"]) > 0: + modifier_data = extra_files["modifier_data"] if isinstance(modifier_data, bytes): - modifier_data = io.BytesIO(modifier_data) + modifier_data = pickle.loads(modifier_data) # Load the modifier directly from the file-like object - modifier = torch.jit.load(modifier_data, map_location=env.DEVICE) + modifier = BaseModifier.get_class_by_type( + modifier_data["type"] + ).deserialize(modifier_data) self.dp = ModelWrapper(model, modifier=modifier) self.modifier = modifier model_def_script = self.dp.model["Default"].get_model_def_script() diff --git a/deepmd/pt/modifier/base_modifier.py b/deepmd/pt/modifier/base_modifier.py index 5a8c6538b0..db37694305 100644 --- a/deepmd/pt/modifier/base_modifier.py +++ b/deepmd/pt/modifier/base_modifier.py @@ -47,6 +47,7 @@ def serialize(self) -> dict: data = { "@class": "Modifier", "type": self.modifier_type, + "use_cache": self.use_cache, "@version": 3, } return data diff --git a/source/tests/pt/modifier/test_data_modifier.py b/source/tests/pt/modifier/test_data_modifier.py index 18c909af65..ddba54cd05 100644 --- a/source/tests/pt/modifier/test_data_modifier.py +++ b/source/tests/pt/modifier/test_data_modifier.py @@ -36,12 +36,18 @@ freeze, get_trainer, ) +from deepmd.pt.model.model import ( + EnergyModel, +) from deepmd.pt.modifier.base_modifier import ( BaseModifier, ) from deepmd.pt.utils import ( env, ) +from deepmd.pt.utils.serialization import ( + serialize_from_file, +) from deepmd.pt.utils.utils import ( to_numpy_array, ) @@ -85,7 +91,9 @@ def modifier_scaling_tester() -> list[Argument]: doc_sfactor = "The scaling factor for correction." doc_use_cache = "Whether to cache modified frames to improve performance by avoiding recomputation." return [ - Argument("model_name", str, optional=False, doc=doc_model_name), + Argument( + "model_name", str, alias=["model"], optional=False, doc=doc_model_name + ), Argument("sfactor", float, optional=False, doc=doc_sfactor), Argument("use_cache", bool, optional=True, doc=doc_use_cache), ] @@ -181,21 +189,69 @@ def modify_data(self, data: dict[str, Array | float], data_sys: DeepmdData) -> N @BaseModifier.register("scaling_tester") class ModifierScalingTester(BaseModifier): - def __new__(cls, *args, **kwargs): - return super().__new__(cls) + def __new__( + cls, + *args: tuple, + model: str | None = None, + model_name: str | None = None, + **kwargs: dict, + ) -> "ModifierScalingTester": + return super().__new__(cls, model_name if model_name is not None else model) def __init__( self, - model_name: str, + model: torch.nn.Module | None = None, + model_name: str | None = None, sfactor: float = 1.0, use_cache: bool = True, ) -> None: """Initialize a test modifier that applies scaled model predictions using a frozen model.""" super().__init__(use_cache) self.modifier_type = "scaling_tester" - self.model_name = model_name self.sfactor = sfactor - self.model = torch.jit.load(model_name, map_location=env.DEVICE) + + if model_name is None and model is None: + raise AttributeError("`model_name` or `model` should be specified.") + if model_name is not None and model is not None: + raise AttributeError( + "`model_name` and `model` cannot be used simultaneously." + ) + + if model is not None: + self._model = model.to(env.DEVICE) + if model_name is not None: + data = serialize_from_file(model_name) + self._model = EnergyModel.deserialize(data["model"]).to(env.DEVICE) + + # use jit model for inference + self.model = torch.jit.script(self._model) + + def serialize(self) -> dict: + """Serialize the modifier. + + Returns + ------- + dict + The serialized data + """ + dd = BaseModifier.serialize(self) + dd.update( + { + "model": self._model.serialize(), + "sfactor": self.sfactor, + } + ) + return dd + + @classmethod + def deserialize(cls, data: dict) -> "ModifierScalingTester": + data = data.copy() + data.pop("@class", None) + data.pop("type", None) + data.pop("@version", None) + model_obj = EnergyModel.deserialize(data.pop("model")) + data["model"] = model_obj + return cls(**data) def forward( self, From d625fe79eb7785086bfa12f2e4f4f5176de7aca2 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 15 Jan 2026 18:48:52 +0800 Subject: [PATCH 10/25] Feat: Refactor DipoleChargeModifier to support direct model object input and improve serialization. Added support for both model file path and direct model object initialization, along with proper serialize/deserialize methods and caching functionality. --- deepmd/pt/modifier/dipole_charge.py | 62 ++++++++++++++++++++++------- doc/model/dplr.md | 42 ++++++++++++++++++- 2 files changed, 87 insertions(+), 17 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index ac4552e1c7..b4fc462c0e 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -8,12 +8,18 @@ calc_grads, ) +from deepmd.pt.model.model import ( + DipoleModel, +) from deepmd.pt.modifier.base_modifier import ( BaseModifier, ) from deepmd.pt.utils import ( env, ) +from deepmd.pt.utils.serialization import ( + serialize_from_file, +) from deepmd.pt.utils.utils import ( to_torch_tensor, ) @@ -42,18 +48,33 @@ def __new__( def __init__( self, - model_name: str, + model_name: str | None, model_charge_map: list[float], sys_charge_map: list[float], ewald_h: float = 1.0, ewald_beta: float = 1.0, + model: DipoleModel | None = None, + use_cache: bool = True, ) -> None: """Constructor.""" - super().__init__() + super().__init__(use_cache=use_cache) self.modifier_type = "dipole_charge" - self.model_name = model_name - self.model = torch.jit.load(model_name, map_location=env.DEVICE) + if model_name is None and model is None: + raise AttributeError("`model_name` or `model` should be specified.") + if model_name is not None and model is not None: + raise AttributeError( + "`model_name` and `model` cannot be used simultaneously." + ) + + if model is not None: + self._model = model.to(env.DEVICE) + if model_name is not None: + data = serialize_from_file(model_name) + self._model = DipoleModel.deserialize(data["model"]).to(env.DEVICE) + + # use jit model for inference + self.model = torch.jit.script(self._model) self.rcut = self.model.get_rcut() self.type_map = self.model.get_type_map() sel_type = self.model.get_sel_type() @@ -93,17 +114,28 @@ def serialize(self) -> dict: dict The serialized data """ - data = { - "@class": "Modifier", - "type": self.modifier_type, - "@version": 3, - "model_name": self.model_name, - "model_charge_map": self._model_charge_map, - "sys_charge_map": self._sys_charge_map, - "ewald_h": self.ewald_h, - "ewald_beta": self.ewald_beta, - } - return data + dd = BaseModifier.serialize(self) + dd.update( + { + "model": self._model.serialize(), + "model_charge_map": self._model_charge_map, + "sys_charge_map": self._sys_charge_map, + "ewald_h": self.ewald_h, + "ewald_beta": self.ewald_beta, + } + ) + return dd + + @classmethod + def deserialize(cls, data: dict) -> "DipoleChargeModifier": + data = data.copy() + data.pop("@class", None) + data.pop("type", None) + data.pop("@version", None) + model_obj = DipoleModel.deserialize(data.pop("model")) + data["model"] = model_obj + data["model_name"] = None + return cls(**data) def forward( self, diff --git a/doc/model/dplr.md b/doc/model/dplr.md index 2e406a0167..d262b778e4 100644 --- a/doc/model/dplr.md +++ b/doc/model/dplr.md @@ -1,10 +1,10 @@ -# Deep potential long-range (DPLR) {{ tensorflow_icon }} +# Deep potential long-range (DPLR) {{ tensorflow_icon }} {{ pytorch_icon }} :::{note} **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} ::: -Notice: **The interfaces of DPLR are not stable and subject to change. In addition, DP/LAMMPS interface does not yet support PyTorch DPLR models. Use `dp convert` to convert your model to TensorFlow backend for LAMMPS simulations.** +Notice: **The interfaces of DPLR are not stable and subject to change. In addition, DP/LAMMPS interface does not yet support PyTorch DPLR models. Use `dp convert-backend` to convert your model to TensorFlow backend for LAMMPS simulations (see [details](#model-conversion)).** The method of DPLR is described in [this paper][1]. One is recommended to read the paper before using the DPLR. @@ -266,4 +266,42 @@ The MD simulation lasts for only 20 steps. If one runs a longer simulation, it w Another restriction that should be noted is that the energies printed at the zero steps are not correct. This is because at the zero steps the position of the WC has not been updated with the DW model. The energies printed in later steps are correct. +## Model conversion + +### Converting from PyTorch to TensorFlow + +To use DPLR models with LAMMPS, you need to convert them from PyTorch to TensorFlow backend. The conversion process involves the following steps: + +1. **Convert the Deep Wannier (DW) model:** + + ```bash + dp convert-backend dw_model.pth dw_model.pb + ``` + + This converts the PyTorch DW model (`.pth` file) to TensorFlow format (`.pb` file). + +2. **Convert the DPLR energy model:** + + ```bash + dp convert-backend dplr_energy_model.pth dplr_energy_model.pb + ``` + + This converts the PyTorch DPLR energy model (short-range contribution) to TensorFlow backend. + +3. **Combine the models using TensorFlow training:** + ```bash + dp --tf train input.json --init-frz-model dplr_energy_model.pb + ``` + Run TensorFlow training with the DPLR TensorFlow setup, setting the training steps to 0 to combine the models without additional training: + ```json + { + "training": { + "numb_steps": 0, + ... + } + } + ``` + +This process ensures that both the DW model and the DPLR energy model are properly converted and combined for use with LAMMPS simulations. + [1]: https://arxiv.org/abs/2112.13327 From 09183e461c0608b26724ad13f9b370d145d37bae Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Fri, 16 Jan 2026 10:56:30 +0800 Subject: [PATCH 11/25] feat: Add batch processing support to DipoleChargeModifier for improved memory efficiency and performance. This commit introduces configurable batch sizes for both Ewald calculations and dipole model inference, refactors the system extension logic to return atomic dipoles, and optimizes memory usage during large-scale simulations. --- deepmd/pt/modifier/dipole_charge.py | 144 +++++++++--------- .../tests/pt/modifier/test_dipole_charge.py | 61 +++++--- 2 files changed, 111 insertions(+), 94 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index b4fc462c0e..a72bf09420 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -1,4 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import os + import numpy as np import torch from torch_admp.pme import ( @@ -53,6 +55,8 @@ def __init__( sys_charge_map: list[float], ewald_h: float = 1.0, ewald_beta: float = 1.0, + ewald_batch_size: int = 5, + dp_batch_size: int | None = None, model: DipoleModel | None = None, use_cache: bool = True, ) -> None: @@ -72,6 +76,7 @@ def __init__( if model_name is not None: data = serialize_from_file(model_name) self._model = DipoleModel.deserialize(data["model"]).to(env.DEVICE) + self._model.eval() # use jit model for inference self.model = torch.jit.script(self._model) @@ -106,6 +111,11 @@ def __init__( (1), device=env.DEVICE, dtype=torch.float64 ) + self.ewald_batch_size = ewald_batch_size + if dp_batch_size is None: + dp_batch_size = int(os.environ.get("DP_INFER_BATCH_SIZE", 1)) + self.dp_batch_size = dp_batch_size + def serialize(self) -> dict: """Serialize the modifier. @@ -122,6 +132,8 @@ def serialize(self) -> dict: "sys_charge_map": self._sys_charge_map, "ewald_h": self.ewald_h, "ewald_beta": self.ewald_beta, + "ewald_batch_size": self.ewald_batch_size, + "dp_batch_size": self.dp_batch_size, } ) return dd @@ -192,12 +204,12 @@ def forward( detached_box = input_box.detach() sfactor = torch.matmul( - torch.linalg.inv(detached_box.reshape(nframes, 3, 3)), + torch.inverse(detached_box.reshape(nframes, 3, 3)), input_box.reshape(nframes, 3, 3), ) input_coord = torch.matmul(coord, sfactor).reshape(nframes, -1) - extended_coord, extended_charge = self.extend_system( + extended_coord, extended_charge, _atomic_dipole = self.extend_system( input_coord, atype, input_box, @@ -205,16 +217,25 @@ def forward( aparam, ) - tot_e = [] # add Ewald reciprocal correction - for ii in range(nframes): + tot_e: list[torch.Tensor] = [] + chunk_coord = torch.split( + extended_coord.reshape(nframes, -1, 3), self.dp_batch_size, dim=0 + ) + chunk_box = torch.split( + input_box.reshape(nframes, 3, 3), self.dp_batch_size, dim=0 + ) + chunk_charge = torch.split( + extended_charge.reshape(nframes, -1), self.dp_batch_size, dim=0 + ) + for _coord, _box, _charge in zip(chunk_coord, chunk_box, chunk_charge): self.er( - extended_coord[ii].reshape((-1, 3)), - input_box[ii].reshape((3, 3)), + _coord, + _box, self.placeholder_pairs, self.placeholder_ds, self.placeholder_buffer_scales, - {"charge": extended_charge[ii].reshape((-1,))}, + {"charge": _charge}, ) tot_e.append(self.er.reciprocal_energy.unsqueeze(0)) # nframe, @@ -243,7 +264,7 @@ def extend_system( box: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Extend the system with WFCC (Wannier Function Charge Centers). Parameters @@ -261,17 +282,17 @@ def extend_system( Returns ------- - tuple - (extended_coord, extended_charge) - extended_coord : torch.Tensor - Extended coordinates with shape (nframes, (natoms + nsel) * 3) - extended_charge : torch.Tensor - Extended charges with shape (nframes, natoms + nsel) + Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + A tuple containing three tensors: + - extended_coord : torch.Tensor + Extended coordinates with shape (nframes, 2 * natoms * 3) + - extended_charge : torch.Tensor + Extended charges with shape (nframes, 2 * natoms) + - atomic_dipole : torch.Tensor + Atomic dipoles with shape (nframes, natoms, 3) """ - nframes = coord.shape[0] - mask = make_mask(self.sel_type, atype) - - extended_coord = self.extend_system_coord( + # nframes, natoms, 3 + extended_coord, atomic_dipole = self.extend_system_coord( coord, atype, box, @@ -286,11 +307,9 @@ def extend_system( # Assign charges to selected atom types for ii, charge in enumerate(self.model_charge_map): wc_charge[atype == self.sel_type[ii]] = charge - # Get the charges for selected atoms only - wc_charge_selected = wc_charge[mask].reshape(nframes, -1) # Concatenate ion charges and wfcc charges - extended_charge = torch.cat([ion_charge, wc_charge_selected], dim=1) - return extended_coord, extended_charge + extended_charge = torch.cat([ion_charge, wc_charge], dim=1) + return extended_coord, extended_charge, atomic_dipole def extend_system_coord( self, @@ -299,7 +318,7 @@ def extend_system_coord( box: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor]: """Extend the system with WFCC (Wannier Function Charge Centers). This function calculates Wannier Function Charge Centers (WFCC) by adding dipole @@ -321,24 +340,42 @@ def extend_system_coord( Returns ------- - all_coord : torch.Tensor - Extended coordinates with shape (nframes, (natoms + nsel) * 3) - where nsel is the number of selected atoms + Tuple[torch.Tensor, torch.Tensor] + A tuple containing two tensors: + - all_coord : torch.Tensor + Extended coordinates with shape (nframes, 2 * natoms * 3) + where nsel is the number of selected atoms + - dipole_reshaped : torch.Tensor + Atomic dipoles with shape (nframes, natoms, 3) """ - mask = make_mask(self.sel_type, atype) - nframes = coord.shape[0] natoms = coord.shape[1] // 3 - all_dipole = [] - for ii in range(nframes): + all_dipole: list[torch.Tensor] = [] + chunk_coord = torch.split(coord, self.dp_batch_size, dim=0) + chunk_atype = torch.split(atype, self.dp_batch_size, dim=0) + chunk_box = torch.split(box, self.dp_batch_size, dim=0) + # use placeholder to make the jit happy for fparam/aparam is None + chunk_fparam = ( + torch.split(fparam, self.dp_batch_size, dim=0) + if fparam is not None + else chunk_atype + ) + chunk_aparam = ( + torch.split(aparam, self.dp_batch_size, dim=0) + if aparam is not None + else chunk_atype + ) + for _coord, _atype, _box, _fparam, _aparam in zip( + chunk_coord, chunk_atype, chunk_box, chunk_fparam, chunk_aparam + ): dipole_batch = self.model( - coord=coord[ii].reshape(1, -1), - atype=atype[ii].reshape(1, -1), - box=box[ii].reshape(1, -1), + coord=_coord, + atype=_atype, + box=_box, do_atomic_virial=False, - fparam=fparam[ii].reshape(1, -1) if fparam is not None else None, - aparam=aparam[ii].reshape(1, -1) if aparam is not None else None, + fparam=_fparam if fparam is not None else None, + aparam=_aparam if aparam is not None else None, ) # Extract dipole from the output dictionary all_dipole.append(dipole_batch["dipole"]) @@ -352,37 +389,6 @@ def extend_system_coord( dipole_reshaped = dipole.reshape(nframes, natoms, 3) coord_reshaped = coord.reshape(nframes, natoms, 3) - _wfcc_coord = coord_reshaped + dipole_reshaped - # Apply mask and reshape - wfcc_coord = _wfcc_coord[mask.unsqueeze(-1).expand_as(_wfcc_coord)] - wfcc_coord = wfcc_coord.reshape(nframes, -1) - all_coord = torch.cat((coord, wfcc_coord), dim=1) - return all_coord - - -@torch.jit.export -def make_mask( - sel_type: torch.Tensor, - atype: torch.Tensor, -) -> torch.Tensor: - """Create a boolean mask for selected atom types. - - Parameters - ---------- - sel_type : torch.Tensor - The selected atom types to create a mask for - atype : torch.Tensor - The atom types in the system - - Returns - ------- - mask : torch.Tensor - Boolean mask where True indicates atoms of selected types - """ - # Ensure tensors are of the right type - sel_type = sel_type.to(torch.long) - atype = atype.to(torch.long) - - # Create mask using broadcasting for JIT compatibility - mask = (atype.unsqueeze(-1) == sel_type).any(dim=-1) - return mask + wfcc_coord = coord_reshaped + dipole_reshaped + all_coord = torch.cat((coord_reshaped, wfcc_coord), dim=1) + return all_coord, dipole_reshaped diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index 07da63b9bd..de8f763069 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -27,8 +27,9 @@ ) from deepmd.tf.modifier import DipoleChargeModifier as TFDipoleChargeModifier -SEED = 1 -DTYPE = torch.float64 +from ...seed import ( + GLOBAL_SEED, +) def ref_data(): @@ -37,7 +38,7 @@ def ref_data(): str(Path(__file__).parent / "water/data/data_0/set.000/coord.npy") ) nframe = len(all_box) - rng = np.random.default_rng(SEED) + rng = np.random.default_rng(GLOBAL_SEED) selected_id = rng.integers(nframe) coord = all_coord[selected_id].reshape(1, -1) @@ -62,13 +63,10 @@ def setUp(self) -> None: self.sys_charge_map = [6.0, 1.0] self.descriptor_dict = { "type": "se_e2_a", - "sel": [46, 92], + "sel": [12, 24], "rcut_smth": 0.5, "rcut": 4.00, - "neuron": [ - 25, - 50, - ], + "neuron": [6, 12, 24], } # Train DW model @@ -117,38 +115,49 @@ def test_jit(self): def test_consistency(self): coord, box, atype = ref_data() # consistent with the input shape from BaseModifier.modify_data - t_coord = to_torch_tensor(coord).to(DTYPE).reshape(1, -1, 3) - t_box = to_torch_tensor(box).to(DTYPE).reshape(1, 3, 3) - t_atype = to_torch_tensor(atype).to(torch.long) + t_coord = ( + to_torch_tensor(coord).to(env.GLOBAL_PT_FLOAT_PRECISION).reshape(1, -1, 3) + ) + t_box = to_torch_tensor(box).to(env.GLOBAL_PT_FLOAT_PRECISION).reshape(1, 3, 3) + t_atype = to_torch_tensor(atype).to(torch.long).reshape(1, -1) - dm_pred = self.dm_pt( + pt_data = self.dm_pt( coord=t_coord, atype=t_atype, box=t_box, ) + tf_data = {} e, f, v = self.dm_tf.eval( coord=coord, box=box, atype=atype.reshape(-1), ) - - np.testing.assert_allclose( - to_numpy_array(dm_pred["energy"]).reshape(-1), e.reshape(-1), rtol=1e-4 - ) - np.testing.assert_allclose( - to_numpy_array(dm_pred["force"]).reshape(-1), f.reshape(-1), rtol=1e-4 - ) + tf_data["energy"] = e + tf_data["force"] = f + tf_data["virial"] = v + + for kw in ["energy", "virial"]: + np.testing.assert_allclose( + to_numpy_array(pt_data[kw]).reshape(-1), + tf_data[kw].reshape(-1), + atol=1e-6, + ) + kw = "force" np.testing.assert_allclose( - to_numpy_array(dm_pred["virial"]).reshape(-1), v.reshape(-1), rtol=1e-4 + to_numpy_array(pt_data[kw]).reshape(-1), + tf_data[kw].reshape(-1), + rtol=1e-6, ) def test_serialize(self): """Test the serialize method of DipoleChargeModifier.""" coord, box, atype = ref_data() # consistent with the input shape from BaseModifier.modify_data - t_coord = to_torch_tensor(coord).to(DTYPE).reshape(1, -1, 3) - t_box = to_torch_tensor(box).to(DTYPE).reshape(1, 3, 3) - t_atype = to_torch_tensor(atype).to(torch.long) + t_coord = ( + to_torch_tensor(coord).to(env.GLOBAL_PT_FLOAT_PRECISION).reshape(1, -1, 3) + ) + t_box = to_torch_tensor(box).to(env.GLOBAL_PT_FLOAT_PRECISION).reshape(1, 3, 3) + t_atype = to_torch_tensor(atype).to(torch.long).reshape(1, -1) dm0 = self.dm_pt.to(env.DEVICE) dm1 = PTDipoleChargeModifier.deserialize(dm0.serialize()).to(env.DEVICE) @@ -178,8 +187,10 @@ def test_box_none_error(self): """Test that a RuntimeError is raised when box is None.""" coord, _b, atype = ref_data() # consistent with the input shape from BaseModifier.modify_data - t_coord = to_torch_tensor(coord).to(DTYPE).reshape(1, -1, 3) - t_atype = to_torch_tensor(atype).to(torch.long) + t_coord = ( + to_torch_tensor(coord).to(env.GLOBAL_PT_FLOAT_PRECISION).reshape(1, -1, 3) + ) + t_atype = to_torch_tensor(atype).to(torch.long).reshape(1, -1) with self.assertRaises(RuntimeError) as context: self.dm_pt( From 09a05a1bbfd533ed699c0e096717c5aad33cbc95 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Fri, 16 Jan 2026 22:53:44 +0800 Subject: [PATCH 12/25] feat: add DipoleChargeModifier.eval_np --- deepmd/pt/modifier/dipole_charge.py | 41 +++++++++++++++++++ deepmd/pt/train/wrapper.py | 2 +- doc/model/dplr.md | 40 +----------------- .../tests/pt/modifier/test_dipole_charge.py | 34 ++++----------- 4 files changed, 52 insertions(+), 65 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index a72bf09420..ac5a95f197 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -23,6 +23,7 @@ serialize_from_file, ) from deepmd.pt.utils.utils import ( + to_numpy_array, to_torch_tensor, ) @@ -392,3 +393,43 @@ def extend_system_coord( wfcc_coord = coord_reshaped + dipole_reshaped all_coord = torch.cat((coord_reshaped, wfcc_coord), dim=1) return all_coord, dipole_reshaped + + def eval_np( + self, + coord: np.ndarray, + box: np.ndarray, + atype: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + nf = coord.shape[0] + na = coord.reshape(nf, -1, 3).shape[1] + + if fparam is not None: + _fparam = ( + to_torch_tensor(fparam) + .reshape(nf, -1) + .to(env.GLOBAL_PT_FLOAT_PRECISION) + ) + else: + _fparam = None + if aparam is not None: + _aparam = ( + to_torch_tensor(aparam) + .reshape(nf, na, -1) + .to(env.GLOBAL_PT_FLOAT_PRECISION) + ) + else: + _aparam = None + modifier_pred = self.forward( + to_torch_tensor(coord).reshape(nf, -1, 3).to(env.GLOBAL_PT_FLOAT_PRECISION), + to_torch_tensor(atype).reshape(nf, -1).to(torch.long), + to_torch_tensor(box).reshape(nf, 3, 3).to(env.GLOBAL_PT_FLOAT_PRECISION), + _fparam, + _aparam, + ) + return ( + to_numpy_array(modifier_pred["energy"]), + to_numpy_array(modifier_pred["force"]), + to_numpy_array(modifier_pred["virial"]), + ) diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index ddb4a4323d..d2cef75614 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -191,7 +191,7 @@ def forward( if self.modifier is not None: modifier_pred = self.modifier(**input_dict) for k, v in modifier_pred.items(): - model_pred[k] = model_pred[k] + v + model_pred[k] = model_pred[k] + v.reshape(model_pred[k].shape) return model_pred, None, None else: natoms = atype.shape[-1] diff --git a/doc/model/dplr.md b/doc/model/dplr.md index d262b778e4..66b385f155 100644 --- a/doc/model/dplr.md +++ b/doc/model/dplr.md @@ -4,7 +4,7 @@ **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} ::: -Notice: **The interfaces of DPLR are not stable and subject to change. In addition, DP/LAMMPS interface does not yet support PyTorch DPLR models. Use `dp convert-backend` to convert your model to TensorFlow backend for LAMMPS simulations (see [details](#model-conversion)).** +Notice: **The interfaces of DPLR are not stable and subject to change. In addition, DP/LAMMPS interface does not yet support PyTorch DPLR models.** The method of DPLR is described in [this paper][1]. One is recommended to read the paper before using the DPLR. @@ -266,42 +266,4 @@ The MD simulation lasts for only 20 steps. If one runs a longer simulation, it w Another restriction that should be noted is that the energies printed at the zero steps are not correct. This is because at the zero steps the position of the WC has not been updated with the DW model. The energies printed in later steps are correct. -## Model conversion - -### Converting from PyTorch to TensorFlow - -To use DPLR models with LAMMPS, you need to convert them from PyTorch to TensorFlow backend. The conversion process involves the following steps: - -1. **Convert the Deep Wannier (DW) model:** - - ```bash - dp convert-backend dw_model.pth dw_model.pb - ``` - - This converts the PyTorch DW model (`.pth` file) to TensorFlow format (`.pb` file). - -2. **Convert the DPLR energy model:** - - ```bash - dp convert-backend dplr_energy_model.pth dplr_energy_model.pb - ``` - - This converts the PyTorch DPLR energy model (short-range contribution) to TensorFlow backend. - -3. **Combine the models using TensorFlow training:** - ```bash - dp --tf train input.json --init-frz-model dplr_energy_model.pb - ``` - Run TensorFlow training with the DPLR TensorFlow setup, setting the training steps to 0 to combine the models without additional training: - ```json - { - "training": { - "numb_steps": 0, - ... - } - } - ``` - -This process ensures that both the DW model and the DPLR energy model are properly converted and combined for use with LAMMPS simulations. - [1]: https://arxiv.org/abs/2112.13327 diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index de8f763069..2119dafa5f 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -114,40 +114,24 @@ def test_jit(self): def test_consistency(self): coord, box, atype = ref_data() - # consistent with the input shape from BaseModifier.modify_data - t_coord = ( - to_torch_tensor(coord).to(env.GLOBAL_PT_FLOAT_PRECISION).reshape(1, -1, 3) - ) - t_box = to_torch_tensor(box).to(env.GLOBAL_PT_FLOAT_PRECISION).reshape(1, 3, 3) - t_atype = to_torch_tensor(atype).to(torch.long).reshape(1, -1) - pt_data = self.dm_pt( - coord=t_coord, - atype=t_atype, - box=t_box, + pt_data = self.dm_pt.eval_np( + coord=coord, + atype=atype, + box=box, ) - tf_data = {} - e, f, v = self.dm_tf.eval( + tf_data = self.dm_tf.eval( coord=coord, box=box, atype=atype.reshape(-1), ) - tf_data["energy"] = e - tf_data["force"] = f - tf_data["virial"] = v - - for kw in ["energy", "virial"]: + for ii in range(3): np.testing.assert_allclose( - to_numpy_array(pt_data[kw]).reshape(-1), - tf_data[kw].reshape(-1), + pt_data[ii].reshape(-1), + tf_data[ii].reshape(-1), atol=1e-6, + rtol=1e-6, ) - kw = "force" - np.testing.assert_allclose( - to_numpy_array(pt_data[kw]).reshape(-1), - tf_data[kw].reshape(-1), - rtol=1e-6, - ) def test_serialize(self): """Test the serialize method of DipoleChargeModifier.""" From 9f600753a2dfad05d51804540e583a2c0c926ee6 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sat, 17 Jan 2026 20:45:00 +0800 Subject: [PATCH 13/25] fix(modifier): resolve DipoleChargeModifier initialization and execution issues - Remove unnecessary __new__ method in PyTorch DipoleChargeModifier - Fix __new__ method signature in TensorFlow DipoleChargeModifier - Simplify control flow in PyTorch modifier by removing unnecessary else block - Add strict=True to zip() calls for better type checking - Improve test assertions with descriptive error messages - Add checkpoint verification in test case --- deepmd/pt/modifier/dipole_charge.py | 119 +++++++++--------- deepmd/tf/modifier/dipole_charge.py | 4 +- .../tests/pt/modifier/test_dipole_charge.py | 9 +- 3 files changed, 67 insertions(+), 65 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index ac5a95f197..1f4bf4c238 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -44,10 +44,6 @@ class DipoleChargeModifier(BaseModifier): Splitting parameter of the Ewald sum. Unit: A^{-1} """ - def __new__( - cls, *args: tuple, model_name: str | None = None, **kwargs: dict - ) -> "DipoleChargeModifier": - return super().__new__(cls, model_name) def __init__( self, @@ -195,68 +191,67 @@ def forward( raise RuntimeError( "dipole_charge data modifier can only be applied for periodic systems." ) - else: - modifier_pred = {} - nframes = coord.shape[0] - natoms = coord.shape[1] + modifier_pred = {} + nframes = coord.shape[0] + natoms = coord.shape[1] - input_box = box.reshape(nframes, 9) - input_box.requires_grad_(True) + input_box = box.reshape(nframes, 9) + input_box.requires_grad_(True) - detached_box = input_box.detach() - sfactor = torch.matmul( - torch.inverse(detached_box.reshape(nframes, 3, 3)), - input_box.reshape(nframes, 3, 3), - ) - input_coord = torch.matmul(coord, sfactor).reshape(nframes, -1) - - extended_coord, extended_charge, _atomic_dipole = self.extend_system( - input_coord, - atype, - input_box, - fparam, - aparam, - ) + detached_box = input_box.detach() + sfactor = torch.matmul( + torch.inverse(detached_box.reshape(nframes, 3, 3)), + input_box.reshape(nframes, 3, 3), + ) + input_coord = torch.matmul(coord, sfactor).reshape(nframes, -1) - # add Ewald reciprocal correction - tot_e: list[torch.Tensor] = [] - chunk_coord = torch.split( - extended_coord.reshape(nframes, -1, 3), self.dp_batch_size, dim=0 - ) - chunk_box = torch.split( - input_box.reshape(nframes, 3, 3), self.dp_batch_size, dim=0 - ) - chunk_charge = torch.split( - extended_charge.reshape(nframes, -1), self.dp_batch_size, dim=0 - ) - for _coord, _box, _charge in zip(chunk_coord, chunk_box, chunk_charge): - self.er( - _coord, - _box, - self.placeholder_pairs, - self.placeholder_ds, - self.placeholder_buffer_scales, - {"charge": _charge}, - ) - tot_e.append(self.er.reciprocal_energy.unsqueeze(0)) - # nframe, - tot_e = torch.concat(tot_e, dim=0) - # nframe, nat * 3 - tot_f = -calc_grads(tot_e, input_coord) - # nframe, nat, 3 - tot_f = torch.reshape(tot_f, (nframes, natoms, 3)) - # nframe, 9 - tot_v = calc_grads(tot_e, input_box) - tot_v = torch.reshape(tot_v, (nframes, 3, 3)) - # nframe, 3, 3 - tot_v = -torch.matmul( - tot_v.transpose(2, 1), input_box.reshape(nframes, 3, 3) + extended_coord, extended_charge, _atomic_dipole = self.extend_system( + input_coord, + atype, + input_box, + fparam, + aparam, + ) + + # add Ewald reciprocal correction + tot_e: list[torch.Tensor] = [] + chunk_coord = torch.split( + extended_coord.reshape(nframes, -1, 3), self.dp_batch_size, dim=0 + ) + chunk_box = torch.split( + input_box.reshape(nframes, 3, 3), self.dp_batch_size, dim=0 + ) + chunk_charge = torch.split( + extended_charge.reshape(nframes, -1), self.dp_batch_size, dim=0 + ) + for _coord, _box, _charge in zip(chunk_coord, chunk_box, chunk_charge, strict=True): + self.er( + _coord, + _box, + self.placeholder_pairs, + self.placeholder_ds, + self.placeholder_buffer_scales, + {"charge": _charge}, ) + tot_e.append(self.er.reciprocal_energy.unsqueeze(0)) + # nframe, + tot_e = torch.concat(tot_e, dim=0) + # nframe, nat * 3 + tot_f = -calc_grads(tot_e, input_coord) + # nframe, nat, 3 + tot_f = torch.reshape(tot_f, (nframes, natoms, 3)) + # nframe, 9 + tot_v = calc_grads(tot_e, input_box) + tot_v = torch.reshape(tot_v, (nframes, 3, 3)) + # nframe, 3, 3 + tot_v = -torch.matmul( + tot_v.transpose(2, 1), input_box.reshape(nframes, 3, 3) + ) - modifier_pred["energy"] = tot_e - modifier_pred["force"] = tot_f - modifier_pred["virial"] = tot_v - return modifier_pred + modifier_pred["energy"] = tot_e + modifier_pred["force"] = tot_f + modifier_pred["virial"] = tot_v + return modifier_pred def extend_system( self, @@ -368,7 +363,7 @@ def extend_system_coord( else chunk_atype ) for _coord, _atype, _box, _fparam, _aparam in zip( - chunk_coord, chunk_atype, chunk_box, chunk_fparam, chunk_aparam + chunk_coord, chunk_atype, chunk_box, chunk_fparam, chunk_aparam, strict=True ): dipole_batch = self.model( coord=_coord, diff --git a/deepmd/tf/modifier/dipole_charge.py b/deepmd/tf/modifier/dipole_charge.py index a2d3efb353..ad3016c677 100644 --- a/deepmd/tf/modifier/dipole_charge.py +++ b/deepmd/tf/modifier/dipole_charge.py @@ -52,8 +52,8 @@ class DipoleChargeModifier(DeepDipole, BaseModifier): Splitting parameter of the Ewald sum. Unit: A^{-1} """ - def __new__(cls, *args, model_name=None, **kwargs): - return super().__new__(cls, model_name) + def __new__(cls, *args, model_name: str, **kwargs): + return super().__new__(cls, model_file=model_name) def __init__( self, diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index 2119dafa5f..6c08f5cd1c 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -125,12 +125,14 @@ def test_consistency(self): box=box, atype=atype.reshape(-1), ) - for ii in range(3): + output_names = ["energy", "force", "virial"] + for ii, name in enumerate(output_names): np.testing.assert_allclose( pt_data[ii].reshape(-1), tf_data[ii].reshape(-1), atol=1e-6, rtol=1e-6, + err_msg=f"Mismatch in {name}", ) def test_serialize(self): @@ -206,6 +208,11 @@ def test_train(self): trainer = get_trainer(config) trainer.run() + # Verify model checkpoint was created + self.assertTrue( + Path("model.ckpt.pt").exists(), + "Training should produce a model checkpoint", + ) def tearDown(self) -> None: os.chdir(self.orig_dir) From 7c27e964f09dd339eecc3252ccc7e95e814aeddb Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Sat, 17 Jan 2026 20:48:23 +0800 Subject: [PATCH 14/25] update required ver of torch-admp --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 06c5539be0..02aac59844 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,10 +168,11 @@ pin_pytorch_cpu = [ # macos x86 has been deprecated "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'", "torch; platform_machine=='x86_64' and platform_system == 'Darwin'", - "torch-admp==1.1.1", + "torch-admp==1.1.3", ] pin_pytorch_gpu = [ "torch==2.10.0", + "torch-admp==1.1.3", ] pin_jax = [ "jax==0.5.0;python_version>='3.10'", From 9113f90495d9e0bd0fdff791286ba6232728f439 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 17 Jan 2026 12:50:19 +0000 Subject: [PATCH 15/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/modifier/dipole_charge.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index 1f4bf4c238..8997cbee0b 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -44,7 +44,6 @@ class DipoleChargeModifier(BaseModifier): Splitting parameter of the Ewald sum. Unit: A^{-1} """ - def __init__( self, model_name: str | None, @@ -224,7 +223,9 @@ def forward( chunk_charge = torch.split( extended_charge.reshape(nframes, -1), self.dp_batch_size, dim=0 ) - for _coord, _box, _charge in zip(chunk_coord, chunk_box, chunk_charge, strict=True): + for _coord, _box, _charge in zip( + chunk_coord, chunk_box, chunk_charge, strict=True + ): self.er( _coord, _box, @@ -244,9 +245,7 @@ def forward( tot_v = calc_grads(tot_e, input_box) tot_v = torch.reshape(tot_v, (nframes, 3, 3)) # nframe, 3, 3 - tot_v = -torch.matmul( - tot_v.transpose(2, 1), input_box.reshape(nframes, 3, 3) - ) + tot_v = -torch.matmul(tot_v.transpose(2, 1), input_box.reshape(nframes, 3, 3)) modifier_pred["energy"] = tot_e modifier_pred["force"] = tot_f From 4307c7dc88fe26701256209b71cf9086743e2c47 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Wed, 21 Jan 2026 17:16:28 +0800 Subject: [PATCH 16/25] fix bug in tf dipole_charge modifier --- deepmd/tf/modifier/dipole_charge.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/modifier/dipole_charge.py b/deepmd/tf/modifier/dipole_charge.py index ad3016c677..a2d3efb353 100644 --- a/deepmd/tf/modifier/dipole_charge.py +++ b/deepmd/tf/modifier/dipole_charge.py @@ -52,8 +52,8 @@ class DipoleChargeModifier(DeepDipole, BaseModifier): Splitting parameter of the Ewald sum. Unit: A^{-1} """ - def __new__(cls, *args, model_name: str, **kwargs): - return super().__new__(cls, model_file=model_name) + def __new__(cls, *args, model_name=None, **kwargs): + return super().__new__(cls, model_name) def __init__( self, From c54f448e2590a6657f8377f0e6a93906502c9329 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Wed, 21 Jan 2026 19:04:23 +0800 Subject: [PATCH 17/25] add DPModifier --- deepmd/pt/modifier/dp_modifier.py | 92 +++++++++++++++++++ .../tests/pt/modifier/test_data_modifier.py | 47 ++-------- 2 files changed, 101 insertions(+), 38 deletions(-) create mode 100644 deepmd/pt/modifier/dp_modifier.py diff --git a/deepmd/pt/modifier/dp_modifier.py b/deepmd/pt/modifier/dp_modifier.py new file mode 100644 index 0000000000..c9ede815a3 --- /dev/null +++ b/deepmd/pt/modifier/dp_modifier.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import torch + +from deepmd.pt.model.model import ( + BaseModel, +) +from deepmd.pt.modifier.base_modifier import ( + BaseModifier, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.serialization import ( + serialize_from_file, +) + + +class DPModifier(BaseModifier): + def __init__( + self, + dp_model: torch.nn.Module | None = None, + dp_model_file_name: str | None = None, + use_cache: bool = True, + **kwargs, + ) -> None: + """Constructor.""" + super().__init__(use_cache=use_cache) + + if dp_model_file_name is None and dp_model is None: + raise AttributeError("`model_name` or `model` should be specified.") + if dp_model_file_name is not None and dp_model is not None: + raise AttributeError( + "`model_name` and `model` cannot be used simultaneously." + ) + + if dp_model is not None: + self._model = dp_model.to(env.DEVICE) + if dp_model_file_name is not None: + data = serialize_from_file(dp_model_file_name) + assert data["model"]["type"] == "standard" + self._model = ( + BaseModel.get_class_by_type(data["model"]["fitting"]["type"]) + .deserialize(data["model"]) + .to(env.DEVICE) + ) + self._model.eval() + # use jit model for inference + self.model = torch.jit.script(self._model) + + def serialize(self) -> dict: + """Serialize the modifier. + + Returns + ------- + dict + The serialized data + """ + dd = BaseModifier.serialize(self) + dd.update( + { + "dp_model": self._model.serialize(), + } + ) + return dd + + @classmethod + def get_modifier(cls, modifier_params: dict) -> "DPModifier": + """Get the modifier by the parameters. + + By default, all the parameters are directly passed to the constructor. + If not, override this method. + + Parameters + ---------- + modifier_params : dict + The modifier parameters + + Returns + ------- + BaseModifier + The modifier + """ + modifier_params = modifier_params.copy() + modifier_params.pop("type", None) + # convert model_name str + model_name = modifier_params.pop("model_name", None) + if model_name is not None: + modifier_params["dp_model"] = None + modifier_params["dp_model_file_name"] = model_name + modifier = cls(**modifier_params) + return modifier diff --git a/source/tests/pt/modifier/test_data_modifier.py b/source/tests/pt/modifier/test_data_modifier.py index ddba54cd05..37874b0357 100644 --- a/source/tests/pt/modifier/test_data_modifier.py +++ b/source/tests/pt/modifier/test_data_modifier.py @@ -42,11 +42,8 @@ from deepmd.pt.modifier.base_modifier import ( BaseModifier, ) -from deepmd.pt.utils import ( - env, -) -from deepmd.pt.utils.serialization import ( - serialize_from_file, +from deepmd.pt.modifier.dp_modifier import ( + DPModifier, ) from deepmd.pt.utils.utils import ( to_numpy_array, @@ -188,44 +185,19 @@ def modify_data(self, data: dict[str, Array | float], data_sys: DeepmdData) -> N @BaseModifier.register("scaling_tester") -class ModifierScalingTester(BaseModifier): - def __new__( - cls, - *args: tuple, - model: str | None = None, - model_name: str | None = None, - **kwargs: dict, - ) -> "ModifierScalingTester": - return super().__new__(cls, model_name if model_name is not None else model) - +class ModifierScalingTester(DPModifier): def __init__( self, - model: torch.nn.Module | None = None, - model_name: str | None = None, + dp_model: torch.nn.Module | None = None, + dp_model_file_name: str | None = None, sfactor: float = 1.0, use_cache: bool = True, ) -> None: """Initialize a test modifier that applies scaled model predictions using a frozen model.""" - super().__init__(use_cache) + super().__init__(dp_model, dp_model_file_name, use_cache) self.modifier_type = "scaling_tester" self.sfactor = sfactor - if model_name is None and model is None: - raise AttributeError("`model_name` or `model` should be specified.") - if model_name is not None and model is not None: - raise AttributeError( - "`model_name` and `model` cannot be used simultaneously." - ) - - if model is not None: - self._model = model.to(env.DEVICE) - if model_name is not None: - data = serialize_from_file(model_name) - self._model = EnergyModel.deserialize(data["model"]).to(env.DEVICE) - - # use jit model for inference - self.model = torch.jit.script(self._model) - def serialize(self) -> dict: """Serialize the modifier. @@ -234,10 +206,9 @@ def serialize(self) -> dict: dict The serialized data """ - dd = BaseModifier.serialize(self) + dd = super().serialize() dd.update( { - "model": self._model.serialize(), "sfactor": self.sfactor, } ) @@ -249,8 +220,8 @@ def deserialize(cls, data: dict) -> "ModifierScalingTester": data.pop("@class", None) data.pop("type", None) data.pop("@version", None) - model_obj = EnergyModel.deserialize(data.pop("model")) - data["model"] = model_obj + model_obj = EnergyModel.deserialize(data.pop("dp_model")) + data["dp_model"] = model_obj return cls(**data) def forward( From 54d64693963de364f6846bf09f59d5b73740c810 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Wed, 21 Jan 2026 19:49:39 +0800 Subject: [PATCH 18/25] refactor(pt): refactor DipoleChargeModifier to inherit from DPModifier - Refactor DipoleChargeModifier to inherit from DPModifier instead of BaseModifier - Simplify model initialization and batch processing logic - Update torch-admp dependency version to >=1.1.4 - Update tests to match new API and add dipole consistency test --- deepmd/pt/modifier/dipole_charge.py | 136 ++++++------------ pyproject.toml | 4 +- .../tests/pt/modifier/test_dipole_charge.py | 61 +++++++- 3 files changed, 103 insertions(+), 98 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index 8997cbee0b..f6c8854e7b 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import os import numpy as np import torch @@ -16,12 +15,12 @@ from deepmd.pt.modifier.base_modifier import ( BaseModifier, ) +from deepmd.pt.modifier.dp_modifier import ( + DPModifier, +) from deepmd.pt.utils import ( env, ) -from deepmd.pt.utils.serialization import ( - serialize_from_file, -) from deepmd.pt.utils.utils import ( to_numpy_array, to_torch_tensor, @@ -29,7 +28,7 @@ @BaseModifier.register("dipole_charge") -class DipoleChargeModifier(BaseModifier): +class DipoleChargeModifier(DPModifier): """Parameters ---------- model_name @@ -46,36 +45,22 @@ class DipoleChargeModifier(BaseModifier): def __init__( self, - model_name: str | None, + dp_model: DipoleModel | None, model_charge_map: list[float], sys_charge_map: list[float], ewald_h: float = 1.0, ewald_beta: float = 1.0, - ewald_batch_size: int = 5, - dp_batch_size: int | None = None, - model: DipoleModel | None = None, + dp_model_file_name: str | None = None, use_cache: bool = True, ) -> None: """Constructor.""" - super().__init__(use_cache=use_cache) self.modifier_type = "dipole_charge" + super().__init__( + dp_model=dp_model, + dp_model_file_name=dp_model_file_name, + use_cache=use_cache, + ) - if model_name is None and model is None: - raise AttributeError("`model_name` or `model` should be specified.") - if model_name is not None and model is not None: - raise AttributeError( - "`model_name` and `model` cannot be used simultaneously." - ) - - if model is not None: - self._model = model.to(env.DEVICE) - if model_name is not None: - data = serialize_from_file(model_name) - self._model = DipoleModel.deserialize(data["model"]).to(env.DEVICE) - self._model.eval() - - # use jit model for inference - self.model = torch.jit.script(self._model) self.rcut = self.model.get_rcut() self.type_map = self.model.get_type_map() sel_type = self.model.get_sel_type() @@ -95,23 +80,22 @@ def __init__( # init ewald recp self.ewald_h = ewald_h self.ewald_beta = ewald_beta - self.er = CoulombForceModule( + er = CoulombForceModule( rcut=self.rcut, rspace=False, kappa=ewald_beta, spacing=ewald_h, - ) + ).to(env.GLOBAL_PT_FLOAT_PRECISION) + self.er = torch.jit.script(er) + self.er.eval() self.placeholder_pairs = torch.ones((1, 2), device=env.DEVICE, dtype=torch.long) - self.placeholder_ds = torch.ones((1), device=env.DEVICE, dtype=torch.float64) + self.placeholder_ds = torch.ones( + (1), device=env.DEVICE, dtype=env.GLOBAL_PT_FLOAT_PRECISION + ) self.placeholder_buffer_scales = torch.zeros( - (1), device=env.DEVICE, dtype=torch.float64 + (1), device=env.DEVICE, dtype=env.GLOBAL_PT_FLOAT_PRECISION ) - self.ewald_batch_size = ewald_batch_size - if dp_batch_size is None: - dp_batch_size = int(os.environ.get("DP_INFER_BATCH_SIZE", 1)) - self.dp_batch_size = dp_batch_size - def serialize(self) -> dict: """Serialize the modifier. @@ -120,16 +104,13 @@ def serialize(self) -> dict: dict The serialized data """ - dd = BaseModifier.serialize(self) + dd = super().serialize() dd.update( { - "model": self._model.serialize(), "model_charge_map": self._model_charge_map, "sys_charge_map": self._sys_charge_map, "ewald_h": self.ewald_h, "ewald_beta": self.ewald_beta, - "ewald_batch_size": self.ewald_batch_size, - "dp_batch_size": self.dp_batch_size, } ) return dd @@ -140,9 +121,9 @@ def deserialize(cls, data: dict) -> "DipoleChargeModifier": data.pop("@class", None) data.pop("type", None) data.pop("@version", None) - model_obj = DipoleModel.deserialize(data.pop("model")) - data["model"] = model_obj - data["model_name"] = None + model_obj = DipoleModel.deserialize(data.pop("dp_model")) + data["dp_model"] = model_obj + data["dp_model_file_name"] = None return cls(**data) def forward( @@ -213,30 +194,23 @@ def forward( ) # add Ewald reciprocal correction - tot_e: list[torch.Tensor] = [] - chunk_coord = torch.split( - extended_coord.reshape(nframes, -1, 3), self.dp_batch_size, dim=0 + placeholder_pairs = torch.tile( + self.placeholder_pairs.unsqueeze(0), (nframes, 1, 1) ) - chunk_box = torch.split( - input_box.reshape(nframes, 3, 3), self.dp_batch_size, dim=0 + placeholder_ds = torch.tile(self.placeholder_ds.unsqueeze(0), (nframes, 1)) + placeholder_buffer_scales = torch.tile( + self.placeholder_buffer_scales.unsqueeze(0), (nframes, 1) ) - chunk_charge = torch.split( - extended_charge.reshape(nframes, -1), self.dp_batch_size, dim=0 + self.er( + extended_coord.reshape(nframes, 2 * natoms, 3), + input_box.reshape(nframes, 3, 3), + placeholder_pairs, + placeholder_ds, + placeholder_buffer_scales, + {"charge": extended_charge.reshape(nframes, 2 * natoms)}, ) - for _coord, _box, _charge in zip( - chunk_coord, chunk_box, chunk_charge, strict=True - ): - self.er( - _coord, - _box, - self.placeholder_pairs, - self.placeholder_ds, - self.placeholder_buffer_scales, - {"charge": _charge}, - ) - tot_e.append(self.er.reciprocal_energy.unsqueeze(0)) # nframe, - tot_e = torch.concat(tot_e, dim=0) + tot_e = self.er.reciprocal_energy # nframe, nat * 3 tot_f = -calc_grads(tot_e, input_coord) # nframe, nat, 3 @@ -346,37 +320,17 @@ def extend_system_coord( nframes = coord.shape[0] natoms = coord.shape[1] // 3 - all_dipole: list[torch.Tensor] = [] - chunk_coord = torch.split(coord, self.dp_batch_size, dim=0) - chunk_atype = torch.split(atype, self.dp_batch_size, dim=0) - chunk_box = torch.split(box, self.dp_batch_size, dim=0) - # use placeholder to make the jit happy for fparam/aparam is None - chunk_fparam = ( - torch.split(fparam, self.dp_batch_size, dim=0) - if fparam is not None - else chunk_atype - ) - chunk_aparam = ( - torch.split(aparam, self.dp_batch_size, dim=0) - if aparam is not None - else chunk_atype + model_pred = self.model( + coord=coord, + atype=atype, + box=box, + do_atomic_virial=False, + fparam=fparam if fparam is not None else None, + aparam=aparam if aparam is not None else None, ) - for _coord, _atype, _box, _fparam, _aparam in zip( - chunk_coord, chunk_atype, chunk_box, chunk_fparam, chunk_aparam, strict=True - ): - dipole_batch = self.model( - coord=_coord, - atype=_atype, - box=_box, - do_atomic_virial=False, - fparam=_fparam if fparam is not None else None, - aparam=_aparam if aparam is not None else None, - ) - # Extract dipole from the output dictionary - all_dipole.append(dipole_batch["dipole"]) - # nframe x natoms x 3 - dipole = torch.cat(all_dipole, dim=0) + # nframe, natoms, 3 + dipole = model_pred["dipole"] if dipole.shape[0] != nframes: raise RuntimeError( f"Dipole shape mismatch: expected {nframes} frames, got {dipole.shape[0]}" diff --git a/pyproject.toml b/pyproject.toml index 02aac59844..16f3649215 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,11 +168,11 @@ pin_pytorch_cpu = [ # macos x86 has been deprecated "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'", "torch; platform_machine=='x86_64' and platform_system == 'Darwin'", - "torch-admp==1.1.3", + "torch-admp>=1.1.4", ] pin_pytorch_gpu = [ "torch==2.10.0", - "torch-admp==1.1.3", + "torch-admp>=1.1.4", ] pin_jax = [ "jax==0.5.0;python_version>='3.10'", diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index 6c08f5cd1c..a3aa68116a 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -13,6 +13,12 @@ from deepmd.entrypoints.convert_backend import ( convert_backend, ) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer import ( + DeepEval, +) from deepmd.pt.entrypoints.main import ( freeze, get_trainer, @@ -41,8 +47,8 @@ def ref_data(): rng = np.random.default_rng(GLOBAL_SEED) selected_id = rng.integers(nframe) - coord = all_coord[selected_id].reshape(1, -1) - box = all_box[selected_id].reshape(1, -1) + coord = all_coord[selected_id].reshape(1, -1).astype(GLOBAL_NP_FLOAT_PRECISION) + box = all_box[selected_id].reshape(1, -1).astype(GLOBAL_NP_FLOAT_PRECISION) atype = np.loadtxt( str(Path(__file__).parent / "water/data/data_0/type.raw"), dtype=int, @@ -68,6 +74,14 @@ def setUp(self) -> None: "rcut": 4.00, "neuron": [6, 12, 24], } + self.modifier_dict = { + "type": "dipole_charge", + "model_name": "dw_model.pth", + "model_charge_map": self.model_charge_map, + "sys_charge_map": self.sys_charge_map, + "ewald_beta": self.ewald_beta, + "ewald_h": self.ewald_h, + } # Train DW model input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") @@ -95,11 +109,12 @@ def setUp(self) -> None: convert_backend(INPUT="dw_model.pth", OUTPUT="dw_model.pb") self.dm_pt = PTDipoleChargeModifier( - "dw_model.pth", + None, self.model_charge_map, self.sys_charge_map, self.ewald_h, self.ewald_beta, + "dw_model.pth", ) self.dm_tf = TFDipoleChargeModifier( "dw_model.pb", @@ -112,7 +127,39 @@ def setUp(self) -> None: def test_jit(self): torch.jit.script(self.dm_pt) + def test_dipole_consistency(self): + coord, box, atype = ref_data() + tf_model = DeepEval("dw_model.pb") + tf_data = tf_model.eval( + coords=coord, + cells=box, + atom_types=atype.reshape(-1), + ) + + nframes = 1 + input_box = to_torch_tensor(box).reshape(nframes, 9) + input_coord = to_torch_tensor(coord).reshape(nframes, -1) + input_atype = to_torch_tensor(atype).to(torch.long) + _extended_coord, _extended_charge, atomic_dipole = self.dm_pt.extend_system( + input_coord, + input_atype, + input_box, + None, + None, + ) + + np.testing.assert_allclose( + tf_data.reshape(-1, 3), + to_numpy_array(atomic_dipole).reshape(-1, 3), + ) + + @unittest.skipIf( + env.GLOBAL_PT_FLOAT_PRECISION != torch.float64, "run only for double precision" + ) def test_consistency(self): + dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float64) + coord, box, atype = ref_data() pt_data = self.dm_pt.eval_np( @@ -125,16 +172,19 @@ def test_consistency(self): box=box, atype=atype.reshape(-1), ) + tol = 1e-6 output_names = ["energy", "force", "virial"] for ii, name in enumerate(output_names): np.testing.assert_allclose( pt_data[ii].reshape(-1), tf_data[ii].reshape(-1), - atol=1e-6, - rtol=1e-6, + atol=tol, + rtol=tol, err_msg=f"Mismatch in {name}", ) + torch.set_default_dtype(dtype) + def test_serialize(self): """Test the serialize method of DipoleChargeModifier.""" coord, box, atype = ref_data() @@ -206,6 +256,7 @@ def test_train(self): ] config["training"]["numb_steps"] = 1 + config["model"]["modifier"] = self.modifier_dict trainer = get_trainer(config) trainer.run() # Verify model checkpoint was created From f0376995ac9661125100b68aff03b7e667ab5e11 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Wed, 21 Jan 2026 19:51:40 +0800 Subject: [PATCH 19/25] update DPModifier.__init__ --- deepmd/pt/modifier/dp_modifier.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/pt/modifier/dp_modifier.py b/deepmd/pt/modifier/dp_modifier.py index c9ede815a3..c411843bc7 100644 --- a/deepmd/pt/modifier/dp_modifier.py +++ b/deepmd/pt/modifier/dp_modifier.py @@ -22,7 +22,6 @@ def __init__( dp_model: torch.nn.Module | None = None, dp_model_file_name: str | None = None, use_cache: bool = True, - **kwargs, ) -> None: """Constructor.""" super().__init__(use_cache=use_cache) From 86b4e5e57ee9612b84df293ea86152d346f6f4e0 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 22 Jan 2026 10:56:30 +0800 Subject: [PATCH 20/25] fix(pt): resolve issues with dipole charge modifier - Use pickle instead of json for modifier serialization to handle np.ndarray - Fix device placement in EwaldReal class by explicitly setting device and dtype - Add proper error handling in test case to ensure torch default dtype is restored --- deepmd/pt/entrypoints/main.py | 2 + deepmd/pt/modifier/dipole_charge.py | 2 +- .../tests/pt/modifier/test_dipole_charge.py | 45 ++++++++++--------- 3 files changed, 26 insertions(+), 23 deletions(-) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 0f68052287..805fdbd5e8 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -404,6 +404,8 @@ def freeze( extra_files = {"modifier_data": ""} dm = tester.modifier if dm is not None: + # dict from dm.serialize() includes np.ndarray + # use pickle rather than json bytes_data = pickle.dumps(dm.serialize()) extra_files = {"modifier_data": bytes_data} torch.jit.save( diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index f6c8854e7b..1a00ea2f9f 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -85,7 +85,7 @@ def __init__( rspace=False, kappa=ewald_beta, spacing=ewald_h, - ).to(env.GLOBAL_PT_FLOAT_PRECISION) + ).to(device=env.DEVICE, dtype=env.GLOBAL_PT_FLOAT_PRECISION) self.er = torch.jit.script(er) self.er.eval() self.placeholder_pairs = torch.ones((1, 2), device=env.DEVICE, dtype=torch.long) diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index a3aa68116a..f8cd3660fc 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -160,30 +160,31 @@ def test_consistency(self): dtype = torch.get_default_dtype() torch.set_default_dtype(torch.float64) - coord, box, atype = ref_data() + try: + coord, box, atype = ref_data() - pt_data = self.dm_pt.eval_np( - coord=coord, - atype=atype, - box=box, - ) - tf_data = self.dm_tf.eval( - coord=coord, - box=box, - atype=atype.reshape(-1), - ) - tol = 1e-6 - output_names = ["energy", "force", "virial"] - for ii, name in enumerate(output_names): - np.testing.assert_allclose( - pt_data[ii].reshape(-1), - tf_data[ii].reshape(-1), - atol=tol, - rtol=tol, - err_msg=f"Mismatch in {name}", + pt_data = self.dm_pt.eval_np( + coord=coord, + atype=atype, + box=box, ) - - torch.set_default_dtype(dtype) + tf_data = self.dm_tf.eval( + coord=coord, + box=box, + atype=atype.reshape(-1), + ) + tol = 1e-6 + output_names = ["energy", "force", "virial"] + for ii, name in enumerate(output_names): + np.testing.assert_allclose( + pt_data[ii].reshape(-1), + tf_data[ii].reshape(-1), + atol=tol, + rtol=tol, + err_msg=f"Mismatch in {name}", + ) + finally: + torch.set_default_dtype(dtype) def test_serialize(self): """Test the serialize method of DipoleChargeModifier.""" From 57a9593068df5d2ca289684e08c92cb4f5dbb9c7 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Fri, 23 Jan 2026 19:03:40 +0800 Subject: [PATCH 21/25] docs(pt): update docstrings in dipole_charge.py - Updated class docstring to follow NumPy style with comprehensive description - Enhanced __init__ method docstring with detailed parameter descriptions - Improved serialize/deserialize method docstrings - Added detailed docstring for eval_np method - Ensured all docstrings follow the project's NumPy style guidelines --- deepmd/pt/modifier/dipole_charge.py | 101 +++++++++++++++--- deepmd/pt/modifier/dp_modifier.py | 6 +- .../tests/pt/modifier/test_dipole_charge.py | 50 ++++----- 3 files changed, 114 insertions(+), 43 deletions(-) diff --git a/deepmd/pt/modifier/dipole_charge.py b/deepmd/pt/modifier/dipole_charge.py index 1a00ea2f9f..a4a5833e11 100644 --- a/deepmd/pt/modifier/dipole_charge.py +++ b/deepmd/pt/modifier/dipole_charge.py @@ -29,18 +29,29 @@ @BaseModifier.register("dipole_charge") class DipoleChargeModifier(DPModifier): - """Parameters + """Modifier for dipole-charge systems using Wannier Function Charge Centers (WFCC). + + This modifier extends a system with Wannier Function Charge Centers (WFCC) by + adding dipole vectors to atomic coordinates for selected atom types. It then + calculates the electrostatic interactions using Ewald reciprocal summation + to obtain energy, force, and virial corrections. + + Parameters ---------- - model_name - The model file for the DeepDipole model - model_charge_map - Gives the amount of charge for the wfcc - sys_charge_map - Gives the amount of charge for the real atoms - ewald_h - Grid spacing of the reciprocal part of Ewald sum. Unit: A - ewald_beta - Splitting parameter of the Ewald sum. Unit: A^{-1} + dp_model : DipoleModel | None + The DeepDipole model to use for dipole prediction + model_charge_map : List[float] + The amount of charge for the WFCC for each selected atom type + sys_charge_map : List[float] + The amount of charge for the real atoms for each atom type + ewald_h : float, optional + Grid spacing of the reciprocal part of Ewald sum. Unit: Å, default is 1.0 + ewald_beta : float, optional + Splitting parameter of the Ewald sum. Unit: Å^{-1}, default is 1.0 + dp_model_file_name : str | None, optional + Path to the model file, by default None + use_cache : bool, optional + Whether to use cache for computations, by default True """ def __init__( @@ -53,7 +64,30 @@ def __init__( dp_model_file_name: str | None = None, use_cache: bool = True, ) -> None: - """Constructor.""" + """Initialize the DipoleChargeModifier. + + Parameters + ---------- + dp_model : DipoleModel | None + The DeepDipole model to use for dipole prediction + model_charge_map : List[float] + The amount of charge for the WFCC for each selected atom type + sys_charge_map : List[float] + The amount of charge for the real atoms for each atom type + ewald_h : float, optional + Grid spacing of the reciprocal part of Ewald sum. Unit: Å, default is 1.0 + ewald_beta : float, optional + Splitting parameter of the Ewald sum. Unit: Å^{-1}, default is 1.0 + dp_model_file_name : str | None, optional + Path to the model file, by default None + use_cache : bool, optional + Whether to use cache for computations, by default True + + Raises + ------ + ValueError + If model_charge_map and sel_type have mismatching lengths + """ self.modifier_type = "dipole_charge" super().__init__( dp_model=dp_model, @@ -97,12 +131,12 @@ def __init__( ) def serialize(self) -> dict: - """Serialize the modifier. + """Serialize the modifier to a dictionary. Returns ------- dict - The serialized data + The serialized data containing model parameters and configuration """ dd = super().serialize() dd.update( @@ -117,6 +151,18 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "DipoleChargeModifier": + """Deserialize the modifier from a dictionary. + + Parameters + ---------- + data : dict + The serialized data containing model parameters and configuration + + Returns + ------- + DipoleChargeModifier + The deserialized modifier instance + """ data = data.copy() data.pop("@class", None) data.pop("type", None) @@ -342,6 +388,7 @@ def extend_system_coord( all_coord = torch.cat((coord_reshaped, wfcc_coord), dim=1) return all_coord, dipole_reshaped + @torch.jit.unused def eval_np( self, coord: np.ndarray, @@ -350,6 +397,32 @@ def eval_np( fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Evaluate the modifier with NumPy input and output. + + This method converts NumPy inputs to PyTorch tensors, evaluates the modifier, + and converts the results back to NumPy arrays. + + Parameters + ---------- + coord : np.ndarray + The coordinates of atoms with shape (nframes, natoms, 3) + box : np.ndarray + The simulation box with shape (nframes, 3, 3) + atype : np.ndarray + The atom types with shape (nframes, natoms) + fparam : np.ndarray | None, optional + Frame parameters with shape (nframes, nfp), by default None + aparam : np.ndarray | None, optional + Atom parameters with shape (nframes, natoms, nap), by default None + + Returns + ------- + Tuple[np.ndarray, np.ndarray, np.ndarray] + A tuple containing three NumPy arrays: + - energy: Energy correction with shape (nframes, 1) + - force: Force correction with shape (nframes, natoms, 3) + - virial: Virial correction with shape (nframes, 3, 3) + """ nf = coord.shape[0] na = coord.reshape(nf, -1, 3).shape[1] diff --git a/deepmd/pt/modifier/dp_modifier.py b/deepmd/pt/modifier/dp_modifier.py index c411843bc7..c332ffd525 100644 --- a/deepmd/pt/modifier/dp_modifier.py +++ b/deepmd/pt/modifier/dp_modifier.py @@ -37,7 +37,11 @@ def __init__( self._model = dp_model.to(env.DEVICE) if dp_model_file_name is not None: data = serialize_from_file(dp_model_file_name) - assert data["model"]["type"] == "standard" + model_type = data["model"]["type"] + if model_type != "standard": + raise ValueError( + f"DPModifier only support standard model. Unsupported model type: {model_type}" + ) self._model = ( BaseModel.get_class_by_type(data["model"]["fitting"]["type"]) .deserialize(data["model"]) diff --git a/source/tests/pt/modifier/test_dipole_charge.py b/source/tests/pt/modifier/test_dipole_charge.py index f8cd3660fc..1237a9ff81 100644 --- a/source/tests/pt/modifier/test_dipole_charge.py +++ b/source/tests/pt/modifier/test_dipole_charge.py @@ -141,9 +141,9 @@ def test_dipole_consistency(self): input_coord = to_torch_tensor(coord).reshape(nframes, -1) input_atype = to_torch_tensor(atype).to(torch.long) _extended_coord, _extended_charge, atomic_dipole = self.dm_pt.extend_system( - input_coord, + input_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), input_atype, - input_box, + input_box.to(env.GLOBAL_PT_FLOAT_PRECISION), None, None, ) @@ -157,34 +157,28 @@ def test_dipole_consistency(self): env.GLOBAL_PT_FLOAT_PRECISION != torch.float64, "run only for double precision" ) def test_consistency(self): - dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.float64) - - try: - coord, box, atype = ref_data() + coord, box, atype = ref_data() - pt_data = self.dm_pt.eval_np( - coord=coord, - atype=atype, - box=box, - ) - tf_data = self.dm_tf.eval( - coord=coord, - box=box, - atype=atype.reshape(-1), + pt_data = self.dm_pt.eval_np( + coord=coord, + atype=atype, + box=box, + ) + tf_data = self.dm_tf.eval( + coord=coord, + box=box, + atype=atype.reshape(-1), + ) + tol = 1e-6 + output_names = ["energy", "force", "virial"] + for ii, name in enumerate(output_names): + np.testing.assert_allclose( + pt_data[ii].reshape(-1), + tf_data[ii].reshape(-1), + atol=tol, + rtol=tol, + err_msg=f"Mismatch in {name}", ) - tol = 1e-6 - output_names = ["energy", "force", "virial"] - for ii, name in enumerate(output_names): - np.testing.assert_allclose( - pt_data[ii].reshape(-1), - tf_data[ii].reshape(-1), - atol=tol, - rtol=tol, - err_msg=f"Mismatch in {name}", - ) - finally: - torch.set_default_dtype(dtype) def test_serialize(self): """Test the serialize method of DipoleChargeModifier.""" From a6cf06cae632973d0e7f44bc7be6a2bc9180f2a3 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Wed, 28 Jan 2026 18:06:29 +0800 Subject: [PATCH 22/25] update ver requirement for torch-admp --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 16f3649215..21ed7f747e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,11 +168,11 @@ pin_pytorch_cpu = [ # macos x86 has been deprecated "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'", "torch; platform_machine=='x86_64' and platform_system == 'Darwin'", - "torch-admp>=1.1.4", + "torch-admp>=1.1.5", ] pin_pytorch_gpu = [ "torch==2.10.0", - "torch-admp>=1.1.4", + "torch-admp>=1.1.5", ] pin_jax = [ "jax==0.5.0;python_version>='3.10'", From aa1617d6cb3004646d4da76334940200ec836dbe Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Fri, 30 Jan 2026 11:52:07 +0800 Subject: [PATCH 23/25] build: update torch-admp dependency to version 1.1.5 or newer --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 21ed7f747e..aa94462fe1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,7 +141,7 @@ jax = [ # 'jax-ai-stack;python_version>="3.10"', ] torch = [ - "torch-admp", + "torch-admp>=1.1.5", ] [tool.deepmd_build_backend.scripts] From 01157c12b8864f244b09fe4b683b88a004e0f5a3 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Fri, 30 Jan 2026 12:10:11 +0800 Subject: [PATCH 24/25] fix(pt): correct error messages and add fallback for fitting type in DPModifier --- deepmd/pt/modifier/dp_modifier.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/modifier/dp_modifier.py b/deepmd/pt/modifier/dp_modifier.py index c332ffd525..ff45d9e56f 100644 --- a/deepmd/pt/modifier/dp_modifier.py +++ b/deepmd/pt/modifier/dp_modifier.py @@ -27,10 +27,10 @@ def __init__( super().__init__(use_cache=use_cache) if dp_model_file_name is None and dp_model is None: - raise AttributeError("`model_name` or `model` should be specified.") + raise AttributeError("`dp_model_file_name` or `dp_model` should be specified.") if dp_model_file_name is not None and dp_model is not None: raise AttributeError( - "`model_name` and `model` cannot be used simultaneously." + "`dp_model_file_name` and `dp_model` cannot be used simultaneously." ) if dp_model is not None: @@ -42,8 +42,9 @@ def __init__( raise ValueError( f"DPModifier only support standard model. Unsupported model type: {model_type}" ) + fitting_type = data["model"].get("fitting", {}).get("type", "ener") self._model = ( - BaseModel.get_class_by_type(data["model"]["fitting"]["type"]) + BaseModel.get_class_by_type(fitting_type) .deserialize(data["model"]) .to(env.DEVICE) ) From e4644389c4a663cc51984e438fbe16a3e9fae73e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 30 Jan 2026 04:13:26 +0000 Subject: [PATCH 25/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/modifier/dp_modifier.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/modifier/dp_modifier.py b/deepmd/pt/modifier/dp_modifier.py index ff45d9e56f..eaa52edbb7 100644 --- a/deepmd/pt/modifier/dp_modifier.py +++ b/deepmd/pt/modifier/dp_modifier.py @@ -27,7 +27,9 @@ def __init__( super().__init__(use_cache=use_cache) if dp_model_file_name is None and dp_model is None: - raise AttributeError("`dp_model_file_name` or `dp_model` should be specified.") + raise AttributeError( + "`dp_model_file_name` or `dp_model` should be specified." + ) if dp_model_file_name is not None and dp_model is not None: raise AttributeError( "`dp_model_file_name` and `dp_model` cannot be used simultaneously."