fea: site_charges -> partial_charges

Author: CompRhys

pyproject.toml

@@ -41,14 +41,12 @@ test = [
     "torch-sim-atomistic[io,symmetry,vesin]",
     "platformdirs>=4.0.0",
     "psutil>=7.0.0",
-    "pymatgen>=2025.6.14",
     "pytest-cov>=6",
     "pytest>=8",
     "spglib>=2.6",
-    "vesin[torch]>=0.5.3",
 ]
 vesin = ["vesin[torch]>=0.5.3"]
-io = ["ase>=3.26", "phonopy>=2.37.0", "pymatgen>=2025.6.14"]
+io = ["ase>=3.26", "phonopy>=2.37.0", "pymatgen>=2026.3.23"]
 symmetry = ["moyopy>=0.7.8"]
 mace = ["mace-torch>=0.3.15"]
 mattersim = ["mattersim>=0.1.2"]

tests/models/test_electrostatics.py

@@ -30,7 +30,7 @@ def _make_charged_state(
     charges = torch.empty(n, dtype=dtype, device=device)
     charges[::2] = 1.0
     charges[1::2] = -1.0
-    state._atom_extras["site_charges"] = charges  # noqa: SLF001
+    state._atom_extras["partial_charges"] = charges  # noqa: SLF001
     return state
 
 
@@ -49,33 +49,33 @@ def pme_model() -> PMEModel:
     return PMEModel(cutoff=8.0, accuracy=1e-6, device=DEVICE, dtype=DTYPE)
 
 
-def _add_site_charges(state: ts.SimState) -> ts.SimState:
+def _add_partial_charges(state: ts.SimState) -> ts.SimState:
     """Inject alternating +/-0.5 site charges into a state."""
     n = state.n_atoms
     charges = torch.zeros(n, dtype=state.positions.dtype, device=state.device)
     charges[::2] = 0.5
     charges[1::2] = -0.5
-    state._atom_extras["site_charges"] = charges  # noqa: SLF001
+    state._atom_extras["partial_charges"] = charges  # noqa: SLF001
     return state
 
 
 test_dsf_model_outputs = make_validate_model_outputs_test(
     model_fixture_name="dsf_model",
     device=DEVICE,
     dtype=DTYPE,
-    state_modifiers=[_add_site_charges],
+    state_modifiers=[_add_partial_charges],
 )
 test_ewald_model_outputs = make_validate_model_outputs_test(
     model_fixture_name="ewald_model",
     device=DEVICE,
     dtype=DTYPE,
-    state_modifiers=[_add_site_charges],
+    state_modifiers=[_add_partial_charges],
 )
 test_pme_model_outputs = make_validate_model_outputs_test(
     model_fixture_name="pme_model",
     device=DEVICE,
     dtype=DTYPE,
-    state_modifiers=[_add_site_charges],
+    state_modifiers=[_add_partial_charges],
 )
 
 

torch_sim/models/electrostatics.py

@@ -3,7 +3,7 @@
 Wraps the ``nvalchemiops`` Warp-accelerated electrostatics implementations as
 :class:`~torch_sim.models.interface.ModelInterface` subclasses, with full PBC,
 stress (virial), and batched system support.  Per-atom partial charges are read
-from ``state.site_charges`` (a SimState atom extra).
+from ``state.partial_charges`` (a SimState atom extra).
 """
 
 from __future__ import annotations
@@ -80,7 +80,7 @@ class DSFCoulombModel(ModelInterface):
     forces, and (optionally) stress.  All user-facing quantities are in
     metal units (Angstrom / eV); the Coulomb constant ``ke`` is baked in.
 
-    Per-atom partial charges are read from ``state.site_charges``.
+    Per-atom partial charges are read from ``state.partial_charges``.
 
     Args:
         cutoff: Real-space cutoff in Angstrom.
@@ -121,15 +121,18 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
         """Compute DSF electrostatic energy, forces, and stress.
 
         Args:
-            state: Simulation state with ``site_charges`` set as an
+            state: Simulation state with ``partial_charges`` set as an
                 atom extra (shape ``[n_atoms]``).
             **_kwargs: Unused; accepted for interface compatibility.
 
         Returns:
             dict with ``"energy"`` [n_systems], ``"forces"`` [n_atoms, 3],
             and (if ``compute_stress``) ``"stress"`` [n_systems, 3, 3].
         """
-        charges = state.site_charges
+        if not state.has_extras("partial_charges"):
+            raise ValueError("Partial charges are required for DSF Coulomb summation.")
+
+        charges = state.partial_charges
         edge_index, neighbor_ptr, unit_shifts = _build_csr(
             state, self.cutoff, self.neighbor_list_fn
         )
@@ -170,7 +173,7 @@ class EwaldModel(ModelInterface):
     Returns per-atom energies that are aggregated to per-system.  All
     user-facing quantities are in metal units (Angstrom / eV).
 
-    Per-atom partial charges are read from ``state.site_charges``.
+    Per-atom partial charges are read from ``state.partial_charges``.
 
     Requires periodic boundary conditions.
 
@@ -213,7 +216,7 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
         """Compute Ewald electrostatic energy, forces, and stress.
 
         Args:
-            state: Simulation state with ``site_charges`` set as an
+            state: Simulation state with ``partial_charges`` set as an
                 atom extra (shape ``[n_atoms]``).  Returns zeros for
                 non-periodic states.
             **_kwargs: Unused; accepted for interface compatibility.
@@ -222,11 +225,14 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
             dict with ``"energy"`` [n_systems], ``"forces"`` [n_atoms, 3],
             and (if ``compute_stress``) ``"stress"`` [n_systems, 3, 3].
         """
+        if not state.has_extras("partial_charges"):
+            raise ValueError("Partial charges are required for Ewald summation.")
+
         if not state.pbc.any():
             return _zero_result(
                 state, self._dtype, self._compute_forces, self._compute_stress
             )
-        charges = state.site_charges
+        charges = state.partial_charges
         edge_index, neighbor_ptr, unit_shifts = _build_csr(
             state, self.cutoff, self.neighbor_list_fn
         )
@@ -270,7 +276,7 @@ class PMEModel(ModelInterface):
     per-system.  All user-facing quantities are in metal units
     (Angstrom / eV).
 
-    Per-atom partial charges are read from ``state.site_charges``.
+    Per-atom partial charges are read from ``state.partial_charges``.
 
     Requires periodic boundary conditions.
 
@@ -322,7 +328,7 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
         """Compute PME electrostatic energy, forces, and stress.
 
         Args:
-            state: Simulation state with ``site_charges`` set as an
+            state: Simulation state with ``partial_charges`` set as an
                 atom extra (shape ``[n_atoms]``).  Returns zeros for
                 non-periodic states.
             **_kwargs: Unused; accepted for interface compatibility.
@@ -331,11 +337,14 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
             dict with ``"energy"`` [n_systems], ``"forces"`` [n_atoms, 3],
             and (if ``compute_stress``) ``"stress"`` [n_systems, 3, 3].
         """
+        if not state.has_extras("partial_charges"):
+            raise ValueError("Partial charges are required for PME summation.")
+
         if not state.pbc.any():
             return _zero_result(
                 state, self._dtype, self._compute_forces, self._compute_stress
             )
-        charges = state.site_charges
+        charges = state.partial_charges
         edge_index, neighbor_ptr, unit_shifts = _build_csr(
             state, self.cutoff, self.neighbor_list_fn
         )