fix(lmdb): align statistics grouping and decoder shutdown

Author: OutisLi

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.15.22
+    rev: v0.16.0
     hooks:
       - id: ruff
         args: ["--fix"]

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -758,7 +758,7 @@ def change_out_bias(
             delta_bias, out_std = compute_output_stats(
                 sample_merged,
                 self.get_ntypes(),
-                keys=list(self.atomic_output_def().keys()),
+                keys=self.bias_keys,
                 stat_file_path=stat_file_path,
                 model_forward=self._get_forward_wrapper_func(),
                 rcond=self.rcond,
@@ -771,7 +771,7 @@ def change_out_bias(
             bias_out, std_out = compute_output_stats(
                 sample_merged,
                 self.get_ntypes(),
-                keys=list(self.atomic_output_def().keys()),
+                keys=self.bias_keys,
                 stat_file_path=stat_file_path,
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,

deepmd/dpmodel/descriptor/descriptor.py

@@ -110,6 +110,37 @@ def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""
         raise NotImplementedError
 
+    def set_stat_mean_and_stddev(
+        self,
+        mean: Array,
+        stddev: Array,
+    ) -> None:
+        """Update the normalization arrays of the descriptor block.
+
+        Parameters
+        ----------
+        mean
+            Mean of the environment matrix.
+        stddev
+            Standard deviation of the environment matrix.
+
+        Returns
+        -------
+        None
+        """
+        self["davg"] = mean
+        self["dstd"] = stddev
+
+    def get_stat_mean_and_stddev(self) -> tuple[Array, Array]:
+        """Return the normalization arrays of the descriptor block.
+
+        Returns
+        -------
+        tuple[Array, Array]
+            Mean and standard deviation of the environment matrix.
+        """
+        return self["davg"], self["dstd"]
+
     def share_params(
         self, base_class: Any, shared_level: Any, resume: bool = False
     ) -> None:

deepmd/dpmodel/descriptor/dpa1.py

@@ -1461,15 +1461,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self, use_graph=True)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.stddev)

deepmd/dpmodel/descriptor/repflows.py

@@ -494,15 +494,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.stddev)

deepmd/dpmodel/descriptor/repformers.py

@@ -459,15 +459,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.stddev)

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -349,15 +349,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.dstd)

deepmd/dpmodel/descriptor/se_r.py

@@ -328,15 +328,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.dstd)

deepmd/dpmodel/descriptor/se_t.py

@@ -303,15 +303,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.dstd)

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -791,15 +791,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.stddev)