feat(dpa3): decouple charge_spin from fparam (#5431)

`add_chg_spin_ebd=True` previously hijacked `fparam` to smuggle the
  [charge, spin] scalars into DPA3, forcing users to set `numb_fparam=2`
  on the fitting net and blocking real frame parameters from coexisting
  with charge/spin. This PR plumbs `charge_spin: Tensor | None` as a
  first-class kwarg through every forward chain and adds an optional
  `default_chg_spin` fallback on the DPA3 descriptor.

Backends covered: pt, dpmodel, pt_expt. The pd backend is left
untouched.
  The C/C++/LAMMPS layer is unchanged.

  ## Forward chain

  `Calculator / deep_eval / dp test / lmdb_data / training.get_data`
    -> `wrapper.forward`
    -> `ener_model.forward / forward_lower`
    -> `make_model.forward_common / forward_common_lower`
    -> `base_atomic_model.forward_common_atomic`
-> `dp_atomic_model.forward_atomic` # default_chg_spin fallback here
    -> `descriptor.forward`                  # only DPA3 consumes it

  All other descriptors (se_e2_a, se_r, se_t, se_t_tebd, dpa1, dpa2,
  hybrid) only forward the kwarg through their signatures.

  ## New API surface

  On `BaseAtomicModel` and the wrapped model:

  - `has_chg_spin_ebd() -> bool`
  - `get_dim_chg_spin() -> int`               # 2 for DPA3, else 0
  - `has_default_chg_spin() -> bool`
  - `get_default_chg_spin() -> list[float] | Tensor | None`

  DPA3 descriptor gains a `default_chg_spin: list[float] | None = None`
constructor arg (length 2, validated; round-trips through `serialize`).
  `descrpt_dpa3_args` exposes the matching `Argument` and the
  `add_chg_spin_ebd` doc no longer references fparam.

  ## Training data

  `charge_spin` is registered as a `DataRequirementItem(ndof=2,
  atomic=False, must=not has_default_cs, default=cs_default)`. The
  `get_data` path drops it (along with fparam) on frames where
  `find_charge_spin == 0`, so missing per-frame data falls back to
  `default_chg_spin` when one is configured.

  ## pt_expt specifics

  `forward_common_atomic`, `forward_common_lower_exportable`, the
  `make_fx`-traced inner `fn`, `_trace_and_compile`, and all wrapping
  energy/spin/dipole/dos/polar/property/dp_linear/dp_zbl model variants
  gained a `charge_spin` arg in lockstep so the export and inductor-
  compiled paths keep matching signatures. `deep_eval` no longer reuses
  `fparam` for charge/spin — it constructs `charge_spin_t` (with the
  metadata default-fallback) and passes it explicitly.

  ## Tests

  Three `cs_mode` cases are exercised everywhere it matters:
  `no_chg_spin`, `explicit_chg_spin`, `default_chg_spin`.

  - pt UT (`source/tests/pt/model/test_dpa3.py::test_consistency`)
    rewritten over the three modes; default mode also asserts that the
    default-fallback descriptor matches an explicit `[5,1]` peer.
  - pt_expt UT (`source/tests/pt_expt/descriptor/test_dpa3.py`) gains
    `test_consistency_chg_spin` covering explicit and default modes
    against dpmodel.
  - Universal tests: `DescriptorParamDPA3` learns `default_chg_spin`,
    parametrize gains `(None, [5.0, 1.0])`, and the
    `add_chg_spin_ebd` skip rule in `test_model.py` is replaced —
    the universal driver does not feed `charge_spin`, so chg_spin runs
    rely on the `default_chg_spin` fallback. 622 DPA3 model cases pass.
  - Consistent tests: `descriptor/common.py` threads `charge_spin`
    through every `eval_*` (pd ignores it). `test_dpa3.py` swaps
    `self.fparam` for `self.charge_spin`. `test_ener.py::
    TestEnerChgSpinEbdFparam` is reparametrized over the three modes
    and no longer touches `numb_fparam` / `default_fparam`.

  ## Smoke

  `examples/water/dpa3 dp --pt train input_torch_dynamic.json
  --skip-neighbor-stat` runs to batch 600 with monotonically decreasing
  loss.

  ## Test plan

  - [x] pytest source/tests/pt/model/test_dpa3.py -v
  - [x] pytest source/tests/pt_expt/descriptor/test_dpa3.py -v
  - [x] pytest source/tests/consistent/descriptor/test_dpa3.py -v
- [x] pytest
source/tests/consistent/model/test_ener.py::TestEnerChgSpinEbdFparam -v
- [x] pytest source/tests/universal/dpmodel/model/test_model.py -k "DPA3
and 5"
- [x] examples/water/dpa3 smoke training (600 batches, loss decreasing)

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

* **New Features**
* Optional per-frame charge-spin input supported end-to-end: data
readers, batching, inference, training, export/tracing;
prediction/training calls accept and forward it.
* Models/descriptors expose capability-query and default-value helpers
for charge-spin embeddings; exportable/traced APIs honor defaults.

* **Tests**
* Tests updated/expanded to validate charge-spin embedding behavior and
cross-backend consistency.

* **Chores**
* Configuration normalization warns on legacy charge/spin packed into
legacy parameters and documents migration.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>
Co-authored-by: Anyang Peng <137014849+anyangml@users.noreply.github.com>

Author: iProzd

deepmd/calculator.py

@@ -137,8 +137,14 @@ def calculate(
 
         fparam = self.atoms.info.get("fparam", None)
         aparam = self.atoms.info.get("aparam", None)
+        charge_spin = self.atoms.info.get("charge_spin", None)
         e, f, v = self.dp.eval(
-            coords=coord, cells=cell, atom_types=atype, fparam=fparam, aparam=aparam
+            coords=coord,
+            cells=cell,
+            atom_types=atype,
+            fparam=fparam,
+            aparam=aparam,
+            charge_spin=charge_spin,
         )[:3]
         self.results["energy"] = e[0][0]
         # see https://gitlab.com/ase/ase/-/merge_requests/2485

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -156,6 +156,22 @@ def get_default_fparam(self) -> list[float] | None:
         """Get the default frame parameters."""
         return None
 
+    def has_chg_spin_ebd(self) -> bool:
+        """Check if the model has charge spin embedding."""
+        return False
+
+    def get_dim_chg_spin(self) -> int:
+        """Get the dimension of charge_spin input."""
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Check if the model has default charge_spin values."""
+        return False
+
+    def get_default_chg_spin(self) -> list[float] | None:
+        """Get the default charge_spin values."""
+        return None
+
     def reinit_atom_exclude(
         self,
         exclude_types: list[int] = [],
@@ -232,6 +248,7 @@ def forward_common_atomic(
         fparam: Array | None = None,
         aparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         """Common interface for atomic inference.
 
@@ -284,6 +301,7 @@ def forward_common_atomic(
             fparam=fparam,
             aparam=aparam,
             comm_dict=comm_dict,
+            charge_spin=charge_spin,
         )
         ret_dict = self.apply_out_stat(ret_dict, atype)
 
@@ -312,6 +330,7 @@ def call(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         return self.forward_common_atomic(
             extended_coord,
@@ -320,6 +339,7 @@ def call(
             mapping=mapping,
             fparam=fparam,
             aparam=aparam,
+            charge_spin=charge_spin,
         )
 
     def get_intensive(self) -> bool:
@@ -524,6 +544,7 @@ def model_forward(
             box: np.ndarray | None,
             fparam: np.ndarray | None = None,
             aparam: np.ndarray | None = None,
+            charge_spin: 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
@@ -543,6 +564,8 @@ def model_forward(
                 fparam = xp.asarray(fparam, device=device)
             if aparam is not None:
                 aparam = xp.asarray(aparam, device=device)
+            if charge_spin is not None:
+                charge_spin = xp.asarray(charge_spin, device=device)
 
             (
                 extended_coord,
@@ -564,6 +587,7 @@ def model_forward(
                 mapping=mapping,
                 fparam=fparam,
                 aparam=aparam,
+                charge_spin=charge_spin,
             )
             # Convert outputs back to numpy arrays
             return {kk: to_numpy_array(vv) for kk, vv in atomic_ret.items()}

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -79,6 +79,28 @@ def __init__(
         )
         super().init_out_stat()
 
+    def has_chg_spin_ebd(self) -> bool:
+        """Check if the model has charge spin embedding."""
+        return self.add_chg_spin_ebd
+
+    def get_dim_chg_spin(self) -> int:
+        """Get the dimension of charge_spin input."""
+        if self.add_chg_spin_ebd:
+            return self.descriptor.get_dim_chg_spin()
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Check if the model has default charge_spin values."""
+        if self.add_chg_spin_ebd:
+            return self.descriptor.has_default_chg_spin()
+        return False
+
+    def get_default_chg_spin(self) -> list[float] | None:
+        """Get the default charge_spin values."""
+        if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin():
+            return self.descriptor.get_default_chg_spin()
+        return None
+
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
         return self.fitting_net.output_def()
@@ -158,6 +180,7 @@ def forward_atomic(
         fparam: Array | None = None,
         aparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         """Models' atomic predictions.
 
@@ -178,6 +201,8 @@ def forward_atomic(
         comm_dict
             MPI communication metadata for parallel inference. ``None`` for
             non-parallel inference (default). Forwarded to the descriptor.
+        charge_spin
+            charge and spin parameter for descriptor. nf x 2
 
         Returns
         -------
@@ -188,38 +213,29 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = xp_take_first_n(extended_atype, 1, nloc)
 
-        # Handle default fparam if fitting net supports it
-        if (
-            hasattr(self.fitting_net, "get_dim_fparam")
-            and self.fitting_net.get_dim_fparam() > 0
-            and fparam is None
-        ):
-            # use default fparam
-            from deepmd.dpmodel.array_api import (
-                array_api_compat,
-            )
-
-            default_fparam = self.fitting_net.get_default_fparam()
-            assert default_fparam is not None
-            xp = array_api_compat.array_namespace(extended_coord)
-            default_fparam_array = xp.asarray(
-                default_fparam,
-                dtype=extended_coord.dtype,
-                device=array_api_compat.device(extended_coord),
-            )
-            fparam_input_for_des = xp.tile(
-                xp.reshape(default_fparam_array, (1, -1)), (nframes, 1)
-            )
-        else:
-            fparam_input_for_des = fparam
+        # Handle default charge_spin if descriptor supports it
+        if self.add_chg_spin_ebd and charge_spin is None:
+            default_cs = self.descriptor.get_default_chg_spin()
+            if default_cs is not None:
+                from deepmd.dpmodel.array_api import (
+                    array_api_compat,
+                )
+
+                xp = array_api_compat.array_namespace(extended_coord)
+                cs_array = xp.asarray(
+                    default_cs,
+                    dtype=extended_coord.dtype,
+                    device=array_api_compat.device(extended_coord),
+                )
+                charge_spin = xp.tile(xp.reshape(cs_array, (1, -1)), (nframes, 1))
 
         descriptor, rot_mat, g2, h2, sw = self.descriptor(
             extended_coord,
             extended_atype,
             nlist,
             mapping=mapping,
-            fparam=fparam_input_for_des if self.add_chg_spin_ebd else None,
             comm_dict=comm_dict,
+            charge_spin=charge_spin if self.add_chg_spin_ebd else None,
         )
         ret = self.fitting_net(
             descriptor,

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -225,6 +225,7 @@ def forward_atomic(
         fparam: Array | None = None,
         aparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         """Return atomic prediction.
 
@@ -286,6 +287,7 @@ def forward_atomic(
                     fparam,
                     aparam,
                     comm_dict,
+                    charge_spin=charge_spin,
                 )["energy"]
             )
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -138,6 +138,7 @@ def fwd(
             mapping: t_tensor | None = None,
             fparam: t_tensor | None = None,
             aparam: t_tensor | None = None,
+            charge_spin: t_tensor | None = None,
         ) -> dict[str, t_tensor]:
             pass
 

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -254,6 +254,7 @@ def forward_atomic(
         fparam: Array | None = None,
         aparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         del comm_dict  # pairtab is local; no MPI ghost exchange needed.
         xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist)

deepmd/dpmodel/descriptor/dpa1.py

@@ -509,6 +509,7 @@ def call(
         mapping: Array | None = None,
         fparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> Array:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/dpa2.py

@@ -842,6 +842,7 @@ def call(
         mapping: Array | None = None,
         fparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/dpa3.py

@@ -377,6 +377,7 @@ def __init__(
         use_loc_mapping: bool = True,
         type_map: list[str] | None = None,
         add_chg_spin_ebd: bool = False,
+        default_chg_spin: list[float] | None = None,
     ) -> None:
         super().__init__()
 
@@ -433,6 +434,11 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
 
         self.use_econf_tebd = use_econf_tebd
         self.add_chg_spin_ebd = add_chg_spin_ebd
+        if default_chg_spin is not None and len(default_chg_spin) != 2:
+            raise ValueError(
+                "default_chg_spin must have exactly 2 values [charge, spin]"
+            )
+        self.default_chg_spin = default_chg_spin
         self.use_tebd_bias = use_tebd_bias
         self.use_loc_mapping = use_loc_mapping
         self.type_map = type_map
@@ -499,6 +505,18 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_dim_chg_spin(self) -> int:
+        """Returns the dimension of charge_spin input."""
+        return 2 if self.add_chg_spin_ebd else 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Returns whether default charge_spin values are set."""
+        return self.default_chg_spin is not None
+
+    def get_default_chg_spin(self) -> list[float] | None:
+        """Returns the default charge_spin values."""
+        return self.default_chg_spin
+
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.rcut_smth
@@ -647,6 +665,7 @@ def call(
         mapping: Array | None = None,
         fparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 
@@ -702,13 +721,13 @@ def call(
             )
 
         if self.add_chg_spin_ebd:
-            assert fparam is not None
+            assert charge_spin is not None
             assert self.chg_embedding is not None
             assert self.spin_embedding is not None
             chg_tebd = self.chg_embedding.call()
             spin_tebd = self.spin_embedding.call()
-            charge = xp.astype(fparam[:, 0], xp.int64) + 100
-            spin = xp.astype(fparam[:, 1], xp.int64)
+            charge = xp.astype(charge_spin[:, 0], xp.int64) + 100
+            spin = xp.astype(charge_spin[:, 1], xp.int64)
             chg_ebd = xp.reshape(
                 xp.take(chg_tebd, xp.reshape(charge, (-1,)), axis=0),
                 (nframes, self.tebd_dim),
@@ -753,6 +772,7 @@ def serialize(self) -> dict:
             "use_tebd_bias": self.use_tebd_bias,
             "use_loc_mapping": self.use_loc_mapping,
             "add_chg_spin_ebd": self.add_chg_spin_ebd,
+            "default_chg_spin": self.default_chg_spin,
             "type_map": self.type_map,
             "type_embedding": self.type_embedding.serialize(),
         }

deepmd/dpmodel/descriptor/hybrid.py

@@ -123,6 +123,40 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return np.max([descrpt.get_rcut() for descrpt in self.descrpt_list]).item()
 
+    def get_dim_chg_spin(self) -> int:
+        """Returns the dimension of charge_spin input (0 if not supported)."""
+        return max(
+            (descrpt.get_dim_chg_spin() for descrpt in self.descrpt_list), default=0
+        )
+
+    def has_default_chg_spin(self) -> bool:
+        """Returns whether the descriptor has a default charge_spin value."""
+        default_chg_spin = None
+        found_chg_spin = False
+        for descrpt in self.descrpt_list:
+            if descrpt.get_dim_chg_spin() == 0:
+                continue
+            found_chg_spin = True
+            if not descrpt.has_default_chg_spin():
+                return False
+            child_default_chg_spin = descrpt.get_default_chg_spin()
+            if child_default_chg_spin is None:
+                return False
+            if default_chg_spin is None:
+                default_chg_spin = child_default_chg_spin
+            elif child_default_chg_spin != default_chg_spin:
+                return False
+        return found_chg_spin
+
+    def get_default_chg_spin(self) -> list[float] | None:
+        """Returns the default charge_spin value, or None."""
+        if not self.has_default_chg_spin():
+            return None
+        for descrpt in self.descrpt_list:
+            if descrpt.get_dim_chg_spin() > 0:
+                return descrpt.get_default_chg_spin()
+        return None
+
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         # may not be a good idea...
@@ -287,6 +321,7 @@ def call(
         mapping: Array | None = None,
         fparam: Array | None = None,
         comm_dict: dict | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[
         Array,
         Array | None,
@@ -344,7 +379,13 @@ def call(
                 assert nl_distinguish_types is not None
                 nl = nl_distinguish_types[:, :, nci]
             odescriptor, gr, _g2, _h2, _sw = descrpt(
-                coord_ext, atype_ext, nl, mapping, comm_dict=comm_dict
+                coord_ext,
+                atype_ext,
+                nl,
+                mapping,
+                fparam=fparam,
+                comm_dict=comm_dict,
+                charge_spin=charge_spin,
             )
             out_descriptor.append(odescriptor)
             if gr is not None: