feat(pt_expt): atomic model

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import math
+from collections.abc import (
+    Callable,
+)
 from typing import (
     Any,
 )
@@ -30,6 +33,9 @@
     map_atom_exclude_types,
     map_pair_exclude_types,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 
 from .make_base_atomic_model import (
     make_base_atomic_model,
@@ -246,6 +252,180 @@ def call(
             aparam=aparam,
         )
 
+    def get_intensive(self) -> bool:
+        """Whether the fitting property is intensive."""
+        return False
+
+    def get_compute_stats_distinguish_types(self) -> bool:
+        """Get whether the fitting net computes stats which are not distinguished between different types of atoms."""
+        return True
+
+    def compute_or_load_out_stat(
+        self,
+        merged: Callable[[], list[dict]] | list[dict],
+        stat_file_path: DPPath | None = None,
+    ) -> None:
+        """
+        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
+
+        Parameters
+        ----------
+        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. Since the sampling process can be slow and memory-intensive,
+                the lazy function helps by only sampling once.
+        stat_file_path : Optional[DPPath]
+            The path to the stat file.
+
+        """
+        self.change_out_bias(
+            merged,
+            stat_file_path=stat_file_path,
+            bias_adjust_mode="set-by-statistic",
+        )
+
+    def change_out_bias(
+        self,
+        sample_merged: Callable[[], list[dict]] | list[dict],
+        stat_file_path: DPPath | None = None,
+        bias_adjust_mode: str = "change-by-statistic",
+    ) -> None:
+        """Change the output bias according to the input data and the pretrained model.
+
+        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. Since the sampling process can be slow and memory-intensive,
+                the lazy function helps by only sampling once.
+        bias_adjust_mode : str
+            The mode for changing output bias : ['change-by-statistic', 'set-by-statistic']
+            'change-by-statistic' : perform predictions on labels of target dataset,
+                    and do least square on the errors to obtain the target shift as bias.
+            'set-by-statistic' : directly use the statistic output bias in the target dataset.
+        stat_file_path : Optional[DPPath]
+            The path to the stat file.
+        """
+        from deepmd.dpmodel.utils.stat import (
+            compute_output_stats,
+        )
+
+        if bias_adjust_mode == "change-by-statistic":
+            delta_bias, out_std = compute_output_stats(
+                sample_merged,
+                self.get_ntypes(),
+                keys=list(self.atomic_output_def().keys()),
+                stat_file_path=stat_file_path,
+                model_forward=self._get_forward_wrapper_func(),
+                rcond=self.rcond,
+                preset_bias=self.preset_out_bias,
+                stats_distinguish_types=self.get_compute_stats_distinguish_types(),
+                intensive=self.get_intensive(),
+            )
+            self._store_out_stat(delta_bias, out_std, add=True)
+        elif bias_adjust_mode == "set-by-statistic":
+            bias_out, std_out = compute_output_stats(
+                sample_merged,
+                self.get_ntypes(),
+                keys=list(self.atomic_output_def().keys()),
+                stat_file_path=stat_file_path,
+                rcond=self.rcond,
+                preset_bias=self.preset_out_bias,
+                stats_distinguish_types=self.get_compute_stats_distinguish_types(),
+                intensive=self.get_intensive(),
+            )
+            self._store_out_stat(bias_out, std_out)
+        else:
+            raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode)
+
+    def _store_out_stat(
+        self,
+        out_bias: dict[str, np.ndarray],
+        out_std: dict[str, np.ndarray],
+        add: bool = False,
+    ) -> None:
+        """Store output bias and std into the model."""
+        ntypes = self.get_ntypes()
+        out_bias_data = np.array(to_numpy_array(self.out_bias))
+        out_std_data = np.array(to_numpy_array(self.out_std))
+        for kk in out_bias.keys():
+            assert kk in out_std.keys()
+            idx = self._get_bias_index(kk)
+            size = self._varsize(self.atomic_output_def()[kk].shape)
+            if not add:
+                out_bias_data[idx, :, :size] = out_bias[kk].reshape(ntypes, size)
+            else:
+                out_bias_data[idx, :, :size] += out_bias[kk].reshape(ntypes, size)
+            out_std_data[idx, :, :size] = out_std[kk].reshape(ntypes, size)
+        self.out_bias = out_bias_data
+        self.out_std = out_std_data
+
+    def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]:
+        """Get a forward wrapper of the atomic model for output bias calculation."""
+        import array_api_compat
+
+        from deepmd.dpmodel.utils.nlist import (
+            extend_input_and_build_neighbor_list,
+        )
+
+        def model_forward(
+            coord: np.ndarray,
+            atype: np.ndarray,
+            box: np.ndarray | None,
+            fparam: np.ndarray | None = None,
+            aparam: np.ndarray | None = None,
+        ) -> dict[str, np.ndarray]:
+            # Get reference array to determine the target array type and device
+            # Use out_bias as reference since it's always present
+            ref_array = self.out_bias
+            xp = array_api_compat.array_namespace(ref_array)
+
+            # Convert numpy inputs to the model's array type with correct device
+            device = array_api_compat.device(ref_array)
+            coord = xp.asarray(coord, device=device)
+            atype = xp.asarray(atype, device=device)
+            if box is not None:
+                if np.allclose(box, 0.0):
+                    box = None
+                else:
+                    box = xp.asarray(box, device=device)
+            if fparam is not None:
+                fparam = xp.asarray(fparam, device=device)
+            if aparam is not None:
+                aparam = xp.asarray(aparam, device=device)
+
+            (
+                extended_coord,
+                extended_atype,
+                mapping,
+                nlist,
+            ) = extend_input_and_build_neighbor_list(
+                coord,
+                atype,
+                self.get_rcut(),
+                self.get_sel(),
+                mixed_types=self.mixed_types(),
+                box=box,
+            )
+            atomic_ret = self.forward_common_atomic(
+                extended_coord,
+                extended_atype,
+                nlist,
+                mapping=mapping,
+                fparam=fparam,
+                aparam=aparam,
+            )
+            # Convert outputs back to numpy arrays
+            return {kk: to_numpy_array(vv) for kk, vv in atomic_ret.items()}
+
+        return model_forward
+
     def serialize(self) -> dict:
         return {
             "type_map": self.type_map,

deepmd/dpmodel/common.py

@@ -121,7 +121,8 @@ def to_numpy_array(x: Optional["Array"]) -> np.ndarray | None:
     try:
         # asarray is not within Array API standard, so may fail
         return np.asarray(x)
-    except (ValueError, AttributeError, TypeError):
+    except (ValueError, AttributeError, TypeError, RuntimeError):
+        # RuntimeError: handles torch tensors with requires_grad=True
         xp = array_api_compat.array_namespace(x)
         # to fix BufferError: Cannot export readonly array since signalling readonly is unsupported by DLPack.
         # Move to CPU device to ensure numpy compatibility