Clean up StateDict removal and fix standalone batch tests

Remove the residual dict-to-SimState fallback in both forward() methods
now that StateDict support is dropped. Fix batch tests in the standalone
package to use same-sized supercells (different lattice params) instead
of mixed unit cell + supercell, matching the inline test fix.

Author: HaoZeke

python/metatomic_torch/metatomic/torch/torchsim.py

@@ -185,18 +185,9 @@ def forward(self, state: ts.SimState) -> dict[str, torch.Tensor]:  # noqa: C901,
                 - 'stress': System stresses with shape [n_systems, 3, 3] if
                     compute_stress=True
         """
-        sim_state = (
-            state
-            if isinstance(state, ts.SimState)
-            else ts.SimState(**state, masses=torch.ones_like(state["positions"]))
-        )
-
-        # Input validation is already done inside the forward method of the
-        # AtomisticModel class, so we don't need to do it again here.
-
-        atomic_nums = sim_state.atomic_numbers
-        cell = sim_state.row_vector_cell
-        positions = sim_state.positions
+        atomic_nums = state.atomic_numbers
+        cell = state.row_vector_cell
+        positions = state.positions
 
         # Check dtype (metatomic models require a specific input dtype)
         if positions.dtype != self._dtype:
@@ -216,7 +207,7 @@ def forward(self, state: ts.SimState) -> dict[str, torch.Tensor]:  # noqa: C901,
         systems: list[System] = []
         strains = []
         for sys_idx in range(len(cell)):
-            system_mask = sim_state.system_idx == sys_idx
+            system_mask = state.system_idx == sys_idx
             system_positions = positions[system_mask]
             system_cell = cell[sys_idx]
             system_atomic_numbers = atomic_nums[system_mask]
@@ -237,7 +228,7 @@ def forward(self, state: ts.SimState) -> dict[str, torch.Tensor]:  # noqa: C901,
                     positions=system_positions,
                     types=system_atomic_numbers,
                     cell=system_cell,
-                    pbc=sim_state.pbc,
+                    pbc=state.pbc,
                 )
             )
 

python/metatomic_torchsim/metatomic/torchsim/_calculator.py

@@ -184,15 +184,9 @@ def forward(self, state: "ts.SimState") -> Dict[str, torch.Tensor]:
             ``"forces"`` (shape ``[n_atoms, 3]``, if ``compute_forces``), and
             ``"stress"`` (shape ``[n_systems, 3, 3]``, if ``compute_stress``).
         """
-        sim_state = (
-            state
-            if isinstance(state, ts.SimState)
-            else ts.SimState(**state, masses=torch.ones_like(state["positions"]))
-        )
-
-        positions = sim_state.positions
-        cell = sim_state.row_vector_cell
-        atomic_nums = sim_state.atomic_numbers
+        positions = state.positions
+        cell = state.row_vector_cell
+        atomic_nums = state.atomic_numbers
 
         if positions.dtype != self._dtype:
             raise TypeError(
@@ -207,7 +201,7 @@ def forward(self, state: "ts.SimState") -> Dict[str, torch.Tensor]:
         n_systems = len(cell)
 
         for sys_idx in range(n_systems):
-            mask = sim_state.system_idx == sys_idx
+            mask = state.system_idx == sys_idx
             sys_positions = positions[mask]
             sys_cell = cell[sys_idx]
             sys_types = atomic_nums[mask]
@@ -231,7 +225,7 @@ def forward(self, state: "ts.SimState") -> Dict[str, torch.Tensor]:
                     positions=sys_positions,
                     types=sys_types,
                     cell=sys_cell,
-                    pbc=sim_state.pbc,
+                    pbc=state.pbc,
                 )
             )
 

python/metatomic_torchsim/tests/test_torchsim.py

@@ -132,11 +132,21 @@ def test_forward_no_forces(lj_model, ni_atoms):
     assert "stress" in output
 
 
-def test_batched_forward(metatomic_model, ni_atoms):
-    """Forward pass handles batched systems correctly."""
+def _make_ni_atoms_2():
+    """Create a second Ni supercell (same size, different lattice parameter)."""
     import ase.build
 
-    atoms_2 = ase.build.bulk("Ni", "fcc", a=3.6, cubic=True)
+    np.random.seed(0xCAFEBABE)
+    atoms = ase.build.make_supercell(
+        ase.build.bulk("Ni", "fcc", a=3.5, cubic=True), 2 * np.eye(3)
+    )
+    atoms.positions += 0.1 * np.random.rand(*atoms.positions.shape)
+    return atoms
+
+
+def test_batched_forward(metatomic_model, ni_atoms):
+    """Forward pass handles batched systems correctly."""
+    atoms_2 = _make_ni_atoms_2()
     sim_state = ts.io.atoms_to_state([ni_atoms, atoms_2], DEVICE, DTYPE)
     output = metatomic_model(sim_state)
 
@@ -148,9 +158,7 @@ def test_batched_forward(metatomic_model, ni_atoms):
 
 def test_energy_consistency_single_vs_batch(metatomic_model, ni_atoms):
     """Energy from single system matches the corresponding entry in a batch."""
-    import ase.build
-
-    atoms_2 = ase.build.bulk("Ni", "fcc", a=3.6, cubic=True)
+    atoms_2 = _make_ni_atoms_2()
 
     # single
     state_1 = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE)