Merge branch 'master' into 0224-observered-type

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>

Author: iProzd

deepmd/dpmodel/array_api.py

@@ -78,7 +78,7 @@ def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array:
     xp = array_api_compat.array_namespace(input)
 
     # Create flat index array matching input shape
-    idx = xp.arange(input.size, dtype=xp.int64)
+    idx = xp.arange(input.size, dtype=xp.int64, device=array_api_compat.device(input))
     idx = xp.reshape(idx, input.shape)
 
     # Get flat indices where we want to add values
@@ -190,6 +190,10 @@ def xp_bincount(x: Array, weights: Array | None = None, minlength: int = 0) -> A
     else:
         if weights is None:
             weights = xp.ones_like(x)
-        result = xp.zeros((max(minlength, int(xp.max(x)) + 1),), dtype=weights.dtype)
+        result = xp.zeros(
+            (max(minlength, int(xp.max(x)) + 1),),
+            dtype=weights.dtype,
+            device=array_api_compat.device(weights),
+        )
         result = xp_add_at(result, x, weights)
     return result

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import functools
 import math
 from collections.abc import (
     Callable,
@@ -52,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."""
@@ -77,6 +80,14 @@ def init_out_stat(self) -> None:
         self.out_bias = out_bias_data
         self.out_std = out_std_data
 
+    def get_out_bias(self) -> Array:
+        """Get the output bias."""
+        return self.out_bias
+
+    def set_out_bias(self, out_bias: Array) -> None:
+        """Set the output bias."""
+        self.out_bias = out_bias
+
     def __setitem__(self, key: str, value: Array) -> None:
         if key in ["out_bias"]:
             self.out_bias = value
@@ -287,6 +298,57 @@ def compute_or_load_out_stat(
             bias_adjust_mode="set-by-statistic",
         )
 
+    def _make_wrapped_sampler(
+        self,
+        sampled_func: Callable[[], list[dict]],
+    ) -> Callable[[], list[dict]]:
+        """Wrap the sampled function with exclusion types and default fparam.
+
+        The returned callable is cached so that the sampling (which may be
+        expensive) is performed at most once.
+
+        Parameters
+        ----------
+        sampled_func
+            The lazy sampled function to get data frames from different data
+            systems.
+
+        Returns
+        -------
+        Callable[[], list[dict]]
+            A cached wrapper around *sampled_func* that additionally sets
+            ``pair_exclude_types``, ``atom_exclude_types`` and default
+            ``fparam`` on every sample dict when applicable.
+        """
+
+        @functools.lru_cache
+        def wrapped_sampler() -> list[dict]:
+            sampled = sampled_func()
+            if self.pair_excl is not None:
+                pair_exclude_types = self.pair_excl.get_exclude_types()
+                for sample in sampled:
+                    sample["pair_exclude_types"] = list(pair_exclude_types)
+            if self.atom_excl is not None:
+                atom_exclude_types = self.atom_excl.get_exclude_types()
+                for sample in sampled:
+                    sample["atom_exclude_types"] = list(atom_exclude_types)
+            if (
+                "find_fparam" not in sampled[0]
+                and "fparam" not in sampled[0]
+                and self.has_default_fparam()
+            ):
+                default_fparam = self.get_default_fparam()
+                if default_fparam is not None:
+                    default_fparam_np = np.array(default_fparam)
+                    for sample in sampled:
+                        nframe = sample["atype"].shape[0]
+                        sample["fparam"] = np.tile(
+                            default_fparam_np.reshape(1, -1), (nframe, 1)
+                        )
+            return sampled
+
+        return wrapped_sampler
+
     def change_out_bias(
         self,
         sample_merged: Callable[[], list[dict]] | list[dict],

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from collections.abc import (
+    Callable,
+)
 from typing import (
     Any,
 )
@@ -15,6 +18,9 @@
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -26,7 +32,21 @@
 
 @BaseAtomicModel.register("standard")
 class DPAtomicModel(BaseAtomicModel):
-    """Model give atomic prediction of some physical property.
+    r"""Model give atomic prediction of some physical property.
+
+    The atomic model computes atomic properties by first extracting a descriptor
+    from the atomic environment, then passing it through a fitting network:
+
+    .. math::
+        \mathcal{D}^i = \mathcal{D}(\mathbf{R}^i, \mathbf{R}_j, \alpha_j),
+
+    .. math::
+        \mathbf{y}^i = \mathcal{F}(\mathcal{D}^i),
+
+    where :math:`\mathcal{D}^i` is the descriptor for atom :math:`i`,
+    :math:`\alpha_j` is the atom type of neighbor :math:`j`,
+    :math:`\mathcal{F}` is the fitting network, and
+    :math:`\mathbf{y}^i` is the predicted atomic property (energy, dipole, etc.).
 
     Parameters
     ----------
@@ -48,17 +68,16 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         super().__init__(type_map, **kwargs)
-        self.type_map = type_map
         self.descriptor = descriptor
-        self.fitting = fitting
-        if hasattr(self.fitting, "reinit_exclude"):
-            self.fitting.reinit_exclude(self.atom_exclude_types)
+        self.fitting_net = fitting
+        if hasattr(self.fitting_net, "reinit_exclude"):
+            self.fitting_net.reinit_exclude(self.atom_exclude_types)
         self.type_map = type_map
         super().init_out_stat()
 
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
-        return self.fitting.output_def()
+        return self.fitting_net.output_def()
 
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
@@ -73,7 +92,7 @@ def set_case_embd(self, case_idx: int) -> None:
         Set the case embedding of this atomic model by the given case_idx,
         typically concatenated with the output of the descriptor and fed into the fitting net.
         """
-        self.fitting.set_case_embd(case_idx)
+        self.fitting_net.set_case_embd(case_idx)
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -166,7 +185,7 @@ def forward_atomic(
             nlist,
             mapping=mapping,
         )
-        ret = self.fitting(
+        ret = self.fitting_net(
             descriptor,
             atype,
             gr=rot_mat,
@@ -177,6 +196,37 @@ def forward_atomic(
         )
         return ret
 
+    def compute_or_load_stat(
+        self,
+        sampled_func: Callable[[], list[dict]],
+        stat_file_path: DPPath | None = None,
+        compute_or_load_out_stat: bool = True,
+    ) -> None:
+        """Compute or load the statistics parameters of the model,
+        such as mean and standard deviation of descriptors or the energy bias of the fitting net.
+
+        Parameters
+        ----------
+        sampled_func
+            The lazy sampled function to get data frames from different data systems.
+        stat_file_path
+            The path to the stat file.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics
+            (e.g. mean and standard deviation of descriptors).
+        """
+        if stat_file_path is not None and self.type_map is not None:
+            stat_file_path /= " ".join(self.type_map)
+
+        wrapped_sampler = self._make_wrapped_sampler(sampled_func)
+        self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
+        self.fitting_net.compute_input_stats(
+            wrapped_sampler, stat_file_path=stat_file_path
+        )
+        if compute_or_load_out_stat:
+            self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
+
     def change_type_map(
         self, type_map: list[str], model_with_new_type_stat: Any | None = None
     ) -> None:
@@ -193,7 +243,31 @@ def change_type_map(
             if model_with_new_type_stat is not None
             else None,
         )
-        self.fitting.change_type_map(type_map=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_net.compute_input_stats(
+            sample_merged,
+            protection=self.data_stat_protect,
+            stat_file_path=stat_file_path,
+        )
 
     def serialize(self) -> dict:
         dd = super().serialize()
@@ -204,7 +278,7 @@ def serialize(self) -> dict:
                 "@version": 2,
                 "type_map": self.type_map,
                 "descriptor": self.descriptor.serialize(),
-                "fitting": self.fitting.serialize(),
+                "fitting": self.fitting_net.serialize(),
             }
         )
         return dd
@@ -230,19 +304,19 @@ def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel":
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
-        return self.fitting.get_dim_fparam()
+        return self.fitting_net.get_dim_fparam()
 
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
-        return self.fitting.get_dim_aparam()
+        return self.fitting_net.get_dim_aparam()
 
     def has_default_fparam(self) -> bool:
         """Check if the model has default frame parameters."""
-        return self.fitting.has_default_fparam()
+        return self.fitting_net.has_default_fparam()
 
     def get_default_fparam(self) -> list[float] | None:
         """Get the default frame parameters."""
-        return self.fitting.get_default_fparam()
+        return self.fitting_net.get_default_fparam()
 
     def get_sel_type(self) -> list[int]:
         """Get the selected atom types of this model.
@@ -251,7 +325,7 @@ def get_sel_type(self) -> list[int]:
         to the result of the model.
         If returning an empty list, all atom types are selected.
         """
-        return self.fitting.get_sel_type()
+        return self.fitting_net.get_sel_type()
 
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).