Merge remote-tracking branch 'origin/feat-other-full-model' into feat-other-full-model

Author: Han Wang

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -53,13 +53,15 @@ def __init__(
         pair_exclude_types: list[tuple[int, int]] = [],
         rcond: float | None = None,
         preset_out_bias: dict[str, Array] | None = None,
+        data_stat_protect: float = 1e-2,
     ) -> None:
         super().__init__()
         self.type_map = type_map
         self.reinit_atom_exclude(atom_exclude_types)
         self.reinit_pair_exclude(pair_exclude_types)
         self.rcond = rcond
         self.preset_out_bias = preset_out_bias
+        self.data_stat_protect = data_stat_protect
 
     def init_out_stat(self) -> None:
         """Initialize the output bias."""

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -232,6 +232,30 @@ def change_type_map(
         )
         self.fitting_net.change_type_map(type_map=type_map)
 
+    def compute_fitting_input_stat(
+        self,
+        sample_merged: Callable[[], list[dict]] | list[dict],
+        stat_file_path: DPPath | None = None,
+    ) -> None:
+        """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.
+
+        Parameters
+        ----------
+        sample_merged : Union[Callable[[], list[dict]], list[dict]]
+            - list[dict]: A list of data samples from various data systems.
+                Each element, ``merged[i]``, is a data dictionary containing
+                ``keys``: ``np.ndarray`` originating from the ``i``-th data system.
+            - Callable[[], list[dict]]: A lazy function that returns data samples
+                in the above format only when needed.
+        stat_file_path : Optional[DPPath]
+            The path to the stat file.
+        """
+        self.fitting.compute_input_stats(
+            sample_merged,
+            protection=self.data_stat_protect,
+            stat_file_path=stat_file_path,
+        )
+
     def serialize(self) -> dict:
         dd = super().serialize()
         dd.update(

deepmd/dpmodel/fitting/general_fitting.py

@@ -437,11 +437,14 @@ def change_type_map(
         self.ntypes = len(type_map)
         self.reinit_exclude(map_atom_exclude_types(self.exclude_types, remap_index))
         if has_new_type:
+            xp = array_api_compat.array_namespace(self.bias_atom_e)
             extend_shape = [len(type_map), *list(self.bias_atom_e.shape[1:])]
-            extend_bias_atom_e = np.zeros(extend_shape, dtype=self.bias_atom_e.dtype)
-            self.bias_atom_e = np.concatenate(
-                [self.bias_atom_e, extend_bias_atom_e], axis=0
+            extend_bias_atom_e = xp.zeros(
+                extend_shape,
+                dtype=self.bias_atom_e.dtype,
+                device=array_api_compat.device(self.bias_atom_e),
             )
+            self.bias_atom_e = xp.concat([self.bias_atom_e, extend_bias_atom_e], axis=0)
         self.bias_atom_e = self.bias_atom_e[remap_index]
 
     def __setitem__(self, key: str, value: Any) -> None:

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -237,14 +237,21 @@ def change_type_map(
         remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map)
         super().change_type_map(type_map=type_map)
         if has_new_type:
+            xp = array_api_compat.array_namespace(self.scale)
             extend_shape = [len(type_map), *list(self.scale.shape[1:])]
-            extend_scale = np.ones(extend_shape, dtype=self.scale.dtype)
-            self.scale = np.concatenate([self.scale, extend_scale], axis=0)
+            extend_scale = xp.ones(
+                extend_shape,
+                dtype=self.scale.dtype,
+                device=array_api_compat.device(self.scale),
+            )
+            self.scale = xp.concat([self.scale, extend_scale], axis=0)
             extend_shape = [len(type_map), *list(self.constant_matrix.shape[1:])]
-            extend_constant_matrix = np.zeros(
-                extend_shape, dtype=self.constant_matrix.dtype
+            extend_constant_matrix = xp.zeros(
+                extend_shape,
+                dtype=self.constant_matrix.dtype,
+                device=array_api_compat.device(self.constant_matrix),
             )
-            self.constant_matrix = np.concatenate(
+            self.constant_matrix = xp.concat(
                 [self.constant_matrix, extend_constant_matrix], axis=0
             )
         self.scale = self.scale[remap_index]

deepmd/dpmodel/model/make_model.py

@@ -417,6 +417,8 @@ def change_out_bias(
                 'change-by-statistic' or 'set-by-statistic'.
             """
             self.atomic_model.change_out_bias(merged, bias_adjust_mode=bias_adjust_mode)
+            if bias_adjust_mode == "set-by-statistic":
+                self.atomic_model.compute_fitting_input_stat(merged)
 
         def _input_type_cast(
             self,
@@ -616,12 +618,17 @@ def do_grad_c(
             return self.atomic_model.do_grad_c(var_name)
 
         def change_type_map(
-            self, type_map: list[str], model_with_new_type_stat: Any = None
+            self, type_map: list[str], model_with_new_type_stat: Any | None = None
         ) -> None:
             """Change the type related params to new ones, according to `type_map` and the original one in the model.
             If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
             """
-            self.atomic_model.change_type_map(type_map=type_map)
+            self.atomic_model.change_type_map(
+                type_map=type_map,
+                model_with_new_type_stat=model_with_new_type_stat.atomic_model
+                if model_with_new_type_stat is not None
+                else None,
+            )
 
         def serialize(self) -> dict:
             return self.atomic_model.serialize()