fix: get correct intensive property prediction when using virtual atoms (deepmodeling#4869)

When using virtual atoms, the property output of virtual atom is `0`.
- If predicting energy or other extensive properties, it works well,
that's because the virtual atom property `0` do not contribute to the
total energy or other extensive properties.
- However, if predicting intensive properties, there is some error. For
example, a frame has two real atoms and two virtual atoms, the atomic
property contribution is [2, 2, 0, 0](the atomic property of virtual
atoms are always 0), the final property should be `(2+2)/real_atoms =
2`, not be `(2+2)/total_atoms =1`.

This PR is used to solve this bug mentioned above.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Models now provide accessors to retrieve property names and their
fitting network; property fitting nets expose output definitions.

* **Bug Fixes**
* Intensive property reduction respects atom masks so padded/dummy atoms
are ignored, keeping results invariant to padding.

* **Tests**
* Added PyTorch, JAX, and core tests validating consistent behavior with
padded atoms.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 2 people

deepmd/dpmodel/fitting/property_fitting.py

@@ -12,6 +12,10 @@
 from deepmd.dpmodel.fitting.invar_fitting import (
     InvarFitting,
 )
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -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()

deepmd/dpmodel/model/dp_model.py

@@ -45,3 +45,7 @@ def update_sel(
             train_data, type_map, local_jdata["descriptor"]
         )
         return local_jdata_cpy, min_nbor_dist
+
+    def get_fitting_net(self):
+        """Get the fitting network."""
+        return self.atomic_model.fitting

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

deepmd/dpmodel/model/property_model.py

@@ -25,3 +25,7 @@ def __init__(
     ) -> None:
         DPModelCommon.__init__(self)
         DPPropertyModel_.__init__(self, *args, **kwargs)
+
+    def get_var_name(self) -> str:
+        """Get the name of the property."""
+        return self.get_fitting_net().var_name

deepmd/dpmodel/model/transform_output.py

@@ -1,5 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+from typing import (
+    Optional,
+)
+
 import array_api_compat
 import numpy as np
 
@@ -24,6 +28,7 @@ def fit_output_to_model_output(
     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 +43,19 @@ 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:
+                    model_ret[kk_redu] = xp.sum(
+                        vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+                    ) / np.sum(mask, axis=-1, keepdims=True)
+                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

deepmd/jax/model/base_model.py

@@ -46,7 +46,16 @@ 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:
+                    model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis) / jnp.sum(
+                        mask, axis=-1, keepdims=True
+                    )
+                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:
 

deepmd/pt/model/model/make_model.py

@@ -299,6 +299,7 @@ def forward_common_lower(
                 cc_ext,
                 do_atomic_virial=do_atomic_virial,
                 create_graph=self.training,
+                mask=atomic_ret["mask"] if "mask" in atomic_ret else None,
             )
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict

deepmd/pt/model/model/transform_output.py

@@ -158,6 +158,7 @@ def fit_output_to_model_output(
     coord_ext: torch.Tensor,
     do_atomic_virial: bool = False,
     create_graph: bool = True,
+    mask: Optional[torch.Tensor] = None,
 ) -> dict[str, torch.Tensor]:
     """Transform the output of the fitting network to
     the model output.
@@ -172,7 +173,12 @@ def fit_output_to_model_output(
         if vdef.reducible:
             kk_redu = get_reduce_name(kk)
             if vdef.intensive:
-                model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis)
+                if mask is not None:
+                    model_ret[kk_redu] = torch.sum(
+                        vv.to(redu_prec), dim=atom_axis
+                    ) / torch.sum(mask, dim=-1, keepdim=True)
+                else:
+                    model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis)
             else:
                 model_ret[kk_redu] = torch.sum(vv.to(redu_prec), dim=atom_axis)
             if vdef.r_differentiable:

source/tests/common/dpmodel/test_padding_atoms.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+from copy import (
+    deepcopy,
+)
+
+import numpy as np
+
+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()

source/tests/jax/test_padding_atoms.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import sys
+import unittest
+from copy import (
+    deepcopy,
+)
+
+import numpy as np
+
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
+
+if sys.version_info >= (3, 10):
+    from deepmd.jax.common import (
+        to_jax_array,
+    )
+    from deepmd.jax.descriptor.se_e2_a import (
+        DescrptSeA,
+    )
+    from deepmd.jax.env import (
+        jnp,
+    )
+    from deepmd.jax.fitting.fitting import (
+        PropertyFittingNet,
+    )
+    from deepmd.jax.model.property_model import (
+        PropertyModel,
+    )
+
+    dtype = jnp.float64
+
+
+@unittest.skipIf(
+    sys.version_info < (3, 10),
+    "JAX requires Python 3.10 or later",
+)
+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
+
+
+@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)
+
+    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 = [to_jax_array(ii) for ii in [self.coord, self.atype, self.cell]]
+        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,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()