feat(pt-expt): add compact descriptor DPA4C 🎉🎉🎉 (deepmodeling#5972)

## Summary

This PR introduces DPA4C, the compact and compressible degree-wise
member of
the DPA4 family, as a PyTorch Exportable (`pt_expt`) descriptor. DPA4C
is a
strictly local, one-hop model intended for high-throughput molecular
dynamics:
it reads each directed neighbor edge once, performs one destination
reduction,
and converts the resulting degree-wise moments into a fixed invariant
vector
without cross-atom message passing.

The PR includes the complete path from training to deployment:

- a backend-neutral DPA4C descriptor and a native `pt_expt`
implementation;
- graph-native training, serialization, export, compression, and
calibration;
- fused CUDA descriptor, fitting, force, virial, and magnetic-force
paths;
- native-spin conditioning from the descriptor through Python, C, C++,
and
  LAMMPS/Kokkos interfaces;
- frame-level charge and spin-multiplicity conditioning, including
runtime
  re-specialization of compressed artifacts;
- function-preserving fine-tuning from a spin-free checkpoint;
- ragged mixed-size training batches without exposing phantom atoms to
the
  network; and
- user documentation plus non-spin and native-spin examples.

## Why DPA4C

DPA4/SeZM uses equivariant message passing to target the accuracy
frontier.
DPA4C targets a different operating point: a compact local student whose
radial dependence can be tabulated and whose angular computation can be
fused
into bounded per-edge and per-node CUDA kernels.

The descriptor consumes a carry-all cutoff graph rather than a
fixed-capacity
neighbor list. It therefore has no `sel` parameter, no capacity derived
from
the densest training frame, and no neighbor truncation. Its persistent
per-atom state is determined by `channels` and `lmax`, not by the number
of
neighbors.

## Descriptor architecture

### Edge representation

For every directed edge `j -> i`, DPA4C combines:

- the DPA4 Bessel or Gaussian radial basis;
- a bias-free one-hidden-layer SwiGLU radial network;
- ordered PairFiLM scale and shift terms for `(type_i, type_j)`;
- optional pair-conditioned shared radial modes; and
- a C3 cutoff envelope whose value and first three radial derivatives
join
  continuously to zero at `rcut`.

`radial_modes` increases chemical/radial resolution without widening the
per-atom moment state. The portable implementation accepts any
non-negative
mode count; the compressed CUDA path specializes the production profiles
listed below.

### One-reduction degree-wise moments

The edge direction is expanded in real Cartesian harmonics through
`lmax`.
All scalar masses and all angular moments are packed into one edge
payload and
accumulated with one destination segment reduction. Two smooth
neighborhood
masses normalize the scalar and non-scalar blocks and are also emitted
as
descriptor coordinates so the fitting network retains effective
coordination
information.

The channel schedule keeps degree 0 wide, retains several channels for
degrees
1 and 2, and uses one channel for degrees 3 and 4. This bounds the node
state
while preserving the low-degree angular information that dominates the
model.

### Fixed invariant readout

The node-local readout combines:

- exact aligned Gram matrices within each degree;
- normalized low-rank bispectrum contractions across allowed degree
triples;
- the projected `Qv` quartic; and
- the two neighborhood-mass coordinates.

Only O(3)-even invariant scalars reach the standard energy fitting
network.
Energy is therefore invariant under rotations, reflections, and neighbor
permutations, while force and virial remain conservative derivatives of
the
same total energy.

The public structural controls are:

- `channels` in `{8, 16, 32, 64, 128}`;
- `lmax` in `{2, 3, 4}`;
- `basis_type` in `{bessel, gaussian}`;
- `n_radial`;
- `radial_modes`; and
- `use_amp`, which applies bf16 autocast only to the edge-dominated
stage and
restores descriptor precision before reduction and invariant
contraction.

## Frame charge-state conditioning

When `add_chg_spin_ebd` is enabled, DPA4C accepts one frame-level
`[charge, multiplicity]` condition. This condition is independent of the
per-atom native-spin vector. It enters at two finite locations:

1. a shift of the center type embedding; and
2. a bias of the ordered-pair encoder hidden state.

The portable graph path keeps the condition per frame, so one batch may
contain
different charge states. `default_chg_spin` supplies the fallback state
when an
input does not provide one.

Compression folds a single state into the finite type table and
ordered-pair
caches, leaving the radial table, angular equations, and CUDA kernel
layout
unchanged. The exported artifact carries a charge-state fold that
rebuilds only
the affected constants when the evaluator, C/C++ API, or LAMMPS pair
style
selects another state. This keeps the compact canonical inference ABI
free of a
per-edge runtime condition while avoiding a permanently baked-in charge
state.

## Compression and deployment

Compression tabulates the distance-only radial network with quintic
Hermite
splines on `[0, rcut]` and snapshots the finite ordered-type-pair
tables. The
compiled descriptor supports:

```text
channels     in {8, 16, 32, 64, 128}
lmax         in {2, 3, 4}
radial_modes in {0, 2, 4, 8}
precision    = float32
```

The fused implementation includes forward and backward descriptor
operators,
compact canonical graph operators, fitting-network kernels, and
force/virial
assembly. The backward saves the minimum node moment state and
recomputes the
edge-local radial and angular terms, avoiding a persistent per-edge
moment
tensor. Evaluation is tiled so temporary memory stays bounded for large
edge
sets.

`DP_CUDA_INFER=1` enables the fused descriptor/fitting path with
autograd force
assembly. `DP_CUDA_INFER=2` additionally uses the compact canonical
fused
energy/force/virial composition. The export metadata records the graph
ABI and
dtype contract used by the C++ and LAMMPS loaders.

Graph folding now fails explicitly when a topology requests local-owner
folding
but does not provide a valid owner for every ghost. This prevents a
malformed
standalone C++ call from silently dropping halo-edge contributions.
Extended
multi-rank paths keep ghosts as distinct nodes and use reverse
communication as
their force-folding contract.

## Integration surface

- Registers `descriptor.type: dpa4c` for the PyTorch Exportable backend
and
  documents its arguments in `argcheck`.
- Adds model serialization, graph export, compression routing, inference
  metadata, and evaluation inputs for both charge state and native spin.
- Extends C and C++ energy/spin interfaces with charge-state dimensions,
  setters, and per-call inputs.
- Adds non-spin water and native-spin NiO examples and a full user
guide.
- Adds backend-neutral, PyTorch, CUDA, graph-lower, export, fine-tuning,
symmetry, derivative, serialization, compression, and deployment tests.
- Adapts the DPA1 shared graph-kernel helpers without changing DPA1's
public
  descriptor contract.

The final integration commit also replaces the removed
`doc_only_pt_expt_supported` symbol with the current
`supported_backends("pt_expt")` registry introduced on `master` by
deepmodeling#5929.
This is the only modification made after cherry-picking the four DPA4C
commits.

## Current scope and limitations

- DPA4C is implemented for `pt_expt`; other backends are not added here.
- Compressed inference is float32-only and restricted to the structural
profiles listed above. Unsupported profiles continue to use the portable
  path or are rejected by explicit compression validation.
- Descriptor-level excluded type pairs are not supported by the fused
compact
  kernel.
- Native spin requires `scheme: native`; the virtual-atom `deepspin`
scheme is
  not used by DPA4C.
- The symmetric spin invariant basis does not represent the
antisymmetric
  Dzyaloshinskii-Moriya interaction.
- The provided LAMMPS example covers evaluation and spin minimization.
Spin
  dynamics through stock `fix nve/spin` additionally depends on that fix
  recognizing the new pair style.


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

* **New Features**
* Added the DPA4C descriptor with native-spin, charge-state
conditioning, compressed CUDA inference, and canonical graph support.
* Added native-spin LAMMPS pair styles and expanded C/C++ APIs for spin,
charge-state configuration, and GPU graph inference.
* Added compression capability detection and support for analytically
bounded compression domains.
* **Bug Fixes**
* Improved force, virial, magnetic-force, charge-state, and loss
handling consistency.
* **Documentation**
* Added DPA4C guides, training configurations, and spin-enabled LAMMPS
examples.
* **Tests**
* Expanded coverage for DPA4C, CUDA compression, export, validation,
spin, and charge-state behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: OutisLi

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -119,7 +119,7 @@ def __init__(
         self.fitting_net = fitting
         self.fitting_net.reinit_exclude(self.atom_exclude_types)
         self.type_map = type_map
-        self.add_chg_spin_ebd: bool = self.descriptor.get_dim_chg_spin() > 0
+        self.add_chg_spin_ebd: bool = self.descriptor.has_chg_spin_ebd()
         # Structural capability: only descriptors with a native spin
         # conditioning mechanism (currently DPA4) accept a ``spin`` kwarg on
         # ``call_graph`` at all -- unlike ``charge_spin``, which every
@@ -182,6 +182,10 @@ def supports_graph_export(self) -> bool:
         """Delegates to this model's own descriptor."""
         return bool(self.descriptor.supports_graph_export())
 
+    def compression_needs_min_nbor_dist(self) -> bool:
+        """Delegates to this model's own descriptor."""
+        return bool(self.descriptor.compression_needs_min_nbor_dist())
+
     def supports_native_spin(self) -> bool:
         """Delegates to this model's own descriptor (cached at construction)."""
         return self._supports_native_spin

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -303,6 +303,15 @@ def enable_compression(
                 check_frequency,
             )
 
+    def compression_needs_min_nbor_dist(self) -> bool:
+        """Required as soon as ANY child consumes it.
+
+        The statistic is measured once and handed to every child, so a single
+        child that tabulates from the shortest observed distance keeps the
+        neighbor-statistics pass for the whole composition.
+        """
+        return any(m.compression_needs_min_nbor_dist() for m in self.models)
+
     def uses_graph_lower(self) -> bool:
         """Graph-capable iff EVERY child supports the graph lower.
 

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -190,6 +190,17 @@ def enable_compression(
             """
             raise NotImplementedError("This atomi model doesn't support compression!")
 
+        def compression_needs_min_nbor_dist(self) -> bool:
+            """Whether :meth:`enable_compression` consumes ``min_nbor_dist``.
+
+            Returns
+            -------
+            bool
+                Concrete default ``True``, so a model that does not report
+                otherwise keeps the neighbor-statistics pass.
+            """
+            return True
+
         def make_atom_mask(
             self,
             atype: t_tensor,

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -505,3 +505,14 @@ def enable_compression(
     ) -> None:
         """Pairtab model does not support compression."""
         pass
+
+    def compression_needs_min_nbor_dist(self) -> bool:
+        """Return whether compression consumes the minimum neighbor distance.
+
+        Returns
+        -------
+        bool
+            Always ``False``. The tabulated pair potential carries its own
+            domain, so compression is a no-op here.
+        """
+        return False

deepmd/dpmodel/descriptor/__init__.py

@@ -11,6 +11,9 @@
 from .dpa4 import (
     DescrptDPA4,
 )
+from .dpa4c import (
+    DescrptDPA4C,
+)
 from .hybrid import (
     DescrptHybrid,
 )
@@ -38,6 +41,7 @@
     "DescrptDPA2",
     "DescrptDPA3",
     "DescrptDPA4",
+    "DescrptDPA4C",
     "DescrptHybrid",
     "DescrptSeA",
     "DescrptSeAttenV2",

deepmd/dpmodel/descriptor/dpa3.py

@@ -33,6 +33,12 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.charge_state import (
+    CHARGE_OFFSET,
+    CHARGE_TABLE_ROWS,
+    MULTIPLICITY_TABLE_ROWS,
+    validate_charge_state,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -468,11 +474,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.default_chg_spin = (
+            None
+            if default_chg_spin is None
+            else validate_charge_state(default_chg_spin)
+        )
         self.use_tebd_bias = use_tebd_bias
         self.use_loc_mapping = use_loc_mapping
         self.type_map = type_map
@@ -494,18 +500,16 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
 
         if self.add_chg_spin_ebd:
             self.cs_activation_fn = get_activation_fn(activation_function)
-            # -100 ~ 100 is a conservative bound
             self.chg_embedding = TypeEmbedNet(
-                ntypes=200,
+                ntypes=CHARGE_TABLE_ROWS,
                 neuron=[self.tebd_dim],
                 padding=True,
                 activation_function="Linear",
                 precision=precision,
                 seed=child_seed(seed, 3),
             )
-            # 100 is a conservative upper bound
             self.spin_embedding = TypeEmbedNet(
-                ntypes=100,
+                ntypes=MULTIPLICITY_TABLE_ROWS,
                 neuron=[self.tebd_dim],
                 padding=True,
                 activation_function="Linear",
@@ -543,6 +547,10 @@ 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_chg_spin_ebd(self) -> bool:
+        """Return whether a frame charge/spin condition is configured."""
+        return self.add_chg_spin_ebd
+
     def get_default_chg_spin(self) -> list[float] | None:
         """Returns the default charge_spin values."""
         return self.default_chg_spin
@@ -755,7 +763,7 @@ def call(
             assert self.spin_embedding is not None
             chg_tebd = self.chg_embedding.call()
             spin_tebd = self.spin_embedding.call()
-            charge = xp.astype(charge_spin[:, 0], xp.int64) + 100
+            charge = xp.astype(charge_spin[:, 0], xp.int64) + CHARGE_OFFSET
             spin = xp.astype(charge_spin[:, 1], xp.int64)
             chg_ebd = xp.reshape(
                 xp.take(chg_tebd, xp.reshape(charge, (-1,)), axis=0),

deepmd/dpmodel/descriptor/dpa4.py

@@ -73,6 +73,9 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.charge_state import (
+    validate_charge_state,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -784,10 +787,10 @@ def __init__(
         self.edge_cartesian = bool(edge_cartesian)
         self.node_cartesian = str(node_cartesian)
         self.add_chg_spin_ebd = bool(add_chg_spin_ebd)
-        if default_chg_spin is not None and len(default_chg_spin) != 2:
-            raise ValueError("`default_chg_spin` must contain [charge, spin].")
         self.default_chg_spin = (
-            None if default_chg_spin is None else [float(x) for x in default_chg_spin]
+            None
+            if default_chg_spin is None
+            else validate_charge_state(default_chg_spin)
         )
 
         # === Native per-atom spin embedding ===
@@ -2282,6 +2285,10 @@ def get_ntypes(self) -> int:
     def get_type_map(self) -> list[str]:
         return self.type_map if self.type_map is not None else []
 
+    def has_chg_spin_ebd(self) -> bool:
+        """Return whether a frame charge/spin condition is configured."""
+        return self.charge_spin_embedding is not None
+
     def get_dim_chg_spin(self) -> int:
         """Return the charge/spin condition width."""
         return 2 if self.add_chg_spin_ebd else 0

deepmd/dpmodel/descriptor/dpa4_nn/__init__.py

@@ -73,6 +73,10 @@
     merge_lora_into_base,
     strip_lora_from_extra_state,
 )
+from .mlp import (
+    SwiGLUMLP,
+    resolve_swiglu_hidden_width,
+)
 from .norm import (
     EquivariantRMSNorm,
     ReducedEquivariantRMSNorm,
@@ -159,6 +163,7 @@
     "SeZMTypeEmbedding",
     "SpinEmbedding",
     "SwiGLU",
+    "SwiGLUMLP",
     "WignerDCalculator",
     "apply_lora_to_sezm",
     "build_cartesian_basis",
@@ -189,6 +194,7 @@
     "quaternion_z_rotation",
     "resolve_s2_grid_resolution",
     "resolve_so3_grid",
+    "resolve_swiglu_hidden_width",
     "safe_norm",
     "segment_envelope_gated_softmax",
     "so3_packed_index",

deepmd/dpmodel/descriptor/dpa4_nn/embedding.py

@@ -45,6 +45,11 @@
 from deepmd.dpmodel.utils.type_embed import (
     remap_atype_to_padding,
 )
+from deepmd.utils.charge_state import (
+    CHARGE_OFFSET,
+    CHARGE_TABLE_ROWS,
+    MULTIPLICITY_TABLE_ROWS,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -128,26 +133,41 @@ def __init__(
         # === Step 2. Register the embedding table parameter ===
         self.adam_type_embedding = table
 
-    def call(self, atype: Any) -> Any:
+    def call(self, atype: Any | None = None) -> Any:
         """
         Gather type embeddings.
 
         Parameters
         ----------
         atype
-            Atom types with shape (...,). Valid type range is [0, ntypes-1].
+            Atom types with shape (...). Valid type range is [0, ntypes-1].
+            If omitted, return the complete embedding table, including the
+            optional padding row. This form is used by graph-native descriptor
+            ABIs that precompute the table once per forward call.
 
         Returns
         -------
         Array
-            Type embeddings with shape (..., embed_dim).
+            Gathered type embeddings with shape ``(..., embed_dim)`` when
+            ``atype`` is provided. Otherwise, the complete table with shape
+            ``(ntypes + int(padding), embed_dim)``.
         """
+        # === Step 1. Return the complete graph-native lookup table ===
+        if atype is None:
+            xp = array_api_compat.array_namespace(self.adam_type_embedding)
+            return xp_asarray_nodetach(
+                xp,
+                self.adam_type_embedding[...],
+                device=array_api_compat.device(self.adam_type_embedding),
+            )
+
+        # === Step 2. Gather rows for an explicit atom-type tensor ===
         xp = array_api_compat.array_namespace(atype)
         weight = xp_asarray_nodetach(
             xp, self.adam_type_embedding[...], device=array_api_compat.device(atype)
         )
-        # torch.embedding gather: flatten the indices to int64, take the rows,
-        # then restore the original index shape.
+        # Flattening provides one backend-neutral gather while preserving every
+        # leading batch or graph dimension on restoration.
         index = xp.astype(xp.reshape(atype, (-1,)), xp.int64)
         if self.padding:
             index = remap_atype_to_padding(index, self.ntypes + 1)
@@ -869,15 +889,15 @@ def __init__(
             raise ValueError("`embed_dim` must be positive")
 
         self.charge_embedding = SeZMTypeEmbedding(
-            ntypes=200,
+            ntypes=CHARGE_TABLE_ROWS,
             embed_dim=self.embed_dim,
             precision=self.precision,
             seed=child_seed(seed, 0),
             trainable=self.trainable,
             padding=False,
         )
         self.spin_embedding = SeZMTypeEmbedding(
-            ntypes=100,
+            ntypes=MULTIPLICITY_TABLE_ROWS,
             embed_dim=self.embed_dim,
             precision=self.precision,
             seed=child_seed(seed, 1),
@@ -908,7 +928,7 @@ def call(self, charge_spin: Any) -> Any:
             Mixed condition embedding with shape (nf, embed_dim).
         """
         xp = array_api_compat.array_namespace(charge_spin)
-        charge = xp.astype(charge_spin[:, 0], xp.int64) + 100
+        charge = xp.astype(charge_spin[:, 0], xp.int64) + CHARGE_OFFSET
         spin = xp.astype(charge_spin[:, 1], xp.int64)
         charge_embed = self.charge_embedding(charge)
         spin_embed = self.spin_embedding(spin)