feat: add stat for dpmodel's atomic model. implement atomic model for pt_expt

Author: Han Wang

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,196 @@ 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.copy(self.out_bias)
+        out_std_data = np.copy(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 = getattr(ref_array, "device", None)
+            if device is not None:
+                # For torch tensors
+                coord = xp.asarray(coord, device=device)
+                atype = xp.asarray(atype, device=device)
+                if box is not None:
+                    # Check if box is all zeros before converting
+                    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)
+            else:
+                # For numpy arrays
+                coord = xp.asarray(coord)
+                atype = xp.asarray(atype)
+                if box is not None:
+                    if np.allclose(box, 0.0):
+                        box = None
+                    else:
+                        box = xp.asarray(box)
+                if fparam is not None:
+                    fparam = xp.asarray(fparam)
+                if aparam is not None:
+                    aparam = xp.asarray(aparam)
+
+            (
+                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

deepmd/pt_expt/atomic_model/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+from .energy_atomic_model import (
+    DPEnergyAtomicModel,
+)
+
+__all__ = [
+    "DPAtomicModel",
+    "DPEnergyAtomicModel",
+]

deepmd/pt_expt/atomic_model/dp_atomic_model.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+import torch
+
+from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP
+from deepmd.pt_expt.common import (
+    dpmodel_setattr,
+    register_dpmodel_mapping,
+)
+
+
+class DPAtomicModel(DPAtomicModelDP, torch.nn.Module):
+    # Import at class level to set base classes for deserialization
+    # These will be used by the dpmodel deserialize method to create pt_expt instances
+    from deepmd.pt_expt.descriptor.base_descriptor import (
+        BaseDescriptor,
+    )
+    from deepmd.pt_expt.fitting.base_fitting import (
+        BaseFitting,
+    )
+
+    base_descriptor_cls = BaseDescriptor
+    base_fitting_cls = BaseFitting
+
+    def __init__(
+        self, descriptor: Any, fitting: Any, *args: Any, **kwargs: Any
+    ) -> None:
+        torch.nn.Module.__init__(self)
+        # Convert descriptor and fitting to pt_expt versions if they are dpmodel instances
+        # The dpmodel_setattr mechanism will handle this automatically via registry
+        from deepmd.pt_expt.common import (
+            try_convert_module,
+        )
+
+        descriptor_pt = try_convert_module(descriptor)
+        fitting_pt = try_convert_module(fitting)
+        # If conversion failed (not registered), use original (assume already pt_expt)
+        if descriptor_pt is None:
+            descriptor_pt = descriptor
+        if fitting_pt is None:
+            fitting_pt = fitting
+        DPAtomicModelDP.__init__(self, descriptor_pt, fitting_pt, *args, **kwargs)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        # Ensure torch.nn.Module.__call__ drives forward() for export/tracing.
+        return torch.nn.Module.__call__(self, *args, **kwargs)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        handled, value = dpmodel_setattr(self, name, value)
+        if not handled:
+            super().__setattr__(name, value)
+
+    def forward(
+        self,
+        extended_coord: torch.Tensor,
+        extended_atype: torch.Tensor,
+        nlist: torch.Tensor,
+        mapping: torch.Tensor | None = None,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        return self.forward_atomic(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping=mapping,
+            fparam=fparam,
+            aparam=aparam,
+        )
+
+
+register_dpmodel_mapping(
+    DPAtomicModelDP,
+    lambda v: DPAtomicModel.deserialize(v.serialize()),
+)

deepmd/pt_expt/atomic_model/energy_atomic_model.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.atomic_model.energy_atomic_model import (
+    DPEnergyAtomicModel as DPEnergyAtomicModelDP,
+)
+from deepmd.pt_expt.common import (
+    register_dpmodel_mapping,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPEnergyAtomicModel(DPAtomicModel):
+    """Energy atomic model for pt_expt backend.
+
+    This is a thin wrapper around DPAtomicModel that validates
+    the fitting is an EnergyFittingNet or InvarFitting.
+    """
+
+    pass
+
+
+register_dpmodel_mapping(
+    DPEnergyAtomicModelDP,
+    lambda v: DPEnergyAtomicModel.deserialize(v.serialize()),
+)

source/tests/pt_expt/atomic_model/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later