feat(pt_expt): add training infrastructure

deepmd/backend/pt_expt.py

@@ -32,7 +32,7 @@
 class PyTorchExportableBackend(Backend):
     """PyTorch exportable backend."""
 
-    name = "PyTorch Exportable"
+    name = "PyTorch-Exportable"
     """The formal name of the backend."""
     features: ClassVar[Backend.Feature] = (
         Backend.Feature.ENTRY_POINT
@@ -63,7 +63,7 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]:
         Callable[[Namespace], None]
             The entry point hook of the backend.
         """
-        from deepmd.pt.entrypoints.main import main as deepmd_main
+        from deepmd.pt_expt.entrypoints.main import main as deepmd_main
 
         return deepmd_main
 

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -332,19 +332,19 @@ def wrapped_sampler() -> list[dict]:
                 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()
-            ):
+            # For systems where fparam is missing (find_fparam == 0),
+            # fill with default fparam if available and mark as found.
+            if 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)
-                        )
+                        if "find_fparam" in sample and not sample["find_fparam"]:
+                            nframe = sample["atype"].shape[0]
+                            sample["fparam"] = np.tile(
+                                default_fparam_np.reshape(1, -1), (nframe, 1)
+                            )
+                            sample["find_fparam"] = np.bool_(True)
             return sampled
 
         return wrapped_sampler

deepmd/dpmodel/descriptor/dpa1.py

@@ -242,6 +242,8 @@ class DescrptDPA1(NativeOP, BaseDescriptor):
        arXiv preprint arXiv:2208.08236.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -662,7 +664,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, sel = UpdateSel().update_one_sel(
+        min_nbor_dist, sel = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
         )
         local_jdata_cpy["sel"] = sel[0]

deepmd/dpmodel/descriptor/dpa2.py

@@ -441,6 +441,8 @@ class DescrptDPA2(NativeOP, BaseDescriptor):
        Comput Mater 10, 293 (2024). https://doi.org/10.1038/s41524-024-01493-2
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         ntypes: int,
@@ -1114,7 +1116,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        update_sel = UpdateSel()
+        update_sel = cls._update_sel_cls()
         min_nbor_dist, repinit_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/dpmodel/descriptor/dpa3.py

@@ -337,6 +337,8 @@ class DescrptDPA3(NativeOP, BaseDescriptor):
        arXiv preprint arXiv:2506.01686 (2025).
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         ntypes: int,
@@ -729,7 +731,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        update_sel = UpdateSel()
+        update_sel = cls._update_sel_cls()
         min_nbor_dist, repflow_e_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -51,7 +51,7 @@ class BD(ABC, PluginVariant, make_plugin_registry("descriptor")):
         def __new__(cls, *args: Any, **kwargs: Any) -> Any:
             if cls is BD:
                 cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))
-            return super().__new__(cls)
+            return object.__new__(cls)
 
         @abstractmethod
         def get_rcut(self) -> float:

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -149,6 +149,8 @@ class DescrptSeA(NativeOP, BaseDescriptor):
        Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -582,7 +584,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+        min_nbor_dist, local_jdata_cpy["sel"] = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
         )
         return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/descriptor/se_r.py

@@ -128,6 +128,8 @@ class DescrptSeR(NativeOP, BaseDescriptor):
        Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -505,7 +507,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+        min_nbor_dist, local_jdata_cpy["sel"] = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
         )
         return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/descriptor/se_t.py

@@ -116,6 +116,8 @@ class DescrptSeT(NativeOP, BaseDescriptor):
             Not used in this descriptor, only to be compat with input.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -505,7 +507,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+        min_nbor_dist, local_jdata_cpy["sel"] = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
         )
         return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -141,6 +141,8 @@ class DescrptSeTTebd(NativeOP, BaseDescriptor):
 
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -500,7 +502,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, sel = UpdateSel().update_one_sel(
+        min_nbor_dist, sel = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
         )
         local_jdata_cpy["sel"] = sel[0]

deepmd/dpmodel/fitting/general_fitting.py

@@ -266,6 +266,17 @@ def compute_input_stats(
                 )
             else:
                 sampled = merged() if callable(merged) else merged
+                for ii, frame in enumerate(sampled):
+                    if "find_fparam" not in frame:
+                        raise ValueError(
+                            f"numb_fparam > 0 but fparam is not acquired "
+                            f"for system {ii}."
+                        )
+                    if not frame["find_fparam"]:
+                        raise ValueError(
+                            f"numb_fparam > 0 but no fparam data is provided "
+                            f"for system {ii}."
+                        )
                 cat_data = np.concatenate(
                     [frame["fparam"] for frame in sampled], axis=0
                 )
@@ -313,6 +324,17 @@ def compute_input_stats(
                 )
             else:
                 sampled = merged() if callable(merged) else merged
+                for ii, frame in enumerate(sampled):
+                    if "find_aparam" not in frame:
+                        raise ValueError(
+                            f"numb_aparam > 0 but aparam is not acquired "
+                            f"for system {ii}."
+                        )
+                    if not frame["find_aparam"]:
+                        raise ValueError(
+                            f"numb_aparam > 0 but no aparam data is provided "
+                            f"for system {ii}."
+                        )
                 sys_sumv = []
                 sys_sumv2 = []
                 sys_sumn = []

deepmd/dpmodel/fitting/make_base_fitting.py

@@ -42,7 +42,7 @@ class BF(ABC, PluginVariant, make_plugin_registry("fitting")):
         def __new__(cls: type, *args: Any, **kwargs: Any) -> Any:
             if cls is BF:
                 cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))
-            return super().__new__(cls)
+            return object.__new__(cls)
 
         @abstractmethod
         def output_def(self) -> FittingOutputDef:

deepmd/dpmodel/loss/ener.py

@@ -95,10 +95,10 @@ def call(
         label_dict: dict[str, Array],
     ) -> dict[str, Array]:
         """Calculate loss from model results and labeled results."""
-        energy = model_dict["energy_redu"]
-        force = model_dict["energy_derv_r"]
-        virial = model_dict["energy_derv_c_redu"]
-        atom_ener = model_dict["energy"]
+        energy = model_dict["energy"]
+        force = model_dict["force"]
+        virial = model_dict["virial"]
+        atom_ener = model_dict["atom_energy"]
         energy_hat = label_dict["energy"]
         force_hat = label_dict["force"]
         virial_hat = label_dict["virial"]
@@ -212,7 +212,7 @@ def call(
                 )
                 loss += pref_v * l_huber_loss
             more_loss["rmse_v"] = self.display_if_exist(
-                xp.sqrt(l2_virial_loss), find_virial
+                xp.sqrt(l2_virial_loss) * atom_norm, find_virial
             )
         if self.has_ae:
             atom_ener_reshape = xp.reshape(atom_ener, (-1,))

deepmd/dpmodel/loss/loss.py

@@ -53,8 +53,11 @@ def display_if_exist(loss: Array, find_property: float) -> Array:
             the loss scalar or NaN
         """
         xp = array_api_compat.array_namespace(loss)
+        dev = array_api_compat.device(loss)
         return xp.where(
-            xp.asarray(find_property, dtype=xp.bool), loss, xp.asarray(xp.nan)
+            xp.asarray(find_property, dtype=xp.bool, device=dev),
+            loss,
+            xp.asarray(xp.nan, device=dev),
         )
 
     @classmethod

deepmd/dpmodel/model/base_model.py

@@ -42,7 +42,7 @@ def __new__(cls, *args: Any, **kwargs: Any) -> "BaseModel":
                 if model_type == "standard":
                     model_type = kwargs.get("fitting", {}).get("type", "ener")
                 cls = cls.get_class_by_type(model_type)
-            return super().__new__(cls)
+            return object.__new__(cls)
 
         @abstractmethod
         def __call__(self, *args: Any, **kwds: Any) -> Any:

deepmd/dpmodel/utils/batch.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Normalize raw batches from DeepmdDataSystem into canonical format."""
+
+from typing import (
+    Any,
+)
+
+import numpy as np
+
+# Keys that are metadata / not needed by models or loss functions.
+_DROP_KEYS = {"default_mesh", "sid", "fid"}
+
+# Keys that belong to model input (everything else is label).
+_INPUT_KEYS = {"coord", "atype", "spin", "box", "fparam", "aparam"}
+
+
+def normalize_batch(batch: dict[str, Any]) -> dict[str, Any]:
+    """Normalize a raw batch from :class:`DeepmdDataSystem` to canonical format.
+
+    The following conversions are applied:
+
+    * ``"type"`` is renamed to ``"atype"`` (int64).
+    * ``"natoms_vec"`` (1-D) is tiled to 2-D ``[nframes, 2+ntypes]``
+      and stored as ``"natoms"``.
+    * ``find_*`` flags are converted to ``np.bool_``.
+    * Metadata keys (``default_mesh``, ``sid``, ``fid``) are dropped.
+
+    Parameters
+    ----------
+    batch : dict[str, Any]
+        Raw batch dict returned by ``DeepmdDataSystem.get_batch()``.
+
+    Returns
+    -------
+    dict[str, Any]
+        Normalized batch dict (new dict; the input is not mutated).
+    """
+    out: dict[str, Any] = {}
+
+    for key, val in batch.items():
+        if key in _DROP_KEYS:
+            continue
+
+        if key == "type":
+            out["atype"] = val.astype(np.int64)
+        elif key.startswith("find_"):
+            out[key] = np.bool_(float(val) > 0.5)
+        elif key == "natoms_vec":
+            nv = val
+            if nv.ndim == 1 and "coord" in batch:
+                nframes = batch["coord"].shape[0]
+                nv = np.tile(nv, (nframes, 1))
+            out["natoms"] = nv
+        else:
+            out[key] = val
+
+    return out
+
+
+def split_batch(
+    batch: dict[str, Any],
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Split a normalized batch into input and label dicts.
+
+    Parameters
+    ----------
+    batch : dict[str, Any]
+        Normalized batch (output of :func:`normalize_batch`).
+
+    Returns
+    -------
+    input_dict : dict[str, Any]
+        Model inputs (coord, atype, box, fparam, aparam, spin).
+    label_dict : dict[str, Any]
+        Labels and find flags (energy, force, virial, find_*, natoms, …).
+    """
+    input_dict: dict[str, Any] = {}
+    label_dict: dict[str, Any] = {}
+
+    for key, val in batch.items():
+        if key in _INPUT_KEYS:
+            input_dict[key] = val
+        else:
+            label_dict[key] = val
+
+    return input_dict, label_dict

deepmd/dpmodel/utils/env_mat.py

@@ -26,7 +26,10 @@ def compute_smooth_weight(
     if rmin >= rmax:
         raise ValueError("rmin should be less than rmax.")
     xp = array_api_compat.array_namespace(distance)
-    distance = xp.clip(distance, min=rmin, max=rmax)
+    # Use where instead of clip so that make_fx tracing does not
+    # decompose it into boolean-indexed ops with data-dependent sizes.
+    distance = xp.where(distance < rmin, xp.full_like(distance, rmin), distance)
+    distance = xp.where(distance > rmax, xp.full_like(distance, rmax), distance)
     uu = (distance - rmin) / (rmax - rmin)
     uu2 = uu * uu
     vv = uu2 * uu * (-6.0 * uu2 + 15.0 * uu - 10.0) + 1.0
@@ -42,7 +45,8 @@ def compute_exp_sw(
     if rmin >= rmax:
         raise ValueError("rmin should be less than rmax.")
     xp = array_api_compat.array_namespace(distance)
-    distance = xp.clip(distance, min=0.0, max=rmax)
+    distance = xp.where(distance < 0.0, xp.zeros_like(distance), distance)
+    distance = xp.where(distance > rmax, xp.full_like(distance, rmax), distance)
     C = 20
     a = C / rmin
     b = rmin