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

Author: ChiahsinChu

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,20 +15,20 @@
 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,
 )
 
 
 @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]}"

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.7,<2.10",
-  "torch-admp==1.1.3",
+  "torch-admp>=1.1.4",
 ]
 pin_jax = [
   "jax==0.5.0;python_version>='3.10'",

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