[pull] master from deepmodeling:master

deepmd/dpmodel/atomic_model/inter_potential.py → deepmd/dpmodel/atomic_model/inner_potential.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Analytical pair potentials for Zone bridging (backend-agnostic port of
-``deepmd.pt``'s ``InterPotential``). Lives in the atomic-model package:
+``deepmd.pt``'s ``InnerPotential``). Lives in the atomic-model package:
 the atomic layer owns per-atom energy assembly, where the ZBL term is
 injected on the graph route.
 """
@@ -63,7 +63,7 @@
 _A_BOHR = 0.5291772109  # Bohr radius in Å
 
 
-class InterPotential(NativeOP):
+class InnerPotential(NativeOP):
     """Analytical pair potential for Zone bridging.
 
     Supports the Ziegler-Biersack-Littmark (ZBL) screened nuclear repulsion
@@ -72,7 +72,7 @@ class InterPotential(NativeOP):
     contributes ``V_ZBL(r_ij) / 2`` to both atom i and atom j, avoiding
     double-counting from the symmetric neighbor list. Backend-agnostic
     (array-API) port of the reference implementation in
-    ``deepmd.pt.model.model.sezm_model.InterPotential``.
+    ``deepmd.pt.model.model.sezm_model.InnerPotential``.
 
     Parameters
     ----------
@@ -93,7 +93,7 @@ def __init__(self, type_map: list[str], mode: str = "zbl") -> None:
         super().__init__()
         mode = str(mode).upper()
         if mode != "ZBL":
-            raise ValueError(f"Unknown InterPotential mode: {mode}")
+            raise ValueError(f"Unknown InnerPotential mode: {mode}")
         self.mode = mode
         self.type_map = list(type_map)
         self.ntypes_real = len(type_map)
@@ -266,8 +266,8 @@ def call(
         return xp.astype(xp.reshape(atom_energy, (1, n_node, 1)), edge_vec.dtype)
 
 
-@BaseAtomicModel.register("inter_potential")
-class InterPotentialAtomicModel(BaseAtomicModel):
+@BaseAtomicModel.register("inner_potential")
+class InnerPotentialAtomicModel(BaseAtomicModel):
     """Analytical bridging pair potential as an ATOMIC MODEL.
 
     First-principles composition design: the analytical term maps local
@@ -302,7 +302,7 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         super().__init__(type_map, **kwargs)
-        self.potential = InterPotential(type_map=list(type_map), mode=mode)
+        self.potential = InnerPotential(type_map=list(type_map), mode=mode)
         self.mode = self.potential.mode
         self.rcut = float(rcut)
         self.sel = (
@@ -317,7 +317,7 @@ def change_type_map(
         If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
 
         The generic base handles the public map and the stat/exclusion state;
-        the element lookup belongs to :class:`InterPotential`, so the update is
+        the element lookup belongs to :class:`InnerPotential`, so the update is
         delegated there rather than reimplemented here (review 3649295675 --
         without it the lookup keeps the ORIGINAL elements while ``atype``
         values mean new ones, and a longer new map raises ``IndexError``).
@@ -403,7 +403,7 @@ def forward_atomic(
     ) -> dict:
         """Dense route unsupported: the term rides the NeighborGraph route only."""
         raise NotImplementedError(
-            "InterPotentialAtomicModel rides the NeighborGraph route only; "
+            "InnerPotentialAtomicModel rides the NeighborGraph route only; "
             "the dense (nlist) route has no injection site for the term"
         )
 
@@ -455,7 +455,7 @@ def serialize(self) -> dict:
         data.update(
             {
                 "@class": "Model",
-                "type": "inter_potential",
+                "type": "inner_potential",
                 "@version": 1,
                 "mode": self.mode,
                 "rcut": self.rcut,
@@ -465,7 +465,7 @@ def serialize(self) -> dict:
         return data
 
     @classmethod
-    def deserialize(cls, data: dict) -> "InterPotentialAtomicModel":
+    def deserialize(cls, data: dict) -> "InnerPotentialAtomicModel":
         data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class", None)

deepmd/dpmodel/model/dp_linear_model.py

@@ -35,7 +35,7 @@ class LinearEnergyModel(DPModelCommon, DPLinearModel_):
     energies; on the NeighborGraph route every child consumes the same
     graph, so the summed energy differentiates through one shared edge
     backward. Used e.g. for analytical bridging compositions
-    (learned model + :class:`~deepmd.dpmodel.atomic_model.inter_potential.InterPotentialAtomicModel`).
+    (learned model + :class:`~deepmd.dpmodel.atomic_model.inner_potential.InnerPotentialAtomicModel`).
     """
 
     def __init__(

deepmd/dpmodel/model/model.py

@@ -61,7 +61,7 @@ def get_standard_model(data: dict) -> BaseModel:
     data = copy.deepcopy(data)
     # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's
     # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the
-    # atomic model's InterPotential below.
+    # atomic model's InnerPotential below.
     bridging_method = str(data.get("bridging_method", "none"))
     bridging_enabled = bridging_method.lower() not in ("none", "")
     if bridging_enabled:
@@ -77,8 +77,8 @@ def get_standard_model(data: dict) -> BaseModel:
     # Composition, not a flag (first-principles design): the analytical
     # bridging term is its own atomic model, summed with the learned one by the
     # existing linear composition machinery.
-    from deepmd.dpmodel.atomic_model.inter_potential import (
-        InterPotentialAtomicModel,
+    from deepmd.dpmodel.atomic_model.inner_potential import (
+        InnerPotentialAtomicModel,
     )
     from deepmd.dpmodel.atomic_model.linear_atomic_model import (
         LinearEnergyAtomicModel,
@@ -87,7 +87,7 @@ def get_standard_model(data: dict) -> BaseModel:
         LinearEnergyModel,
     )
 
-    zbl_atomic = InterPotentialAtomicModel(
+    zbl_atomic = InnerPotentialAtomicModel(
         type_map=data["type_map"],
         mode=bridging_method,
         rcut=descriptor.get_rcut(),
@@ -140,7 +140,7 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     exclusions and the analytical-bridging composition -- so ``spin`` and
     ``bridging_method`` combine for free: the wrapper re-classes whatever
     atomic model came back, be it a single learned model or a
-    ``LinearEnergyAtomicModel`` over ``[learned, InterPotential]`` (the
+    ``LinearEnergyAtomicModel`` over ``[learned, InnerPotential]`` (the
     analytical child accepts and ignores ``spin``; the learned child consumes
     it).
 

deepmd/dpmodel/model/spin_model.py

@@ -83,6 +83,20 @@ def _to_xp(self, arr: Any, xp: Any, ref_arr: Any) -> Any:
         """Convert a numpy array to the same namespace as ref_arr."""
         return xp.asarray(arr, device=array_api_compat.device(ref_arr))
 
+    def _lookup_type_values(self, values: Any, atype: Array, ref_arr: Array) -> Array:
+        """Gather per-type values while mapping virtual atom types to zero.
+
+        Negative atom types are padding placeholders, not Python-style indices
+        from the end of the type table. Their spin scale and mask must remain
+        zero until the backbone model applies its normal virtual-atom mask.
+        """
+        xp = array_api_compat.array_namespace(ref_arr)
+        values = self._to_xp(values, xp, ref_arr)
+        real_atom = atype >= 0
+        safe_atype = xp.where(real_atom, atype, xp.zeros_like(atype))
+        gathered = values[safe_atype]
+        return xp.where(real_atom, gathered, xp.zeros_like(gathered))
+
     def process_spin_input(
         self, coord: Array, atype: Array, spin: Array
     ) -> tuple[Array, Array, Array]:
@@ -99,9 +113,12 @@ def process_spin_input(
         """
         xp = array_api_compat.array_namespace(coord)
         nframes, nloc = coord.shape[:-1]
-        atype_spin = xp.concat([atype, atype + self.ntypes_real], axis=-1)
-        vsm = self._to_xp(self.virtual_scale_mask, xp, coord)
-        spin_dist = spin * xp.reshape(vsm[atype], (nframes, nloc, 1))
+        virtual_atype = xp.where(atype >= 0, atype + self.ntypes_real, atype)
+        atype_spin = xp.concat([atype, virtual_atype], axis=-1)
+        spin_dist = spin * xp.reshape(
+            self._lookup_type_values(self.virtual_scale_mask, atype, coord),
+            (nframes, nloc, 1),
+        )
         virtual_coord = coord + spin_dist
         coord_spin = xp.concat([coord, virtual_coord], axis=-2)
         # for spin virial correction
@@ -153,12 +170,18 @@ def process_spin_input_lower(
         xp = array_api_compat.array_namespace(extended_coord)
         nframes, nall = extended_coord.shape[:2]
         nloc = nlist.shape[1]
-        vsm = self._to_xp(self.virtual_scale_mask, xp, extended_coord)
         extended_spin_dist = extended_spin * xp.reshape(
-            vsm[extended_atype], (nframes, nall, 1)
+            self._lookup_type_values(
+                self.virtual_scale_mask, extended_atype, extended_coord
+            ),
+            (nframes, nall, 1),
         )
         virtual_extended_coord = extended_coord + extended_spin_dist
-        virtual_extended_atype = extended_atype + self.ntypes_real
+        virtual_extended_atype = xp.where(
+            extended_atype >= 0,
+            extended_atype + self.ntypes_real,
+            extended_atype,
+        )
         extended_coord_updated = self.concat_switch_virtual(
             extended_coord, virtual_extended_coord, nloc
         )
@@ -224,9 +247,18 @@ def process_spin_output(
         if virtual_scale:
             mask = self._to_xp(self.virtual_scale_mask, xp, out_tensor)
         else:
-            mask = self._to_xp(self.spin_mask, xp, out_tensor)
-        atomic_mask = xp.reshape(mask[atype], (nframes, nloc, 1))
-        out_real, out_mag = out_tensor[:, :nloc], out_tensor[:, nloc:]
+            # spin_mask is integral; it multiplies out_mag below, and the array
+            # API does not promote across kinds.
+            mask = xp.astype(
+                self._to_xp(self.spin_mask, xp, out_tensor), out_tensor.dtype
+            )
+        atomic_mask = xp.reshape(
+            self._lookup_type_values(mask, atype, out_tensor),
+            (nframes, nloc, 1),
+        )
+        # Trailing ellipsis: the array API does not specify numpy's implicit
+        # expansion of a partial multi-axis index.
+        out_real, out_mag = out_tensor[:, :nloc, ...], out_tensor[:, nloc:, ...]
         if add_mag:
             out_real = out_real + out_mag
         out_mag = xp.reshape(
@@ -250,19 +282,27 @@ def process_spin_output_lower(
         if virtual_scale:
             mask = self._to_xp(self.virtual_scale_mask, xp, extended_out_tensor)
         else:
-            mask = self._to_xp(self.spin_mask, xp, extended_out_tensor)
-        atomic_mask = xp.reshape(mask[extended_atype], (nframes, nall, 1))
+            # spin_mask is integral; it multiplies extended_out_mag below, and
+            # the array API does not promote across kinds.
+            mask = xp.astype(
+                self._to_xp(self.spin_mask, xp, extended_out_tensor),
+                extended_out_tensor.dtype,
+            )
+        atomic_mask = xp.reshape(
+            self._lookup_type_values(mask, extended_atype, extended_out_tensor),
+            (nframes, nall, 1),
+        )
         extended_out_real = xp.concat(
             [
-                extended_out_tensor[:, :nloc],
-                extended_out_tensor[:, nloc + nloc : nloc + nall],
+                extended_out_tensor[:, :nloc, ...],
+                extended_out_tensor[:, nloc + nloc : nloc + nall, ...],
             ],
             axis=1,
         )
         extended_out_mag = xp.concat(
             [
-                extended_out_tensor[:, nloc : nloc + nloc],
-                extended_out_tensor[:, nloc + nall :],
+                extended_out_tensor[:, nloc : nloc + nloc, ...],
+                extended_out_tensor[:, nloc + nall :, ...],
             ],
             axis=1,
         )
@@ -700,8 +740,10 @@ def call_common(
         if "mask_mag" not in model_ret:
             xp = array_api_compat.array_namespace(atype)
             nframes_m, nloc_m = atype.shape[:2]
-            vsm = self._to_xp(self.virtual_scale_mask, xp, atype)
-            atomic_mask = xp.reshape(vsm[atype], (nframes_m, nloc_m, 1))
+            atomic_mask = xp.reshape(
+                self._lookup_type_values(self.virtual_scale_mask, atype, atype),
+                (nframes_m, nloc_m, 1),
+            )
             model_ret["mask_mag"] = atomic_mask > 0.0
         return model_ret
 
@@ -883,8 +925,12 @@ def call_common_lower(
         if "mask_mag" not in model_ret:
             xp = array_api_compat.array_namespace(extended_atype)
             nall = extended_atype.shape[1]
-            vsm = self._to_xp(self.virtual_scale_mask, xp, extended_atype)
-            atomic_mask = xp.reshape(vsm[extended_atype], (nframes, nall, 1))
+            atomic_mask = xp.reshape(
+                self._lookup_type_values(
+                    self.virtual_scale_mask, extended_atype, extended_atype
+                ),
+                (nframes, nall, 1),
+            )
             model_ret["mask_mag"] = atomic_mask > 0.0
         return model_ret
 

deepmd/dpmodel/train/__init__.py

@@ -11,6 +11,10 @@
     AbstractTrainEntrypoint,
     TrainEntrypointOptions,
 )
+from .schedule import (
+    StepSchedule,
+    resolve_step_schedule,
+)
 from .trainer import (
     DEFAULT_TASK_KEY,
     AbstractTrainer,
@@ -30,6 +34,7 @@
     "AbstractTrainer",
     "LearningCurveWriter",
     "RankContext",
+    "StepSchedule",
     "TrainEntrypointOptions",
     "TrainStepResult",
     "TrainerConfig",
@@ -41,4 +46,5 @@
     "iter_training_task_configs",
     "make_task_maps",
     "print_data_summaries",
+    "resolve_step_schedule",
 ]