From ced2a3ab2b113cfd89183ba4a1ebb1be3d039f71 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 28 Feb 2024 00:33:18 -0500 Subject: [PATCH] add tests Signed-off-by: Jinzhe Zeng --- deepmd/pt/model/task/ener.py | 4 +- deepmd/tf/fit/dos.py | 4 +- deepmd/tf/fit/ener.py | 4 +- deepmd/utils/data_system.py | 4 +- deepmd/utils/out_stat.py | 2 +- source/tests/common/test_out_stat.py | 81 ++++++++++++++++++++++++++++ source/tests/tf/common.py | 4 +- 7 files changed, 92 insertions(+), 11 deletions(-) create mode 100644 source/tests/common/test_out_stat.py diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index d236a0d42f..ff7ae6f8ec 100644 --- a/deepmd/pt/model/task/ener.py +++ b/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, diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index a09878f56b..0cc5a7df62 100644 --- a/deepmd/tf/fit/dos.py +++ b/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, diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index d5b22ecc1f..5f17f95e6d 100644 --- a/deepmd/tf/fit/ener.py +++ b/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, diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index a4272c7de8..592b1f9748 100644 --- a/deepmd/utils/data_system.py +++ b/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, diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index b76cddc954..8f68e32417 100644 --- a/deepmd/utils/out_stat.py +++ b/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, diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py new file mode 100644 index 0000000000..df8bfdf4d8 --- /dev/null +++ b/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, + ) diff --git a/source/tests/tf/common.py b/source/tests/tf/common.py index 5c6ced4a77..0bcb29b4b5 100644 --- a/source/tests/tf/common.py +++ b/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,