Solve dptest UT, dpatomicmodel UT, code scannisang

deepmd/dpmodel/atomic_model/__init__.py

@@ -42,6 +42,9 @@
 from .polar_atomic_model import (
     DPPolarAtomicModel,
 )
+from .property_atomic_model import (
+    DPPropertyAtomicModel,
+)
 
 __all__ = [
     "BaseAtomicModel",
@@ -54,4 +57,5 @@
     "LinearEnergyAtomicModel",
     "PairTabAtomicModel",
     "make_base_atomic_model",
+    "DPPropertyAtomicModel",
 ]

deepmd/dpmodel/atomic_model/property_atomic_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
 from deepmd.dpmodel.fitting.property_fitting import (
     PropertyFittingNet,
 )
@@ -15,3 +16,23 @@ def __init__(self, descriptor, fitting, type_map, **kwargs):
                 "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel"
             )
         super().__init__(descriptor, fitting, type_map, **kwargs)
+
+    def apply_out_stat(
+        self,
+        ret: dict[str, np.ndarray],
+        atype: np.ndarray,
+    ):
+        """Apply the stat to each atomic output.
+
+        Parameters
+        ----------
+        ret
+            The returned dict by the forward_atomic method
+        atype
+            The atom types. nf x nloc
+
+        """
+        out_bias, out_std = self._fetch_out_stat(self.bias_keys)
+        for kk in self.bias_keys:
+            ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0]
+        return ret

deepmd/dpmodel/model/property_model.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from deepmd.dpmodel.atomic_model.dp_atomic_model import (
-    DPAtomicModel,
+from deepmd.dpmodel.atomic_model import (
+    DPPropertyAtomicModel,
 )
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
@@ -13,7 +13,7 @@
     make_model,
 )
 
-DPPropertyModel_ = make_model(DPAtomicModel)
+DPPropertyModel_ = make_model(DPPropertyAtomicModel)
 
 
 @BaseModel.register("property")

deepmd/entrypoints/test.py

@@ -779,9 +779,15 @@ def test_property(
     tuple[list[np.ndarray], list[int]]
         arrays with results and their shapes
     """
-    data.add("property", dp.task_dim, atomic=False, must=True, high_prec=True)
-    if has_atom_property:
-        data.add("atom_property", dp.task_dim, atomic=True, must=False, high_prec=True)
+    property_name = dp.get_property_name()
+    property_dim = dp.get_property_dim()
+    assert isinstance(property_name, list)
+    assert isinstance(property_dim, list)
+    assert sum(property_dim) == dp.task_dim
+    for (name, dim) in zip(property_name, property_dim):
+        data.add(name, dim, atomic=False, must=True, high_prec=True)
+        if has_atom_property:
+            data.add(f"atom_{name}", dim, atomic=True, must=False, high_prec=True)
 
     if dp.get_dim_fparam() > 0:
         data.add(
@@ -832,6 +838,18 @@ def test_property(
         aproperty = ret[1]
         aproperty = aproperty.reshape([numb_test, natoms * dp.task_dim])
 
+    concat_property = []
+    concat_aproperty = []
+    for (name, dim) in zip(property_name, property_dim):
+        test_data[name] = test_data[name].reshape([numb_test, dim])
+        test_data[f"atom_{name}"] = test_data[f"atom_{name}"].reshape(
+            [numb_test, natoms * dim]
+        )
+        concat_property.append(test_data[name])
+        concat_aproperty.append(test_data[f"atom_{name}"])
+    test_data["property"] = np.concatenate(concat_property, axis=1)
+    test_data["atom_property"] = np.concatenate(concat_aproperty, axis=1)
+
     diff_property = property - test_data["property"][:numb_test]
     mae_property = mae(diff_property)
     rmse_property = rmse(diff_property)

deepmd/infer/deep_property.py

@@ -141,5 +141,13 @@ def get_intensive(self) -> bool:
         """Get whether the property is intensive."""
         return self.deep_eval.get_intensive()
 
+    def get_property_name(self) -> Union[list[str], str]:
+        """Get the name of the property."""
+        return self.deep_eval.get_property_name()
+
+    def get_property_dim(self) -> Union[list[int], int]:
+        """Get the dimension of the property."""
+        return self.deep_eval.get_property_dim()
+
 
 __all__ = ["DeepProperty"]

deepmd/pt/infer/deep_eval.py

@@ -184,6 +184,14 @@ def get_dim_aparam(self) -> int:
     def get_intensive(self) -> bool:
         return self.dp.model["Default"].get_intensive()
 
+    def get_property_name(self) -> Union[list[str], str]:
+        """Get the name of the property."""
+        return self.dp.model["Default"].get_property_name()
+
+    def get_property_dim(self) -> Union[list[int], int]:
+        """Get the dimension of the property."""
+        return self.dp.model["Default"].get_property_dim()
+
     @property
     def model_type(self) -> type["DeepEvalWrapper"]:
         """The the evaluator of the model type."""

deepmd/pt/model/model/property_model.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Optional,
+    Union
 )
 
 import torch
@@ -78,6 +79,16 @@ def get_intensive(self) -> bool:
         """Get whether the property is intensive."""
         return self.model_output_def()["property"].intensive
 
+    @torch.jit.export
+    def get_property_name(self) -> Union[list[str], str]:
+        """Get the name of the property."""
+        return self.get_fitting_net().property_name
+
+    @torch.jit.export
+    def get_property_dim(self) -> Union[list[int], int]:
+        """Get the dimension of the property."""
+        return self.get_fitting_net().property_dim
+
     @torch.jit.export
     def forward_lower(
         self,

deepmd/pt/model/task/property.py

@@ -84,6 +84,7 @@ def __init__(
         intensive: bool = False,
         bias_method: str = "normal",
         property_name: Union[str, list] = "property",
+        property_dim: Union[int, list] = 1,
         resnet_dt: bool = True,
         numb_fparam: int = 0,
         numb_aparam: int = 0,
@@ -97,7 +98,12 @@ def __init__(
         self.task_dim = task_dim
         self.intensive = intensive
         self.bias_method = bias_method
+        if isinstance(property_name, str):
+            property_name = [property_name]
         self.property_name = property_name
+        if isinstance(property_dim, int):
+            property_dim = [property_dim]
+        self.property_dim = property_dim
         super().__init__(
             var_name="property",
             ntypes=ntypes,

deepmd/pt/utils/stat.py

@@ -296,7 +296,8 @@ def compute_output_stats(
             for key in keys:
                 if atomic_output.var_defs[key].sub_var_name is not None:
                     sub_keys.extend(atomic_output.var_defs[key].sub_var_name)
-            keys.extend(sub_keys)
+            del keys
+            keys = sub_keys
         new_keys = [
             ii
             for ii in keys

source/tests/pt/model/test_permutation.py

@@ -332,6 +332,7 @@
         "type": "property",
         "task_dim": 3,
         "property_name": ["band_property"],
+        "property_dim": 3,
         "neuron": [24, 24, 24],
         "resnet_dt": True,
         "bias_method": "normal",

source/tests/pt/property/input.json

@@ -28,6 +28,7 @@
       "type": "property",
       "intensive": true,
       "property_name": "band_property",
+      "property_dim": 3,
       "task_dim": 3,
       "neuron": [
         100,