Add unittest

Author: Chengqian-Zhang

deepmd/dpmodel/fitting/property_fitting.py

@@ -88,11 +88,13 @@ def __init__(
         exclude_types: list[int] = [],
         type_map: list[str] | None = None,
         default_fparam: list | None = None,
+        distinguish_types: bool = True,
         # not used
         seed: int | None = None,
     ) -> None:
         self.task_dim = task_dim
         self.intensive = intensive
+        self.distinguish_types = distinguish_types
         super().__init__(
             var_name=property_name,
             ntypes=ntypes,
@@ -131,7 +133,7 @@ def output_def(self) -> FittingOutputDef:
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 5, 1)
+        check_version_compatibility(data.pop("@version"), 6, 1)
         data.pop("dim_out")
         data["property_name"] = data.pop("var_name")
         data.pop("tot_ener_zero")
@@ -150,7 +152,12 @@ def serialize(self) -> dict:
             "type": "property",
             "task_dim": self.task_dim,
             "intensive": self.intensive,
+            "distinguish_types": self.distinguish_types,
         }
-        dd["@version"] = 5
+        dd["@version"] = 6
 
         return dd
+
+    def get_distinguish_types(self) -> bool:
+        """Get whether the fitting net computes stats which are distinguished between different types of atoms."""
+        return self.distinguish_types

deepmd/pt/model/atomic_model/property_atomic_model.py

@@ -26,7 +26,7 @@ def __init__(
 
     def get_compute_stats_distinguish_types(self) -> bool:
         """Get whether the fitting net computes stats which are not distinguished between different types of atoms."""
-        return True
+        return self.fitting_net.get_distinguish_types()
 
     def get_intensive(self) -> bool:
         """Whether the fitting property is intensive."""

deepmd/pt/model/task/property.py

@@ -70,6 +70,8 @@ class PropertyFittingNet(InvarFitting):
         different fitting nets for different atom types.
     seed : int, optional
         Random seed.
+    distinguish_types : bool
+        Whether to distinguish atom types when computing output statistics.
     """
 
     def __init__(
@@ -91,10 +93,12 @@ def __init__(
         trainable: bool | list[bool] = True,
         seed: int | None = None,
         default_fparam: list | None = None,
+        distinguish_types: bool = True,
         **kwargs: Any,
     ) -> None:
         self.task_dim = task_dim
         self.intensive = intensive
+        self.distinguish_types = distinguish_types
         super().__init__(
             var_name=property_name,
             ntypes=ntypes,
@@ -133,10 +137,14 @@ def get_intensive(self) -> bool:
         """Whether the fitting property is intensive."""
         return self.intensive
 
+    def get_distinguish_types(self) -> bool:
+        """Get whether to distinguish atom types when computing output statistics."""
+        return self.distinguish_types
+
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 1), 5, 1)
+        check_version_compatibility(data.pop("@version", 1), 6, 1)
         data.pop("dim_out")
         data["property_name"] = data.pop("var_name")
         obj = super().deserialize(data)
@@ -150,8 +158,9 @@ def serialize(self) -> dict:
             "type": "property",
             "task_dim": self.task_dim,
             "intensive": self.intensive,
+            "distinguish_types": self.distinguish_types,
         }
-        dd["@version"] = 5
+        dd["@version"] = 6
 
         return dd
 

deepmd/utils/argcheck.py

@@ -1925,6 +1925,7 @@ def fitting_property() -> list[Argument]:
     doc_seed = "Random seed for parameter initialization of the fitting net"
     doc_task_dim = "The dimension of outputs of fitting net"
     doc_intensive = "Whether the fitting property is intensive"
+    doc_distinguish_types = "Whether to distinguish atom types when computing output statistics."
     doc_property_name = "The names of fitting property, which should be consistent with the property name in the dataset."
     doc_trainable = "Whether the parameters in the fitting net are trainable. This option can be\n\n\
 - bool: True if all parameters of the fitting net are trainable, False otherwise.\n\n\
@@ -1966,6 +1967,7 @@ def fitting_property() -> list[Argument]:
         Argument("seed", [int, None], optional=True, doc=doc_seed),
         Argument("task_dim", int, optional=True, default=1, doc=doc_task_dim),
         Argument("intensive", bool, optional=True, default=False, doc=doc_intensive),
+        Argument("distinguish_types", bool, optional=True, default=True, doc=doc_distinguish_types),
         Argument(
             "property_name",
             str,

source/tests/common/test_out_stat.py

@@ -91,7 +91,7 @@ def test_compute_stats_from_redu_with_assigned_bias(self) -> None:
         )
 
     def test_compute_stats_do_not_distinguish_types_intensive(self) -> None:
-        """Test compute_stats_property function with intensive scenario."""
+        """Test compute_stats_do_not_distinguish function with intensive scenario."""
         bias, std = compute_stats_do_not_distinguish_types(
             self.output_redu, self.natoms, intensive=True
         )
@@ -110,7 +110,7 @@ def test_compute_stats_do_not_distinguish_types_intensive(self) -> None:
             )
 
     def test_compute_stats_do_not_distinguish_types_extensive(self) -> None:
-        """Test compute_stats_property function with extensive scenario."""
+        """Test compute_stats_do_not_distinguish function with extensive scenario."""
         bias, std = compute_stats_do_not_distinguish_types(
             self.output_redu, self.natoms
         )
@@ -142,6 +142,35 @@ def test_compute_stats_do_not_distinguish_types_extensive(self) -> None:
                 rtol=1e-7,
             )
 
+    def test_compute_stats_from_redu_intensive(self) -> None:
+        """Test compute_stats_from_redu function with intensive scenario."""
+        bias, std = compute_stats_from_redu(
+            self.output_redu, self.natoms, intensive=True,
+        )
+        # Test shapes
+        assert bias.shape == (len(self.mean), self.output_redu.shape[1])
+        assert std.shape == (self.output_redu.shape[1],)
+
+        # Test values
+        np.testing.assert_allclose(
+            bias,
+            np.array(
+                [
+                    [8926338.68432182, 8750110.71559034, 2045325.12109175, 1392024.84192495, 6714978.25878314],
+                    [554163.59820041, 5965821.3924394, 2171555.69509784, 8050760.64873761, 5277414.78728998],
+                    [9180265.02004177, 6836013.36530394, 9121797.79540738, 7801570.3259364, 4095707.84597587]
+                ]
+            ),
+            rtol=1e-6,
+        )
+        np.testing.assert_allclose(
+            std,
+            np.array(
+                [0.01700638, 0.01954897, 0.02028186, 0.01074124, 0.02025821]
+            ),
+            rtol=1e-6,
+        )
+
     def test_compute_stats_from_atomic(self) -> None:
         bias, std = compute_stats_from_atomic(self.output, self.atype)
         np.testing.assert_allclose(bias, self.mean)