add tests

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/pt/model/task/ener.py

@@ -32,7 +32,7 @@
     to_numpy_array,
 )
 from deepmd.utils.out_stat import (
-    compute_bias_from_redu,
+    compute_stats_from_redu,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -160,7 +160,7 @@ def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None):
                 )
             else:
                 assigned_atom_ener = None
-            bias_atom_e, _ = compute_bias_from_redu(
+            bias_atom_e, _ = compute_stats_from_redu(
                 merged_energy,
                 merged_natoms,
                 assigned_bias=assigned_atom_ener,

deepmd/tf/fit/dos.py

@@ -44,7 +44,7 @@
     one_layer_rand_seed_shift,
 )
 from deepmd.utils.out_stat import (
-    compute_bias_from_redu,
+    compute_stats_from_redu,
 )
 
 log = logging.getLogger(__name__)
@@ -228,7 +228,7 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
         sys_tynatom = np.reshape(sys_tynatom, [nsys, -1])
         sys_tynatom = sys_tynatom[:, 2:]
 
-        dos_shift, _ = compute_bias_from_redu(
+        dos_shift, _ = compute_stats_from_redu(
             sys_dos,
             sys_tynatom,
             rcond=rcond,

deepmd/tf/fit/ener.py

@@ -54,7 +54,7 @@
     Spin,
 )
 from deepmd.utils.out_stat import (
-    compute_bias_from_redu,
+    compute_stats_from_redu,
 )
 
 if TYPE_CHECKING:
@@ -299,7 +299,7 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
             )
         else:
             assigned_atom_ener = None
-        energy_shift, _ = compute_bias_from_redu(
+        energy_shift, _ = compute_stats_from_redu(
             sys_ener.reshape(-1, 1),
             sys_tynatom,
             assigned_bias=assigned_atom_ener,

deepmd/utils/data_system.py

@@ -23,7 +23,7 @@
     DeepmdData,
 )
 from deepmd.utils.out_stat import (
-    compute_bias_from_redu,
+    compute_stats_from_redu,
 )
 
 log = logging.getLogger(__name__)
@@ -251,7 +251,7 @@ def compute_energy_shift(self, rcond=None, key="energy"):
         sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION)
         sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1])
         sys_tynatom = sys_tynatom[:, 2:]
-        energy_shift, _ = compute_bias_from_redu(
+        energy_shift, _ = compute_stats_from_redu(
             sys_ener.reshape(-1, 1),
             sys_tynatom,
             rcond=rcond,

deepmd/utils/out_stat.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 
-def compute_bias_from_redu(
+def compute_stats_from_redu(
     output_redu: np.ndarray,
     natoms: np.ndarray,
     assigned_bias: Optional[np.ndarray] = None,

source/tests/common/test_out_stat.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.utils.out_stat import (
+    compute_stats_from_atomic,
+    compute_stats_from_redu,
+)
+
+
+class TestOutStat(unittest.TestCase):
+    def setUp(self) -> None:
+        rng = np.random.default_rng(20240227)
+        ndim = 5
+        nframes = 1000
+        ntypes = 3
+        nloc = 1000
+        self.atype = rng.integers(0, ntypes, size=(nframes, nloc))
+        # compute the number of atoms for each type in each frame
+        self.natoms = np.zeros((nframes, ntypes), dtype=np.int64)
+        for i in range(ntypes):
+            self.natoms[:, i] = (self.atype == i).sum(axis=1)
+        self.mean = rng.random((ntypes, ndim)) * 1e4
+        self.std = rng.random((ntypes, ndim)) * 1e-3
+
+        # generate random output
+        self.output = rng.normal(
+            loc=self.mean[self.atype, :],
+            scale=self.std[self.atype, :],
+            size=(nframes, nloc, ndim),
+        )
+        self.output_redu = self.output.sum(axis=1)
+
+        return super().setUp()
+
+    def test_compute_stats_from_redu(self):
+        bias, std = compute_stats_from_redu(self.output_redu, self.natoms)
+        np.testing.assert_allclose(bias, self.mean, rtol=1e-7)
+        np.testing.assert_allclose(
+            std,
+            np.sqrt(self.natoms.mean(axis=0) @ np.square(self.std)),
+            rtol=1e-1,
+        )
+        # ensure the sum is close
+        np.testing.assert_allclose(
+            self.output_redu,
+            self.natoms @ bias,
+            rtol=1e-7,
+        )
+
+    def test_compute_stats_from_redu_with_assigned_bias(self):
+        assigned_bias = np.full_like(self.mean, np.nan)
+        assigned_bias[0] = self.mean[0]
+        bias, std = compute_stats_from_redu(
+            self.output_redu,
+            self.natoms,
+            assigned_bias=assigned_bias,
+        )
+        np.testing.assert_allclose(bias, self.mean, rtol=1e-7)
+        np.testing.assert_allclose(bias[0], self.mean[0], rtol=1e-14)
+        np.testing.assert_allclose(
+            std,
+            np.sqrt(self.natoms.mean(axis=0) @ np.square(self.std)),
+            rtol=1e-1,
+        )
+        # ensure the sum is close
+        np.testing.assert_allclose(
+            self.output_redu,
+            self.natoms @ bias,
+            rtol=1e-7,
+        )
+
+    def test_compute_stats_from_atomic(self):
+        bias, std = compute_stats_from_atomic(self.output, self.atype)
+        np.testing.assert_allclose(bias, self.mean)
+        np.testing.assert_allclose(
+            std,
+            self.std,
+            rtol=1e-2,
+        )

source/tests/tf/common.py

@@ -18,7 +18,7 @@
 )
 from deepmd.tf.utils import random as dp_random
 from deepmd.utils.out_stat import (
-    compute_bias_from_redu,
+    compute_stats_from_redu,
 )
 
 if GLOBAL_NP_FLOAT_PRECISION == np.float32:
@@ -1044,7 +1044,7 @@ def compute_energy_shift(self):
         sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION)
         sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1])
         sys_tynatom = sys_tynatom[:, 2:]
-        energy_shift, _ = compute_bias_from_redu(
+        energy_shift, _ = compute_stats_from_redu(
             sys_ener.reshape(-1, 1),
             sys_tynatom,
             rcond=None,