fix(charge-state): enforce the shared table domain at every host boundary

DPA3, DPA4 and DPA4C all embed the frame condition by casting it to int64 and
gathering one row of a 200-entry charge table and one of a 100-entry
multiplicity table. The domain is therefore a property of the feature, not of
any descriptor, so the per-descriptor declaration chain added earlier is
replaced by one shared contract in `deepmd.utils.charge_state`. All three
descriptors size their tables from it and validate their configured default
against it, which closes the silent truncation that let DPA3 accept a
fractional state.

The archive keeps carrying the ranges so the C++ boundary enforces the same
rule without a Python model and without recompiling when the domain changes.
`check_call_charge_spin` and `set_charge_spin` now reject a state that
addresses no row on the per-call, folded and default paths; an archive frozen
before the ranges were recorded is checked on width alone, as before.

Author: OutisLi

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -254,10 +254,6 @@ def get_default_chg_spin(self) -> list[float] | None:
         """Get the default charge_spin values."""
         return None
 
-    def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
-        """Get the row range each charge_spin value indexes, or None."""
-        return None
-
     def reinit_atom_exclude(
         self,
         exclude_types: list[int] = [],

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -154,12 +154,6 @@ def get_default_chg_spin(self) -> list[float] | None:
             return self.descriptor.get_default_chg_spin()
         return None
 
-    def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
-        """Get the row range each charge_spin value indexes, or None."""
-        if self.add_chg_spin_ebd:
-            return self.descriptor.get_chg_spin_table_ranges()
-        return None
-
     def uses_graph_lower(self) -> bool:
         """Delegates to this model's own descriptor."""
         return bool(self.descriptor.uses_graph_lower())

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -740,20 +740,6 @@ def get_default_chg_spin(self) -> "Array | None":
             lambda m: m.get_default_chg_spin(),
         )[1]
 
-    def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
-        """The shared table row ranges, if the children agree.
-
-        A condition reaches every consuming child, so it must address the
-        tables of all of them. Children that disagree share no acceptable
-        state, which the composition reports as an unconstrained domain
-        rather than silently enforcing one child's tables on the others.
-        """
-        return self._agreed_default(
-            self._chg_spin_consumers(),
-            lambda m: m.get_chg_spin_table_ranges() is not None,
-            lambda m: m.get_chg_spin_table_ranges(),
-        )[1]
-
     def has_default_fparam(self) -> bool:
         """Whether every active child shares one default frame parameter."""
         return self._agreed_default(

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",
@@ -755,7 +759,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 ===

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,
 )
@@ -884,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),
@@ -923,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)

deepmd/dpmodel/descriptor/dpa4c.py

@@ -58,6 +58,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,
 )
@@ -73,7 +76,6 @@
     resolve_swiglu_hidden_width,
 )
 from .dpa4c_nn import (
-    CHARGE_STATE_TABLE_RANGES,
     ChargeStateEmbedding,
     InvariantReadout,
     OrderedPairFiLM,
@@ -85,7 +87,6 @@
     derive_bispectrum_ranks,
     derive_degree_channels,
     derive_spin_channels,
-    validate_charge_state,
 )
 
 if TYPE_CHECKING:
@@ -2004,19 +2005,6 @@ def get_default_chg_spin(self) -> list[float] | None:
         """Return the fallback ``[charge, multiplicity]``, if configured."""
         return self.default_chg_spin
 
-    def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
-        """Return the row range each value of the frame condition indexes.
-
-        The condition is embedded by gathering one row of the charge table and
-        one of the multiplicity table, so an acceptable state is a pair of
-        integers inside these half-open ranges. A folded condition indexes the
-        same tables at rebuild time, so the ranges hold whether or not the
-        descriptor is compressed.
-        """
-        if self.charge_spin_embedding is None:
-            return None
-        return [tuple(rng) for rng in CHARGE_STATE_TABLE_RANGES]
-
     def has_message_passing_across_ranks(self) -> bool:
         """Return whether intermediate halo communication is required."""
         return False

deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py

@@ -8,10 +8,8 @@
     enumerate_degree_triples,
 )
 from .charge_state import (
-    CHARGE_STATE_TABLE_RANGES,
     ChargeStateEmbedding,
     canonicalize_charge_spin,
-    validate_charge_state,
 )
 from .geometry import (
     MAX_ANGULAR_DEGREE,
@@ -34,7 +32,6 @@
 )
 
 __all__ = [
-    "CHARGE_STATE_TABLE_RANGES",
     "MAX_ANGULAR_DEGREE",
     "NEIGHBOR_QUADRUPOLE_CHANNELS",
     "BispectrumLayout",
@@ -52,5 +49,4 @@
     "derive_spin_channels",
     "enumerate_degree_triples",
     "packed_l2_to_stf",
-    "validate_charge_state",
 ]

deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py

@@ -52,6 +52,11 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
+from deepmd.utils.charge_state import (
+    CHARGE_OFFSET,
+    CHARGE_TABLE_ROWS,
+    MULTIPLICITY_TABLE_ROWS,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -69,78 +74,6 @@
         Array,
     )
 
-#: Rows of the charge table, covering integer charges in ``[-100, 99]``.
-CHARGE_TABLE_ROWS = 200
-
-#: Index of the neutral charge row, so row ``CHARGE_OFFSET + Q`` holds ``Q``.
-CHARGE_OFFSET = 100
-
-#: Rows of the spin table, covering integer multiplicities below this bound.
-MULTIPLICITY_TABLE_ROWS = 100
-
-#: Half-open range of representable total charges, in units of the elementary
-#: charge.
-CHARGE_RANGE = (-CHARGE_OFFSET, CHARGE_TABLE_ROWS - CHARGE_OFFSET)
-
-#: Half-open range of representable spin multiplicities.
-MULTIPLICITY_RANGE = (0, MULTIPLICITY_TABLE_ROWS)
-
-#: Name of each value of a charge state, in order, for diagnostics.
-CHARGE_STATE_FIELDS = ("charge", "multiplicity")
-
-#: Half-open row range addressed by each value of a charge state, in order.
-#: A condition is a pair of table row indices, so a host-side boundary that
-#: knows these ranges can reject an unaddressable state without knowing which
-#: descriptor holds the tables.
-CHARGE_STATE_TABLE_RANGES = (CHARGE_RANGE, MULTIPLICITY_RANGE)
-
-
-def validate_charge_state(charge_spin: Any) -> list[float]:
-    """Check that a frame condition addresses a row of each embedding table.
-
-    Both tables are indexed directly by the condition, and neither the gather
-    nor the compiled kernel bounds-checks that index, so an out-of-range value
-    would read past the table. Every host-side boundary that accepts a charge
-    state therefore passes it through here first. The per-forward path is
-    deliberately not guarded: its values come from the data pipeline, which
-    owns their validity exactly as it owns the validity of an atom type.
-
-    Parameters
-    ----------
-    charge_spin
-        A pair ``[charge, multiplicity]``, in any sequence form.
-
-    Returns
-    -------
-    list[float]
-        The same pair, as two floats.
-
-    Raises
-    ------
-    ValueError
-        If the pair does not hold exactly two integral values within the
-        representable ranges.
-    """
-    values = [float(value) for value in np.reshape(np.asarray(charge_spin), (-1,))]
-    if len(values) != 2:
-        raise ValueError(
-            f"A charge state must be a `[charge, multiplicity]` pair, got "
-            f"{len(values)} values"
-        )
-    for value, name, (low, high) in zip(
-        values,
-        CHARGE_STATE_FIELDS,
-        CHARGE_STATE_TABLE_RANGES,
-        strict=True,
-    ):
-        if not np.isfinite(value) or value != int(value):
-            raise ValueError(f"The {name} must be an integer, got {value}")
-        if not low <= value < high:
-            raise ValueError(
-                f"The {name} must lie in [{low}, {high}), got {int(value)}"
-            )
-    return values
-
 
 class ChargeStateEmbedding(NativeOP):
     r"""Embed the frame charge and spin multiplicity into two condition vectors.

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -113,17 +113,6 @@ def get_default_chg_spin(self) -> Any:
             """Returns the default charge_spin value, or None."""
             return None
 
-        def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
-            """Returns the row range each charge_spin value indexes, or None.
-
-            A descriptor that embeds the condition by indexing tables reports
-            one half-open range per value, which makes every acceptable state
-            an integer tuple inside those ranges. ``None``, the default, means
-            the condition enters as a continuous quantity and only its width
-            is constrained.
-            """
-            return None
-
         def get_geo_compress(self) -> bool:
             """Return whether geometric tabulated compression is active.