Support dpmodel and jax

Chengqian-Zhang · Chengqian-Zhang · commit aa386a646114 · 2025-08-07T11:01:14.000Z
diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py
@@ -6,6 +6,10 @@
 
 import numpy as np
 
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
 from deepmd.dpmodel.common import (
     DEFAULT_PRECISION,
 )
@@ -108,6 +112,20 @@ def __init__(
             type_map=type_map,
         )
 
+    def output_def(self) -> FittingOutputDef:
+        return FittingOutputDef(
+            [
+                OutputVariableDef(
+                    self.var_name,
+                    [self.dim_out],
+                    reducible=True,
+                    r_differentiable=False,
+                    c_differentiable=False,
+                    intensive=self.intensive,
+                ),
+            ]
+        )
+
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
         data = data.copy()
diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py
@@ -355,6 +355,7 @@ def forward_common_atomic(
                 self.atomic_output_def(),
                 extended_coord,
                 do_atomic_virial=do_atomic_virial,
+                mask=atomic_ret["mask"] if "mask" in atomic_ret else None,
             )
 
         forward_lower = call_lower
diff --git a/deepmd/dpmodel/model/transform_output.py b/deepmd/dpmodel/model/transform_output.py
@@ -17,13 +17,15 @@
     get_hessian_name,
     get_reduce_name,
 )
+from typing import Optional
 
 
 def fit_output_to_model_output(
     fit_ret: dict[str, np.ndarray],
     fit_output_def: FittingOutputDef,
     coord_ext: np.ndarray,
     do_atomic_virial: bool = False,
+    mask: Optional[np.ndarray] = None
 ) -> dict[str, np.ndarray]:
     """Transform the output of the fitting network to
     the model output.
@@ -38,9 +40,21 @@ def fit_output_to_model_output(
         if vdef.reducible:
             kk_redu = get_reduce_name(kk)
             # cast to energy prec before reduction
-            model_ret[kk_redu] = xp.sum(
-                vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
-            )
+            if vdef.intensive:
+                if (mask is not None) and (mask == 0.0).any():
+                    mask = mask.astype(bool)
+                    model_ret[kk_redu] = xp.stack([
+                        xp.mean(vv[ii].astype(GLOBAL_ENER_FLOAT_PRECISION)[mask[ii]], axis=atom_axis)
+                        for ii in range(mask.shape[0])
+                    ])
+                else:
+                    model_ret[kk_redu] = xp.mean(
+                        vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+                    )
+            else:
+                model_ret[kk_redu] = xp.sum(
+                    vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+                )
             if vdef.r_differentiable:
                 kk_derv_r, kk_derv_c = get_deriv_name(kk)
                 # name-holders
diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py
@@ -46,7 +46,18 @@ def forward_common_atomic(
         atom_axis = -(len(shap) + 1)
         if vdef.reducible:
             kk_redu = get_reduce_name(kk)
-            model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis)
+            if vdef.intensive:
+                mask = atomic_ret["mask"] if "mask" in atomic_ret else None
+                if (mask is not None) and (mask==0.0).any():
+                    mask = mask.astype(jnp.bool_)
+                    model_predict[kk_redu] = jnp.stack([
+                        jnp.mean(vv[ii][mask[ii]], axis=atom_axis)
+                        for ii in range(mask.shape[0])
+                    ])
+                else:
+                    model_predict[kk_redu] = jnp.mean(vv, axis=atom_axis)
+            else:
+                model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis)
             kk_derv_r, kk_derv_c = get_deriv_name(kk)
             if vdef.r_differentiable:
 
diff --git a/source/tests/common/dpmodel/test_padding_atoms.py b/source/tests/common/dpmodel/test_padding_atoms.py
@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import sys
+import unittest
+
+import numpy as np
+from copy import deepcopy
+
+from deepmd.dpmodel.descriptor.se_e2_a import (
+    DescrptSeA,
+)
+from deepmd.dpmodel.fitting import (
+    PropertyFittingNet,
+)
+from deepmd.dpmodel.model.property_model import (
+    PropertyModel,
+)
+
+
+class TestCaseSingleFrameWithoutNlist:
+    def setUp(self) -> None:
+        # nf=2, nloc == 3
+        self.nloc = 3
+        self.nt = 2
+        self.coord = np.array(
+            [
+                [
+                    [0, 0, 0],
+                    [0, 1, 0],
+                    [0, 0, 1],
+                ],
+                [
+                    [1, 0, 1],
+                    [0, 1, 1],
+                    [1, 1, 0],
+                ]
+            ],
+            dtype=np.float64,
+        )
+        self.atype = np.array([[0, 0, 1],[1, 1, 0]], dtype=int).reshape([2, self.nloc])
+        self.cell = 2.0 * np.eye(3).reshape([1, 9])
+        self.cell = np.array([self.cell,self.cell]).reshape(2, 9)
+        self.sel = [16, 8]
+        self.rcut = 2.2
+        self.rcut_smth = 0.4
+        self.atol = 1e-12
+
+
+class TestPaddingAtoms(unittest.TestCase, TestCaseSingleFrameWithoutNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithoutNlist.setUp(self)
+
+    def test_padding_atoms_consistency(self):
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        )
+        ft = PropertyFittingNet(
+            self.nt,
+            ds.get_dim_out(),
+            mixed_types=ds.mixed_types(),
+            intensive=True,
+        )
+        type_map = ["foo", "bar"]
+        model = PropertyModel(ds, ft, type_map=type_map)
+        var_name = model.get_var_name()
+        args = [self.coord, self.atype, self.cell]
+        result = model.call(*args)
+        # test intensive
+        np.testing.assert_allclose(
+            result[f"{var_name}_redu"],
+            np.mean(result[f"{var_name}"],axis=1),
+            atol=self.atol,
+        )
+        # test padding atoms
+        padding_atoms_list = [1, 5, 10]
+        for padding_atoms in padding_atoms_list:
+            coord = deepcopy(self.coord)
+            atype = deepcopy(self.atype)
+            atype_padding = np.pad(
+                atype,
+                pad_width=((0, 0), (0, padding_atoms)),
+                mode='constant',
+                constant_values=-1
+            )
+            coord_padding = np.pad(
+                coord,
+                pad_width=((0, 0), (0, padding_atoms), (0, 0)),
+                mode='constant',
+                constant_values=0
+            )
+            args = [coord_padding, atype_padding, self.cell]
+            result_padding = model.call(*args)
+            np.testing.assert_allclose(
+                result[f"{var_name}_redu"],
+                result_padding[f"{var_name}_redu"],
+                atol=self.atol,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/jax/test_padding_atoms.py b/source/tests/jax/test_padding_atoms.py
@@ -29,36 +29,43 @@
     dtype = jnp.float64
 
 
-#@unittest.skipIf(
-#    sys.version_info < (3, 10),
-#    "JAX requires Python 3.10 or later",
-#)
+@unittest.skipIf(
+    sys.version_info < (3, 10),
+    "JAX requires Python 3.10 or later",
+)
 class TestCaseSingleFrameWithoutNlist:
     def setUp(self) -> None:
-        # nloc == 3, nall == 4
+        # nf=2, nloc == 3
         self.nloc = 3
-        self.nf, self.nt = 1, 2
+        self.nt = 2
         self.coord = np.array(
             [
-                [0, 0, 0],
-                [0, 1, 0],
-                [0, 0, 1],
+                [
+                    [0, 0, 0],
+                    [0, 1, 0],
+                    [0, 0, 1],
+                ],
+                [
+                    [1, 0, 1],
+                    [0, 1, 1],
+                    [1, 1, 0],
+                ]
             ],
             dtype=np.float64,
-        ).reshape([1, self.nloc * 3])
-        self.atype = np.array([0, 0, 1], dtype=int).reshape([1, self.nloc])
+        )
+        self.atype = np.array([[0, 0, 1],[1, 1, 0]], dtype=int).reshape([2, self.nloc])
         self.cell = 2.0 * np.eye(3).reshape([1, 9])
-        # sel = [5, 2]
+        self.cell = np.array([self.cell,self.cell]).reshape(2, 9)
         self.sel = [16, 8]
         self.rcut = 2.2
         self.rcut_smth = 0.4
         self.atol = 1e-12
 
 
-#@unittest.skipIf(
-#    sys.version_info < (3, 10),
-#    "JAX requires Python 3.10 or later",
-#)
+@unittest.skipIf(
+    sys.version_info < (3, 10),
+    "JAX requires Python 3.10 or later",
+)
 class TestPaddingAtoms(unittest.TestCase, TestCaseSingleFrameWithoutNlist):
     def setUp(self):
         TestCaseSingleFrameWithoutNlist.setUp(self)
@@ -77,12 +84,40 @@ def test_padding_atoms_consistency(self):
         )
         type_map = ["foo", "bar"]
         model = PropertyModel(ds, ft, type_map=type_map)
+        var_name = model.get_var_name()
         args = [to_jax_array(ii) for ii in [self.coord, self.atype, self.cell]]
-        ret_base = model.call(*args)
+        result = model.call(*args)
+        # test intensive
+        np.testing.assert_allclose(
+            to_numpy_array(result[f"{var_name}_redu"]),
+            np.mean(to_numpy_array(result[f"{var_name}"]),axis=1),
+            atol=self.atol,
+        )
+        # test padding atoms
+        padding_atoms_list = [1, 5, 10]
+        for padding_atoms in padding_atoms_list:
+            coord = deepcopy(self.coord)
+            atype = deepcopy(self.atype)
+            atype_padding = np.pad(
+                atype,
+                pad_width=((0, 0), (0, padding_atoms)),
+                mode='constant',
+                constant_values=-1
+            )
+            coord_padding = np.pad(
+                coord,
+                pad_width=((0, 0), (0, padding_atoms), (0, 0)),
+                mode='constant',
+                constant_values=0
+            )
+            args = [to_jax_array(ii) for ii in [coord_padding, atype_padding, self.cell]]
+            result_padding = model.call(*args)
+            np.testing.assert_allclose(
+                to_numpy_array(result[f"{var_name}_redu"]),
+                to_numpy_array(result_padding[f"{var_name}_redu"]),
+                atol=self.atol,
+            )
 
 
-        #np.testing.assert_allclose(
-        #    to_numpy_array(ret0[model.get_var_name()]),
-        #    to_numpy_array(ret1[md1.get_var_name()]),
-        #    atol=self.atol,
-        #)
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/test_padding_atoms.py b/source/tests/pt/test_padding_atoms.py

Original file line number	Diff line number	Diff line change
`@@ -355,6 +355,7 @@ def forward_common_atomic(`
`355`	`355`	`self.atomic_output_def(),`
`356`	`356`	`extended_coord,`
`357`	`357`	`do_atomic_virial=do_atomic_virial,`
	`358`	`+ mask=atomic_ret["mask"] if "mask" in atomic_ret else None,`
`358`	`359`	`)`
`359`	`360`
`360`	`361`	`forward_lower = call_lower`