feat(pt_expt): add dipole, polar, dos, property and dp-zbl models with cross-backend consistency tests (#5260)

## Summary

- Add 5 new pt_expt model types: DipoleModel, PolarModel, DOSModel,
PropertyModel, and DPZBLModel, completing pt_expt's model coverage to
parity with pt
- Refactor dpmodel base architecture so that pt_expt models inherit
directly from dpmodel via `make_model()`, removing the intermediate
pt_expt atomic model layer
- Consolidate scattered model-level methods (`get_out_bias`,
`set_out_bias`, `get_observed_type_list`, `compute_or_load_stat`) into
shared dpmodel base classes
- Move `compute_fitting_input_stat` for `set-by-statistic` mode from
model-level `change_out_bias` to training-level `model_change_out_bias`
(pt and pd backends), keeping the `change_out_bias` logic focused on
bias only (copied from #5266)
- Fix array-api-compat violations in `general_fitting.change_type_map`
(bare `np.zeros`/`np.ones`/`np.concatenate` → `xp` equivalents with
device)
- Fix dpmodel `change_type_map` not forwarding
`model_with_new_type_stat` through the call chain
- Add comprehensive cross-backend consistency tests for all model types
(dp vs pt vs pt_expt), covering: model output, serialization round-trip,
`change_out_bias`, `change_type_map`, `compute_or_load_stat`, model API
methods.

  ## Changes

  ### New pt_expt models
  - `deepmd/pt_expt/model/dipole_model.py`
  - `deepmd/pt_expt/model/polar_model.py`
  - `deepmd/pt_expt/model/dos_model.py`
  - `deepmd/pt_expt/model/property_model.py`
  - `deepmd/pt_expt/model/dp_zbl_model.py`

  ### Architecture refactoring
- Remove `deepmd/pt_expt/atomic_model/` layer — models now wrap dpmodel
atomic models directly
- Clean up `BaseModel`: remove concrete methods/data, add plugin
registry
- Refactor `make_model` so backends (dp, pt_expt) inherit shared model
logic from dpmodel
- Consolidate `get_out_bias`/`set_out_bias` into `base_atomic_model.py`
- Add `get_observed_type_list` to abstract API and implement in dpmodel,
pt, pd
- Move fitting input stat update to `model_change_out_bias` in pt/pd
training code (#5266)

  ### Bug fixes
- `general_fitting.change_type_map`: use array-api-compat ops instead of
bare numpy (breaks pt_expt)
- `make_model.change_type_map`: properly forward
`model_with_new_type_stat` to atomic model
  - `stat.py`: fix in-place mutation issue

  ### Tests
- New cross-backend consistency tests: `test_dipole.py`,
`test_polar.py`, `test_dos.py`, `test_property.py` and
`test_zbl_ener.py` (~1400 lines each)
  - Expanded `test_ener.py` with pt_expt and full model API coverage
- New pt_expt unit tests: `test_dipole_model.py`, `test_polar_model.py`,
`test_dos_model.py`, `test_property_model.py`, `test_dp_zbl_model.py`
- Added `test_get_model_def_script`, `test_get_min_nbor_dist`,
`test_set_case_embd` across all 6 model test files
  - Moved atomic model output stat tests from pt_expt to dpmodel
  - Added `model_change_out_bias` tests in pt/pd training tests (#5266)



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* File-backed compute/load for descriptor and fitting statistics; new
compute-or-load stat APIs, get/set output-bias, and observed-type
discovery.
* Exportable/traceable lower-level inference paths for dipole, dos,
polar, property, zbl, and energy models.

* **Refactor**
* Model factory and generated models support extensible base-class
composition and unified fitting backend wiring.

* **Tests**
* Large expansion of cross-backend (DP/PT/PT-EXPT) parity, statistics,
bias, and exportability tests.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

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,
 )
@@ -62,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."""
@@ -87,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
@@ -180,7 +185,7 @@ def forward_atomic(
             nlist,
             mapping=mapping,
         )
-        ret = self.fitting(
+        ret = self.fitting_net(
             descriptor,
             atype,
             gr=rot_mat,
@@ -191,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:
@@ -207,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()
@@ -218,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
@@ -244,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.
@@ -265,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).

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from collections.abc import (
+    Callable,
+)
 from typing import (
     Any,
 )
@@ -17,6 +20,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -338,6 +344,38 @@ def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel":
         data["models"] = models
         return super().deserialize(data)
 
+    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.
+
+        For LinearEnergyAtomicModel, this first computes input stats for each
+        sub-model (without output stats), then computes its own output stats.
+
+        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.
+        """
+        for md in self.models:
+            md.compute_or_load_stat(
+                sampled_func, stat_file_path, compute_or_load_out_stat=False
+            )
+
+        if stat_file_path is not None and self.type_map is not None:
+            stat_file_path /= " ".join(self.type_map)
+
+        if compute_or_load_out_stat:
+            wrapped_sampler = self._make_wrapped_sampler(sampled_func)
+            self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
+
     def _compute_weight(
         self,
         extended_coord: Array,
@@ -523,4 +561,4 @@ def _compute_weight(
         # to handle masked atoms
         coef = xp.where(sigma != 0, coef, xp.zeros_like(coef))
         self.zbl_weight = coef
-        return [1 - xp.expand_dims(coef, -1), xp.expand_dims(coef, -1)]
+        return [1 - xp.expand_dims(coef, axis=-1), xp.expand_dims(coef, axis=-1)]