From 3f6e6fd73d0a06507346207e69a78546d8dbbfa7 Mon Sep 17 00:00:00 2001
From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com>
Date: Sat, 25 Jul 2026 12:24:43 +0900
Subject: [PATCH 1/5] feat(state): make pbc a per-system attribute

---
 tests/models/test_electrostatics.py | 40 +++++++++++++++
 tests/models/test_lennard_jones.py  | 24 +++++++++
 tests/test_autobatching.py          | 35 ++++++++++---
 tests/test_elastic.py               |  9 ++++
 tests/test_io.py                    | 34 +++++++++++++
 tests/test_state.py                 | 78 ++++++++++++++++++++++++++++-
 tests/test_trajectory.py            | 24 ++++++++-
 torch_sim/__init__.py               |  4 ++
 torch_sim/autobatching.py           | 27 ++--------
 torch_sim/elastic.py                |  3 +-
 torch_sim/io.py                     | 30 ++++-------
 torch_sim/models/electrostatics.py  | 32 ++++++++++--
 torch_sim/state.py                  | 76 +++++++++++++++++++++++++---
 torch_sim/trajectory.py             |  5 +-
 torch_sim/transforms.py             |  9 ++--
 15 files changed, 356 insertions(+), 74 deletions(-)

diff --git a/tests/models/test_electrostatics.py b/tests/models/test_electrostatics.py
index 8cd8bbcef..9903cfd83 100644
--- a/tests/models/test_electrostatics.py
+++ b/tests/models/test_electrostatics.py
@@ -93,6 +93,46 @@ def test_dsf_nonzero_energy() -> None:
     assert out["energy"].abs().item() > 0
 
 
+@pytest.mark.parametrize(
+    ("model_cls", "kwargs"),
+    [
+        pytest.param(EwaldModel, {"cutoff": 8.0, "accuracy": 1e-6}, id="ewald"),
+        pytest.param(
+            PMEModel,
+            {"cutoff": 8.0, "accuracy": 1e-6, "mesh_spacing": 1.0},
+            id="pme",
+        ),
+    ],
+)
+def test_mixed_pbc_batch_zero_fills_non_periodic(
+    model_cls: type[EwaldModel | PMEModel],
+    kwargs: dict[str, float],
+    benzene_sim_state: ts.SimState,
+) -> None:
+    """Mixed batches compute periodic systems and zero-fill non-periodic ones."""
+    model = model_cls(
+        **kwargs,
+        device=DEVICE,
+        dtype=DTYPE,
+        compute_forces=True,
+        compute_stress=True,
+    )
+    periodic = _make_charged_state()
+    # Charges only make the molecule batchable with the charged system; the
+    # values are never used because the non-periodic system is zero-filled.
+    molecule = _add_partial_charges(benzene_sim_state)
+    mixed = ts.concatenate_states([periodic, molecule])
+
+    out = model(mixed)
+    ref = model(periodic)
+    torch.testing.assert_close(out["energy"][0], ref["energy"][0])
+    torch.testing.assert_close(out["forces"][: periodic.n_atoms], ref["forces"])
+    torch.testing.assert_close(out["stress"][0], ref["stress"][0])
+    assert out["energy"][1] == 0
+    assert torch.all(out["forces"][periodic.n_atoms :] == 0)
+    assert torch.all(out["stress"][1] == 0)
+
+
 @pytest.mark.parametrize(
     ("model_cls", "kwargs"),
     [
diff --git a/tests/models/test_lennard_jones.py b/tests/models/test_lennard_jones.py
index a3e5ed5a5..78ea2d4f5 100644
--- a/tests/models/test_lennard_jones.py
+++ b/tests/models/test_lennard_jones.py
@@ -116,6 +116,30 @@ def test_lennard_jones_model_evaluation(si_double_sim_state: ts.SimState) -> Non
     assert results["stress"].shape == (si_double_sim_state.n_systems, 3, 3)
 
 
+def test_lennard_jones_model_mixed_pbc_batch(
+    ar_supercell_sim_state: ts.SimState, benzene_sim_state: ts.SimState
+) -> None:
+    """A batch mixing periodic and non-periodic systems matches per-system results."""
+    model = LennardJonesModel(
+        sigma=3.405,
+        epsilon=0.0104,
+        cutoff=2.5 * 3.405,
+        dtype=torch.float64,
+        compute_forces=True,
+    )
+    mixed = ts.concatenate_states([ar_supercell_sim_state, benzene_sim_state])
+    assert mixed.pbc.tolist() == [[True, True, True], [False, False, False]]
+
+    batched = model(mixed)
+    for sys_idx, single in enumerate((ar_supercell_sim_state, benzene_sim_state)):
+        single_results = model(single)
+        torch.testing.assert_close(
+            batched["energy"][sys_idx], single_results["energy"][0]
+        )
+        mask = mixed.system_idx == sys_idx
+        torch.testing.assert_close(batched["forces"][mask], single_results["forces"])
+
+
 def test_lennard_jones_model_force_conservation(
     si_double_sim_state: ts.SimState,
 ) -> None:
diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py
index 2e99e7750..a94994c72 100644
--- a/tests/test_autobatching.py
+++ b/tests/test_autobatching.py
@@ -122,11 +122,7 @@ def test_calculate_scaling_metric_non_periodic(benzene_sim_state: ts.SimState) -
         benzene_sim_state.positions.max(dim=0).values
         - benzene_sim_state.positions.min(dim=0).values
     ).clone()
-    pbc_tensor = torch.as_tensor(
-        benzene_sim_state.pbc, device=benzene_sim_state.device, dtype=torch.bool
-    )
-    if pbc_tensor.ndim == 0:
-        pbc_tensor = pbc_tensor.repeat(3)
+    pbc_tensor = benzene_sim_state.pbc.reshape(-1)
     for idx, p in enumerate(pbc_tensor):
         if not p:
             bbox[idx] += 2.0
@@ -147,7 +143,7 @@ def test_calculate_scaling_metric_mixed_pbc_uses_per_system_path(
             split_state.positions.max(dim=0).values
             - split_state.positions.min(dim=0).values
         ).clone()
-        split_state_pbc = torch.as_tensor(split_state.pbc, dtype=torch.bool).tolist()
+        split_state_pbc = split_state.pbc.reshape(-1).tolist()
         for axis_idx, is_periodic in enumerate(split_state_pbc):
             if not is_periodic:
                 bbox[axis_idx] += 2.0
@@ -158,6 +154,33 @@ def test_calculate_scaling_metric_mixed_pbc_uses_per_system_path(
     assert metric_values == pytest.approx(expected_values, rel=1e-5)
 
 
+def test_calculate_scaling_metric_mixed_pbc_batch(
+    si_sim_state: ts.SimState, benzene_sim_state: ts.SimState
+) -> None:
+    """Full-pbc, partial-pbc, and molecule systems each use the right volume."""
+    partial = ts.SimState.from_state(si_sim_state, pbc=[True, False, True])
+    state = ts.concatenate_states([si_sim_state, partial, benzene_sim_state])
+
+    metric_values = calculate_memory_scalers(state, "n_atoms_x_density")
+    expected_values: list[float] = []
+    for split_state in state.split():
+        if split_state.pbc.all():
+            volume = torch.abs(torch.linalg.det(split_state.cell[0])) / 1000
+        else:
+            bbox = (
+                split_state.positions.max(dim=0).values
+                - split_state.positions.min(dim=0).values
+            ).clone()
+            for axis_idx, is_periodic in enumerate(split_state.pbc.reshape(-1).tolist()):
+                if not is_periodic:
+                    bbox[axis_idx] += 2.0
+            volume = bbox.prod() / 1000
+        expected_values.append(
+            split_state.n_atoms * (split_state.n_atoms / volume.item())
+        )
+    assert metric_values == pytest.approx(expected_values, rel=1e-5)
+
+
 def test_n_edges_scalers_periodic(si_sim_state: ts.SimState) -> None:
     """n_edges scalers for a single periodic system have correct shape and type."""
     result = _n_edges_scalers(si_sim_state, cutoff=5.0)
diff --git a/tests/test_elastic.py b/tests/test_elastic.py
index 2ac174410..554e5b5c9 100644
--- a/tests/test_elastic.py
+++ b/tests/test_elastic.py
@@ -13,6 +13,7 @@
     get_elementary_deformations,
     get_strain,
 )
+from torch_sim.models.lennard_jones import LennardJonesModel
 from torch_sim.typing import BravaisType
 from torch_sim.units import UnitConversion
 
@@ -25,6 +26,14 @@
     pytest.skip(f"MACE not installed: {traceback.format_exc()}", allow_module_level=True)
 
 
+def test_calculate_elastic_tensor_requires_full_pbc(
+    benzene_sim_state: ts.SimState, lj_model: LennardJonesModel
+) -> None:
+    """Elastic tensor calculation rejects systems that are not fully periodic."""
+    with pytest.raises(ValueError, match="fully periodic"):
+        calculate_elastic_tensor(benzene_sim_state, lj_model)
+
+
 def test_get_strain_zero_deformation(cu_sim_state: ts.SimState) -> None:
     """Test that zero deformation produces zero strain."""
     # Test with same state as reference and deformed - should give zero strain
diff --git a/tests/test_io.py b/tests/test_io.py
index 63616273b..9a8e7e8eb 100644
--- a/tests/test_io.py
+++ b/tests/test_io.py
@@ -34,6 +34,40 @@ def test_single_atoms_to_state(si_atoms: Atoms) -> None:
     assert state.atomic_numbers.dtype == torch.int
 
 
+def test_atoms_to_state_mixed_pbc(si_atoms: Atoms, benzene_atoms: Atoms) -> None:
+    """Atoms with different pbc convert to per-system pbc and round-trip."""
+    slab_atoms = si_atoms.copy()
+    slab_atoms.set_pbc([True, True, False])
+    atoms_list = [si_atoms, slab_atoms, benzene_atoms]
+    state = ts.io.atoms_to_state(atoms_list, DEVICE, DTYPE)
+    assert state.pbc.tolist() == [
+        [True, True, True],
+        [True, True, False],
+        [False, False, False],
+    ]
+
+    round_trip = state.to_atoms()
+    for converted, original in zip(round_trip, atoms_list, strict=True):
+        np.testing.assert_array_equal(converted.pbc, original.pbc)
+
+
+def test_structures_to_state_mixed_pbc(si_structure: Any) -> None:
+    """Structures with different pbc convert to per-system pbc and round-trip."""
+    from pymatgen.core import Lattice, Structure
+
+    slab = Structure(
+        Lattice(si_structure.lattice.matrix, pbc=(True, True, False)),
+        [site.specie for site in si_structure],
+        si_structure.frac_coords,
+    )
+    state = ts.io.structures_to_state([si_structure, slab], DEVICE, DTYPE)
+    assert state.pbc.tolist() == [[True, True, True], [True, True, False]]
+
+    round_trip = state.to_structures()
+    assert round_trip[0].lattice.pbc == (True, True, True)
+    assert round_trip[1].lattice.pbc == (True, True, False)
+
+
 @pytest.mark.parametrize(
     ("system_extras_map", "atom_extras_map", "expected_sys", "expected_atom"),
     [
diff --git a/tests/test_state.py b/tests/test_state.py
index 7c73a6498..2e68b923a 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -37,9 +37,83 @@ def test_get_attrs_for_scope(si_sim_state: SimState) -> None:
     per_atom_attrs = dict(get_attrs_for_scope(si_sim_state, "per-atom"))
     assert set(per_atom_attrs) == {"positions", "masses", "atomic_numbers", "system_idx"}
     per_system_attrs = dict(get_attrs_for_scope(si_sim_state, "per-system"))
-    assert set(per_system_attrs) == {"cell"}
+    assert set(per_system_attrs) == {"cell", "pbc"}
     global_attrs = dict(get_attrs_for_scope(si_sim_state, "global"))
-    assert set(global_attrs) == {"pbc", "_rng"}
+    assert set(global_attrs) == {"_rng"}
+
+
+def test_mixed_pbc_batch_operations(
+    ar_supercell_sim_state: SimState, benzene_sim_state: SimState
+) -> None:
+    """Per-system pbc survives concatenation, slicing, splitting, and popping."""
+    periodic, molecule = ar_supercell_sim_state, benzene_sim_state
+    assert periodic.pbc.shape == (1, 3)
+    assert not molecule.pbc.any()
+
+    state = ts.concatenate_states([periodic, molecule])
+    assert torch.equal(state.pbc, torch.cat([periodic.pbc, molecule.pbc]))
+    assert torch.equal(state[1].pbc, molecule.pbc)
+
+    split = state.split()
+    assert torch.equal(split[0].pbc, periodic.pbc)
+    assert torch.equal(split[1].pbc, molecule.pbc)
+
+    popped = state.pop(0)
+    assert torch.equal(popped[0].pbc, periodic.pbc)
+    assert torch.equal(state.pbc, molecule.pbc)
+
+
+def test_pbc_broadcast_and_assignment(si_double_sim_state: SimState) -> None:
+    """Scalar and (3,) pbc inputs broadcast to shape (n_systems, 3)."""
+    state = si_double_sim_state
+    assert state.pbc.shape == (2, 3)
+    assert state.pbc.all()
+
+    state.pbc = [True, False, True]
+    assert torch.equal(state.pbc, torch.tensor([[True, False, True]] * 2, device=DEVICE))
+
+    state.pbc = False
+    assert state.pbc.shape == (2, 3)
+    assert not state.pbc.any()
+
+    # Only (3,) is broadcast; 2D pbc must match (n_systems, 3)
+    with pytest.raises(ValueError, match="pbc must have shape"):
+        SimState(
+            positions=state.positions,
+            masses=state.masses,
+            cell=state.cell,
+            pbc=torch.ones(3, 3, dtype=torch.bool),
+            atomic_numbers=state.atomic_numbers,
+            system_idx=state.system_idx,
+        )
+
+
+def test_require_pbc_helpers(
+    ar_supercell_sim_state: SimState, benzene_sim_state: SimState
+) -> None:
+    """require_uniform_pbc and require_full_pbc validate per-system pbc."""
+    uniform = ar_supercell_sim_state
+    row = ts.require_uniform_pbc(uniform, "test feature")
+    assert torch.equal(row, torch.tensor([True, True, True], device=DEVICE))
+    ts.require_full_pbc(uniform, "test feature")
+
+    mixed = ts.concatenate_states([uniform, benzene_sim_state])
+    with pytest.raises(ValueError, match="mixed pbc"):
+        ts.require_uniform_pbc(mixed, "test feature")
+    with pytest.raises(ValueError, match="fully periodic"):
+        ts.require_full_pbc(mixed, "test feature")
+
+
+def test_wrap_positions_mixed_pbc_batch(
+    ar_supercell_sim_state: SimState, benzene_sim_state: SimState
+) -> None:
+    """Wrapping only touches the periodic systems of a mixed batch."""
+    state = ts.concatenate_states([ar_supercell_sim_state, benzene_sim_state])
+    state.positions = state.positions + 100.0
+    wrapped = state.wrap_positions
+    n_periodic = ar_supercell_sim_state.n_atoms
+    assert not torch.allclose(wrapped[:n_periodic], state.positions[:n_periodic])
+    assert torch.equal(wrapped[n_periodic:], state.positions[n_periodic:])
 
 
 def test_all_attributes_must_be_specified_in_scopes() -> None:
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index 144b2c63b..b1a2c61c3 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -458,7 +458,7 @@ def test_get_atoms(trajectory: TorchSimTrajectory, random_state: MDState) -> Non
     np.testing.assert_allclose(
         atoms.get_atomic_numbers(), random_state.atomic_numbers.numpy()
     )
-    np.testing.assert_array_equal(atoms.pbc, random_state.pbc.detach().cpu().numpy())
+    np.testing.assert_array_equal(atoms.pbc, random_state.pbc[0].detach().cpu().numpy())
 
 
 def test_get_state(trajectory: TorchSimTrajectory, random_state: MDState) -> None:
@@ -531,7 +531,7 @@ def test_write_ase_trajectory(
         np.testing.assert_allclose(
             atoms.get_atomic_numbers(), random_state.atomic_numbers.numpy()
         )
-        np.testing.assert_array_equal(atoms.pbc, random_state.pbc.numpy())
+        np.testing.assert_array_equal(atoms.pbc, random_state.pbc[0].numpy())
 
     # Clean up
     ase_traj.close()
@@ -756,6 +756,26 @@ def test_property_model_consistency(
         traj.close()
 
 
+def test_reporter_mixed_pbc_batch(
+    ar_supercell_sim_state: SimState, benzene_sim_state: SimState, tmp_path: Path
+) -> None:
+    """Each per-system trajectory file stores that system's (3,) pbc row."""
+    state = ts.concatenate_states([ar_supercell_sim_state, benzene_sim_state])
+    filenames = [tmp_path / "periodic.hdf5", tmp_path / "molecule.hdf5"]
+    reporter = TrajectoryReporter(filenames, state_frequency=1)
+    reporter.report(state, 0)
+    reporter.close()
+
+    expected_rows = ([True] * 3, [False] * 3)
+    for filename, expected in zip(filenames, expected_rows, strict=True):
+        with TorchSimTrajectory(filename, mode="r") as traj:
+            pbc_arr = traj.get_array("pbc")
+            assert pbc_arr.shape == (3,)
+            assert pbc_arr.tolist() == expected
+            round_trip = traj.get_state(frame=0)
+        assert round_trip.pbc.tolist() == [expected]
+
+
 def test_reporter_with_model(
     si_double_sim_state: SimState, tmp_path: Path, lj_model: LennardJonesModel
 ) -> None:
diff --git a/torch_sim/__init__.py b/torch_sim/__init__.py
index 9a848f9f2..e1c2a146c 100644
--- a/torch_sim/__init__.py
+++ b/torch_sim/__init__.py
@@ -98,6 +98,8 @@
     concatenate_states,
     detach_state_graph,
     initialize_state,
+    require_full_pbc,
+    require_uniform_pbc,
 )
 from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter
 
@@ -200,6 +202,8 @@
     "optimize",
     "optimizers",
     "quantities",
+    "require_full_pbc",
+    "require_uniform_pbc",
     "runners",
     "state",
     "static",
diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py
index fa8a92094..c32aae721 100644
--- a/torch_sim/autobatching.py
+++ b/torch_sim/autobatching.py
@@ -340,11 +340,7 @@ def calculate_memory_scalers(
     if memory_scales_with == "n_atoms":
         return state.n_atoms_per_system.tolist()
     if memory_scales_with == "n_atoms_x_density":
-        pbc_all = (
-            state.pbc.all().item()
-            if torch.is_tensor(state.pbc)
-            else (state.pbc if isinstance(state.pbc, bool) else all(state.pbc))
-        )
+        pbc_all = state.pbc.all().item()
         if state.n_systems > 1 and pbc_all:
             # Vectorized volume only valid when all axes periodic
             n_atoms = state.n_atoms_per_system.to(state.volume.dtype)
@@ -354,31 +350,14 @@ def calculate_memory_scalers(
         scalers = []
         for system_idx in range(state.n_systems):
             system_state = state[system_idx]
-            system_pbc_all = (
-                system_state.pbc
-                if isinstance(system_state.pbc, bool)
-                else (
-                    all(system_state.pbc)
-                    if isinstance(system_state.pbc, (list, tuple))
-                    else system_state.pbc.all().item()
-                )
-            )
-            if system_pbc_all:
+            if system_state.pbc.all().item():
                 volume = torch.abs(torch.linalg.det(system_state.cell[0])) / 1000
             else:
                 bbox = (
                     system_state.positions.max(dim=0).values
                     - system_state.positions.min(dim=0).values
                 )
-                pbc_iter: tuple[bool, ...] | list[bool] = (
-                    (system_state.pbc,) * 3
-                    if isinstance(system_state.pbc, bool)
-                    else (
-                        system_state.pbc.tolist()
-                        if torch.is_tensor(system_state.pbc)
-                        else system_state.pbc
-                    )
-                )
+                pbc_iter = system_state.pbc.reshape(-1).tolist()
                 for axis_idx, periodic in enumerate(pbc_iter):
                     if not periodic:
                         bbox[axis_idx] += 2.0
diff --git a/torch_sim/elastic.py b/torch_sim/elastic.py
index 72341ebb1..ec21b8f51 100644
--- a/torch_sim/elastic.py
+++ b/torch_sim/elastic.py
@@ -29,7 +29,7 @@
 from torch_sim.autobatching import BinningAutoBatcher
 from torch_sim.models.interface import ModelInterface
 from torch_sim.optimizers import OptimState
-from torch_sim.state import SimState
+from torch_sim.state import SimState, require_full_pbc
 from torch_sim.typing import BravaisType
 
 
@@ -1132,6 +1132,7 @@ def calculate_elastic_tensor(
     Returns:
         torch.Tensor: Elastic tensor
     """
+    require_full_pbc(state, "Elastic tensor calculation")
     device, dtype = state.device, state.dtype
 
     # Calculate deformations for the bravais type
diff --git a/torch_sim/io.py b/torch_sim/io.py
index ee5a46e5e..3f504e933 100644
--- a/torch_sim/io.py
+++ b/torch_sim/io.py
@@ -171,13 +171,8 @@ def state_to_structures(state: ts.SimState) -> list[Structure]:
         species = [Element.from_Z(z) for z in system_numbers]
 
         # Create structure for this system
-        if torch.is_tensor(state.pbc):
-            pbc_tup = tuple(state.pbc.tolist())
-        elif isinstance(state.pbc, bool):
-            pbc_tup = (state.pbc, state.pbc, state.pbc)
-        else:
-            pbc_tup = (bool(state.pbc[0]), bool(state.pbc[1]), bool(state.pbc[2]))
-        pbc_tup = (pbc_tup[0], pbc_tup[1], pbc_tup[2])  # ensure tuple[bool, bool, bool]
+        pbc_row = state.pbc[uniq_sys_idx].tolist()
+        pbc_tup = (bool(pbc_row[0]), bool(pbc_row[1]), bool(pbc_row[2]))
         struct = Structure(
             lattice=Lattice(system_cell, pbc=pbc_tup),
             species=species,
@@ -283,12 +278,10 @@ def atoms_to_state(
 
     Raises:
         ImportError: If ASE is not installed.
-        ValueError: If systems have inconsistent periodic boundary conditions.
 
     Notes:
         - Input positions and cell should be in Å
         - Input masses should be in amu
-        - All systems must have consistent periodic boundary conditions
     """
     try:
         from ase import Atoms
@@ -319,9 +312,9 @@ def atoms_to_state(
         torch.arange(len(atoms_list), device=device), atoms_per_system
     )
 
-    # Verify consistent pbc
-    if not all(np.all(np.equal(at.pbc, atoms_list[0].pbc)) for at in atoms_list[1:]):
-        raise ValueError("All systems must have the same periodic boundary conditions")
+    pbc = torch.tensor(
+        np.stack([at.pbc for at in atoms_list]), dtype=torch.bool, device=device
+    )
 
     _system_extras: dict[str, torch.Tensor] = {}
     if system_extras_map:
@@ -347,7 +340,7 @@ def atoms_to_state(
         positions=positions,
         masses=masses,
         cell=cell,
-        pbc=atoms_list[0].pbc,
+        pbc=pbc,
         atomic_numbers=atomic_numbers,
         system_idx=system_idx,
         _system_extras=_system_extras,
@@ -423,19 +416,14 @@ def structures_to_state(
         torch.arange(len(struct_list), device=device), atoms_per_system
     )
 
-    # Verify consistent pbc
-    if not all(tuple(s.pbc) == tuple(struct_list[0].pbc) for s in struct_list[1:]):
-        raise ValueError("All systems must have the same periodic boundary conditions")
-
-    pbc_struct = struct_list[0].pbc
-    pbc_state: torch.Tensor | list[bool] | bool = (
-        list(pbc_struct) if isinstance(pbc_struct, (list, tuple)) else pbc_struct
+    pbc = torch.tensor(
+        np.array([s.pbc for s in struct_list]), dtype=torch.bool, device=device
     )
     return ts.SimState(
         positions=positions,
         masses=masses,
         cell=cell,
-        pbc=pbc_state,
+        pbc=pbc,
         atomic_numbers=atomic_numbers,
         system_idx=system_idx,
     )
diff --git a/torch_sim/models/electrostatics.py b/torch_sim/models/electrostatics.py
index f9c69e620..56727502f 100644
--- a/torch_sim/models/electrostatics.py
+++ b/torch_sim/models/electrostatics.py
@@ -48,6 +48,20 @@ def _zero_result(
     return results
 
 
+def _masked_forward(
+    model: ModelInterface, state: SimState, has_pbc: torch.Tensor
+) -> dict[str, torch.Tensor]:
+    """Evaluate the model on systems with any periodic axis, zero-filling the rest."""
+    sub_results = model.forward(state[torch.where(has_pbc)[0]])
+    results = _zero_result(state, model.dtype, model.compute_forces, model.compute_stress)
+    results["energy"][has_pbc] = sub_results["energy"]
+    if model.compute_forces:
+        results["forces"][has_pbc[state.system_idx]] = sub_results["forces"]
+    if model.compute_stress:
+        results["stress"][has_pbc] = sub_results["stress"]
+    return results
+
+
 class DSFCoulombModel(ModelInterface):
     """Damped Shifted Force electrostatics as a :class:`ModelInterface`.
 
@@ -205,7 +219,7 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
         Args:
             state: Simulation state with ``partial_charges`` set as an
                 atom extra (shape ``[n_atoms]``).  Returns zeros for
-                non-periodic states.
+                non-periodic systems.
             **_kwargs: Unused; accepted for interface compatibility.
 
         Returns:
@@ -215,10 +229,15 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
         if not state.has_extras("partial_charges"):
             raise ValueError("Partial charges are required for Ewald summation.")
 
-        if not state.pbc.any():
+        # Systems with any periodic axis go to the 3D kernels; only systems
+        # with no periodic axis are zero-filled.
+        has_pbc = state.pbc.any(dim=1)
+        if not has_pbc.any():
             return _zero_result(
                 state, self._dtype, self._compute_forces, self._compute_stress
             )
+        if not has_pbc.all():
+            return _masked_forward(self, state, has_pbc)
         charges = state.partial_charges
         edge_index, _mapping, unit_shifts = self.neighbor_list_fn(
             state.positions,
@@ -329,7 +348,7 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
         Args:
             state: Simulation state with ``partial_charges`` set as an
                 atom extra (shape ``[n_atoms]``).  Returns zeros for
-                non-periodic states.
+                non-periodic systems.
             **_kwargs: Unused; accepted for interface compatibility.
 
         Returns:
@@ -339,10 +358,15 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor]
         if not state.has_extras("partial_charges"):
             raise ValueError("Partial charges are required for PME summation.")
 
-        if not state.pbc.any():
+        # Systems with any periodic axis go to the 3D kernels; only systems
+        # with no periodic axis are zero-filled.
+        has_pbc = state.pbc.any(dim=1)
+        if not has_pbc.any():
             return _zero_result(
                 state, self._dtype, self._compute_forces, self._compute_stress
             )
+        if not has_pbc.all():
+            return _masked_forward(self, state, has_pbc)
         charges = state.partial_charges
         edge_index, _mapping, unit_shifts = self.neighbor_list_fn(
             state.positions,
diff --git a/torch_sim/state.py b/torch_sim/state.py
index 768e0f9bd..512290f4c 100644
--- a/torch_sim/state.py
+++ b/torch_sim/state.py
@@ -141,8 +141,9 @@ class SimState:
             the row vector convention `[[a1, a2, a3], [b1, b2, b3], [c1, c2, c3]]`
             used by ASE.
         pbc (bool | list[bool] | torch.Tensor): indicates periodic boundary
-            conditions in each axis. If a boolean is provided, all axes are
-            assumed to have the same periodic boundary conditions.
+            conditions per system and axis, stored with shape (n_systems, 3).
+            A boolean, a list of three booleans, or a tensor of shape (3,) is
+            broadcast to all systems.
         atomic_numbers (torch.Tensor): Atomic numbers with shape (n_atoms,)
         system_idx (torch.Tensor): Maps each atom index to its system index.
             Has shape (n_atoms,), must be unique consecutive integers starting from 0.
@@ -207,8 +208,8 @@ def __init__(  # noqa: D107
         "atomic_numbers",
         "system_idx",
     }
-    _system_attributes: ClassVar[set[str]] = {"cell"}
-    _global_attributes: ClassVar[set[str]] = {"pbc", "_rng"}
+    _system_attributes: ClassVar[set[str]] = {"cell", "pbc"}
+    _global_attributes: ClassVar[set[str]] = {"_rng"}
 
     @property
     def rng(self) -> torch.Generator:
@@ -237,10 +238,16 @@ def __setattr__(self, name: str, value: object) -> None:  # noqa: C901
                     else:
                         del extras[name]
                     return
-        if name == "pbc" and not isinstance(value, torch.Tensor):
-            if isinstance(value, bool):
-                value = [value] * 3
-            value = torch.tensor(value, dtype=torch.bool, device=self.device)
+        if name == "pbc":
+            if not isinstance(value, torch.Tensor):
+                if isinstance(value, bool):
+                    value = [value] * 3
+                value = torch.tensor(value, dtype=torch.bool, device=self.device)
+            # Broadcast a single (3,) row to all systems. During __init__ the pbc
+            # field is assigned before system_idx, so broadcasting then happens
+            # in __post_init__ instead.
+            if value.ndim == 1 and getattr(self, "system_idx", None) is not None:
+                value = value.unsqueeze(0).expand(self.n_systems, -1).contiguous()
         elif name == "system_idx":
             if value is None:
                 if hasattr(self, "positions"):
@@ -286,6 +293,16 @@ def __post_init__(self) -> None:  # noqa: C901
                 f"got {self.cell.shape}"
             )
 
+        # pbc is coerced to a tensor in __setattr__ but assigned before
+        # system_idx, so a single (3,) row is broadcast to all systems here
+        if self.pbc.ndim == 1:
+            self.pbc = self.pbc.unsqueeze(0).expand(n_systems, -1).contiguous()
+        if self.pbc.shape != (n_systems, 3):
+            raise ValueError(
+                f"pbc must have shape (n_systems={n_systems}, 3), "
+                f"got {tuple(self.pbc.shape)}"
+            )
+
         # if devices aren't all the same, raise an error, in a clean way
         devices = {
             attr: getattr(self, attr).device
@@ -904,6 +921,49 @@ def deform_grad(self) -> torch.Tensor:
         return self._deform_grad(self.reference_row_vector_cell, self.row_vector_cell)
 
 
+def require_uniform_pbc(state: SimState, feature: str) -> torch.Tensor:
+    """Return the pbc row shared by all systems, raising if systems differ.
+
+    For consumers that cannot handle per-system pbc, e.g. model featurizers
+    that treat pbc as a single global setting.
+
+    Args:
+        state (SimState): The state to validate.
+        feature (str): Name of the feature used in the error message.
+
+    Returns:
+        torch.Tensor: The shared pbc row with shape (3,).
+
+    Raises:
+        ValueError: If systems have different pbc.
+    """
+    pbc = state.pbc
+    if (pbc != pbc[0]).any():
+        raise ValueError(
+            f"{feature} does not support mixed pbc across systems; "
+            f"got per-system pbc {pbc.tolist()}. Split the batch by pbc instead."
+        )
+    return pbc[0]
+
+
+def require_full_pbc(state: SimState, feature: str) -> None:
+    """Raise if any system is not fully periodic.
+
+    Args:
+        state (SimState): The state to validate.
+        feature (str): Name of the feature used in the error message.
+
+    Raises:
+        ValueError: If any system has a non-periodic axis.
+    """
+    if not state.pbc.all():
+        bad = torch.where(~state.pbc.all(dim=1))[0]
+        raise ValueError(
+            f"{feature} requires fully periodic systems; systems "
+            f"{bad.tolist()} have pbc {state.pbc[bad].tolist()}."
+        )
+
+
 def _normalize_system_indices(
     system_indices: int | Sequence[int] | slice | torch.Tensor,
     n_systems: int,
diff --git a/torch_sim/trajectory.py b/torch_sim/trajectory.py
index 338a827e3..279298933 100644
--- a/torch_sim/trajectory.py
+++ b/torch_sim/trajectory.py
@@ -969,13 +969,14 @@ def write_state(  # noqa: C901
             self.write_arrays({"atomic_numbers": state[0].atomic_numbers}, 0)
 
         if "pbc" not in self.array_registry:
-            pbc_val = state[0].pbc
+            pbc_val = sub_states[0].pbc
             pbc_arr = (
                 pbc_val
                 if torch.is_tensor(pbc_val)
                 else torch.tensor([pbc_val] * 3 if isinstance(pbc_val, bool) else pbc_val)
             )
-            self.write_global_array("pbc", pbc_arr)
+            # Trajectory files hold a single system, so store its (3,) pbc row
+            self.write_global_array("pbc", pbc_arr.reshape(-1))
 
         # Write all arrays to file
         self.write_arrays(data, steps)
diff --git a/torch_sim/transforms.py b/torch_sim/transforms.py
index cea65f0ef..72576ce03 100644
--- a/torch_sim/transforms.py
+++ b/torch_sim/transforms.py
@@ -130,9 +130,10 @@ def pbc_wrap_batched(
             lattice vectors as column vectors.
         system_idx (torch.Tensor): Tensor of shape (n_atoms,) containing system
             indices for each atom.
-        pbc (torch.Tensor | bool): Tensor of shape (3,) containing boolean values
-            indicating whether periodic boundary conditions are applied in each dimension.
-            Can also be a bool. Defaults to True.
+        pbc (torch.Tensor | bool): Tensor of shape (n_systems, 3) or (3,) containing
+            boolean values indicating whether periodic boundary conditions are applied
+            in each dimension. Can also be a bool. A (3,) tensor or bool is broadcast
+            to all systems. Defaults to True.
 
     Returns:
         torch.Tensor: Wrapped positions in real space with same shape as input positions.
@@ -150,7 +151,7 @@ def pbc_wrap_batched(
             f"Number of unique systems ({n_systems}) doesn't "
             f"match number of cells ({cell.shape[0]})"
         )
-    pbc_batched = pbc.unsqueeze(0).expand(n_systems, -1)
+    pbc_batched = pbc.unsqueeze(0).expand(n_systems, -1) if pbc.ndim == 1 else pbc
     wrapped, _ = pbc_wrap_batched_and_get_lattice_shifts(
         positions, cell.mT, system_idx, pbc_batched
     )

From f0cd65f7f10d1a8bfc893413407aed83e0decdc3 Mon Sep 17 00:00:00 2001
From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com>
Date: Sat, 25 Jul 2026 13:57:54 +0900
Subject: [PATCH 2/5] refactor(state): declare pbc after system_idx to
 broadcast in one place

During __init__ the dataclass assigned pbc before system_idx, so the (3,) to (n_systems, 3) broadcast in __setattr__ could not run at construction time and had to be duplicated in __post_init__. Declaring pbc after system_idx makes n_systems available when pbc is assigned, leaving a single broadcast site in __setattr__ and only shape validation in __post_init__. kw_only=True keeps the public API unchanged.
---
 torch_sim/state.py | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/torch_sim/state.py b/torch_sim/state.py
index 512290f4c..5805fee39 100644
--- a/torch_sim/state.py
+++ b/torch_sim/state.py
@@ -178,9 +178,9 @@ class SimState:
     positions: torch.Tensor
     masses: torch.Tensor
     cell: torch.Tensor
-    pbc: torch.Tensor  # coerced from bool/list[bool] by __setattr__
     atomic_numbers: torch.Tensor
     system_idx: torch.Tensor = field(default=None)  # type: ignore[assignment]  # coerced from None by __setattr__
+    pbc: torch.Tensor  # coerced from bool/list[bool] by __setattr__
     _constraints: list["Constraint"] = field(default_factory=list)
     _system_extras: dict[str, torch.Tensor] = field(default_factory=dict)
     _atom_extras: dict[str, torch.Tensor] = field(default_factory=dict)
@@ -194,8 +194,8 @@ def __init__(  # noqa: D107
             positions: torch.Tensor,
             masses: torch.Tensor,
             cell: torch.Tensor,
-            pbc: torch.Tensor | list[bool] | bool,
             atomic_numbers: torch.Tensor,
+            pbc: torch.Tensor | list[bool] | bool,
             system_idx: torch.Tensor | None = None,
             _constraints: list[Constraint] | None = None,
             _rng: PRNGLike = None,
@@ -243,9 +243,8 @@ def __setattr__(self, name: str, value: object) -> None:  # noqa: C901
                 if isinstance(value, bool):
                     value = [value] * 3
                 value = torch.tensor(value, dtype=torch.bool, device=self.device)
-            # Broadcast a single (3,) row to all systems. During __init__ the pbc
-            # field is assigned before system_idx, so broadcasting then happens
-            # in __post_init__ instead.
+            # Broadcast a single (3,) row to all systems. pbc is declared after
+            # system_idx, so n_systems is available already during __init__.
             if value.ndim == 1 and getattr(self, "system_idx", None) is not None:
                 value = value.unsqueeze(0).expand(self.n_systems, -1).contiguous()
         elif name == "system_idx":
@@ -293,10 +292,6 @@ def __post_init__(self) -> None:  # noqa: C901
                 f"got {self.cell.shape}"
             )
 
-        # pbc is coerced to a tensor in __setattr__ but assigned before
-        # system_idx, so a single (3,) row is broadcast to all systems here
-        if self.pbc.ndim == 1:
-            self.pbc = self.pbc.unsqueeze(0).expand(n_systems, -1).contiguous()
         if self.pbc.shape != (n_systems, 3):
             raise ValueError(
                 f"pbc must have shape (n_systems={n_systems}, 3), "

From a5137ae1335a7a489be5b9c89705e24f2ab19988 Mon Sep 17 00:00:00 2001
From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com>
Date: Sat, 25 Jul 2026 18:37:13 +0900
Subject: [PATCH 3/5] fix: adapt example and model wrappers to per-system pbc

---
 examples/scripts/7_others.py         |  7 ++----
 tests/test_state.py                  | 32 ++++++++++++++++++++--------
 torch_sim/__init__.py                |  2 --
 torch_sim/models/metatomic.py        | 23 +++++++++++++++++++-
 torch_sim/models/nequip_framework.py | 14 ++++++++++++
 torch_sim/models/orb.py              |  9 ++++++--
 torch_sim/state.py                   | 23 +++++++++++---------
 7 files changed, 81 insertions(+), 29 deletions(-)

diff --git a/examples/scripts/7_others.py b/examples/scripts/7_others.py
index 18cc629e7..3e1a25391 100644
--- a/examples/scripts/7_others.py
+++ b/examples/scripts/7_others.py
@@ -56,9 +56,6 @@
 cutoff = 4.0
 self_interaction = False
 
-# Ensure pbc has the correct shape [n_systems, 3]
-pbc_tensor = torch.tensor(pbc).repeat(state.n_systems, 1)
-
 log.info(f"Batched system with {state.n_systems} structures:")
 for i, atoms in enumerate(atoms_list):
     log.info(f"  Structure {i}: {atoms.get_chemical_formula()} ({len(atoms)} atoms)")
@@ -66,7 +63,7 @@
 # Method 1: Linked cell neighbor list (efficient for large systems)
 log.info("Calculating neighbor lists with linked cell method...")
 mapping, mapping_system, shifts_idx = torch_nl_linked_cell(
-    pos, cell, pbc_tensor, cutoff, system_idx, self_interaction
+    pos, cell, pbc, cutoff, system_idx, self_interaction
 )
 cell_shifts = transforms.compute_cell_shifts(cell, shifts_idx, mapping_system)
 dds = transforms.compute_distances_with_cell_shifts(pos, mapping, cell_shifts)
@@ -82,7 +79,7 @@
 # Method 2: N^2 neighbor list (simple but slower)
 log.info("Calculating neighbor lists with N^2 method...")
 mapping_n2, mapping_system_n2, shifts_idx_n2 = torch_nl_n2(
-    pos, cell, pbc_tensor, cutoff, system_idx, self_interaction
+    pos, cell, pbc, cutoff, system_idx, self_interaction
 )
 cell_shifts_n2 = transforms.compute_cell_shifts(cell, shifts_idx_n2, mapping_system_n2)
 dds_n2 = transforms.compute_distances_with_cell_shifts(pos, mapping_n2, cell_shifts_n2)
diff --git a/tests/test_state.py b/tests/test_state.py
index 2e68b923a..044a3eb37 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -21,6 +21,7 @@
     _normalize_system_indices,
     _pop_states,
     _slice_state,
+    _to_legacy_pbc_state,
     coerce_prng,
     get_attrs_for_scope,
 )
@@ -88,22 +89,35 @@ def test_pbc_broadcast_and_assignment(si_double_sim_state: SimState) -> None:
         )
 
 
-def test_require_pbc_helpers(
+def test_require_full_pbc(
     ar_supercell_sim_state: SimState, benzene_sim_state: SimState
 ) -> None:
-    """require_uniform_pbc and require_full_pbc validate per-system pbc."""
-    uniform = ar_supercell_sim_state
-    row = ts.require_uniform_pbc(uniform, "test feature")
-    assert torch.equal(row, torch.tensor([True, True, True], device=DEVICE))
-    ts.require_full_pbc(uniform, "test feature")
+    """require_full_pbc validates per-system pbc."""
+    ts.require_full_pbc(ar_supercell_sim_state, "test feature")
 
-    mixed = ts.concatenate_states([uniform, benzene_sim_state])
-    with pytest.raises(ValueError, match="mixed pbc"):
-        ts.require_uniform_pbc(mixed, "test feature")
+    mixed = ts.concatenate_states([ar_supercell_sim_state, benzene_sim_state])
     with pytest.raises(ValueError, match="fully periodic"):
         ts.require_full_pbc(mixed, "test feature")
 
 
+def test_to_legacy_pbc_state(
+    si_double_sim_state: SimState, benzene_sim_state: SimState
+) -> None:
+    """_to_legacy_pbc_state shares tensors and exposes the shared (3,) pbc row."""
+    state = SimState.from_state(si_double_sim_state, pbc=[True, False, True])
+    assert state.pbc.shape == (2, 3)
+
+    view = _to_legacy_pbc_state(state)
+    assert view.pbc.shape == (3,)
+    assert torch.equal(view.pbc, torch.tensor([True, False, True], device=DEVICE))
+    assert view.positions is state.positions
+    assert state.pbc.shape == (2, 3)
+
+    mixed = ts.concatenate_states([state, benzene_sim_state])
+    with pytest.raises(ValueError, match="mixed pbc"):
+        _to_legacy_pbc_state(mixed)
+
+
 def test_wrap_positions_mixed_pbc_batch(
     ar_supercell_sim_state: SimState, benzene_sim_state: SimState
 ) -> None:
diff --git a/torch_sim/__init__.py b/torch_sim/__init__.py
index e1c2a146c..4b5c8d5d1 100644
--- a/torch_sim/__init__.py
+++ b/torch_sim/__init__.py
@@ -99,7 +99,6 @@
     detach_state_graph,
     initialize_state,
     require_full_pbc,
-    require_uniform_pbc,
 )
 from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter
 
@@ -203,7 +202,6 @@
     "optimizers",
     "quantities",
     "require_full_pbc",
-    "require_uniform_pbc",
     "runners",
     "state",
     "static",
diff --git a/torch_sim/models/metatomic.py b/torch_sim/models/metatomic.py
index f39e9e852..db4d13e4d 100644
--- a/torch_sim/models/metatomic.py
+++ b/torch_sim/models/metatomic.py
@@ -9,7 +9,28 @@
 
 
 try:
-    from metatomic_torchsim import MetatomicModel  # pyright: ignore[reportMissingImports]
+    from metatomic_torchsim import (  # pyright: ignore[reportMissingImports]
+        MetatomicModel as _MetatomicTorchSimModel,
+    )
+
+    from torch_sim.state import SimState, _to_legacy_pbc_state
+
+    class MetatomicModel(_MetatomicTorchSimModel):
+        """Metatomic model wrapper for torch-sim."""
+
+        def forward(self, *args: Any, **kwargs: Any) -> Any:
+            """Run forward pass with the legacy global pbc row.
+
+            The metatomic integration reads state.pbc as a single global
+            setting, so present the shared (3,) row until it supports
+            per-system pbc.
+            """
+            if args and isinstance(args[0], SimState):
+                args = (_to_legacy_pbc_state(args[0]), *args[1:])
+            elif isinstance(kwargs.get("state"), SimState):
+                kwargs["state"] = _to_legacy_pbc_state(kwargs["state"])
+            return super().forward(*args, **kwargs)
+
 except ImportError as exc:
     warnings.warn(
         f"metatomic-torchsim import failed: {traceback.format_exc()}", stacklevel=2
diff --git a/torch_sim/models/nequip_framework.py b/torch_sim/models/nequip_framework.py
index cdb846bda..5e55dc736 100644
--- a/torch_sim/models/nequip_framework.py
+++ b/torch_sim/models/nequip_framework.py
@@ -15,6 +15,8 @@
 try:
     from nequip.integrations.torchsim import NequIPTorchSimCalc
 
+    from torch_sim.state import SimState, _to_legacy_pbc_state
+
     # Re-export with backward-compatible name
     class NequIPFrameworkModel(NequIPTorchSimCalc):
         """NequIP model framework wrapper for torch-sim.
@@ -24,6 +26,18 @@ class NequIPFrameworkModel(NequIPTorchSimCalc):
         model will cast to the correct dtype internally.
         """
 
+        def forward(self, *args: Any, **kwargs: Any) -> Any:
+            """Run forward pass with the legacy global pbc row.
+
+            The NequIP integration reads state.pbc as a single global setting,
+            so present the shared (3,) row until it supports per-system pbc.
+            """
+            if args and isinstance(args[0], SimState):
+                args = (_to_legacy_pbc_state(args[0]), *args[1:])
+            elif isinstance(kwargs.get("state"), SimState):
+                kwargs["state"] = _to_legacy_pbc_state(kwargs["state"])
+            return super().forward(*args, **kwargs)
+
 except ImportError as exc:
     _nequip_import_error = exc  # capture before except block ends (exc is deleted)
     warnings.warn(f"NequIP import failed: {traceback.format_exc()}", stacklevel=2)
diff --git a/torch_sim/models/orb.py b/torch_sim/models/orb.py
index e8736fe02..3b9eb5c9e 100644
--- a/torch_sim/models/orb.py
+++ b/torch_sim/models/orb.py
@@ -20,6 +20,7 @@
 
     import torch_sim as ts
     from torch_sim.elastic import voigt_6_to_full_3x3_stress
+    from torch_sim.state import _to_legacy_pbc_state
 
     # Re-export with backward-compatible name
     class OrbModel(OrbTorchSimModel):
@@ -83,10 +84,14 @@ def _get_results(self, out: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
 
         def forward(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
             """Run forward pass, detaching outputs unless retain_graph is True."""
+            # orb-models reads state.pbc as a single global setting, so present
+            # the shared (3,) row until it supports per-system pbc.
             if args and isinstance(args[0], ts.SimState):
-                args = (self._normalize_charge_spin(args[0]), *args[1:])
+                state = self._normalize_charge_spin(args[0])
+                args = (_to_legacy_pbc_state(state), *args[1:])
             elif isinstance(kwargs.get("state"), ts.SimState):
-                kwargs["state"] = self._normalize_charge_spin(kwargs["state"])
+                state = self._normalize_charge_spin(kwargs["state"])
+                kwargs["state"] = _to_legacy_pbc_state(state)
             output = super().forward(*args, **kwargs)
             return {  # detach tensors as energy is not detached by default
                 k: v.detach() if hasattr(v, "detach") else v for k, v in output.items()
diff --git a/torch_sim/state.py b/torch_sim/state.py
index 5805fee39..6dfe3e0c4 100644
--- a/torch_sim/state.py
+++ b/torch_sim/state.py
@@ -916,29 +916,32 @@ def deform_grad(self) -> torch.Tensor:
         return self._deform_grad(self.reference_row_vector_cell, self.row_vector_cell)
 
 
-def require_uniform_pbc(state: SimState, feature: str) -> torch.Tensor:
-    """Return the pbc row shared by all systems, raising if systems differ.
+def _to_legacy_pbc_state(state: SimState) -> SimState:
+    """Return a shallow copy of state with pbc in the legacy (3,) shape.
 
-    For consumers that cannot handle per-system pbc, e.g. model featurizers
-    that treat pbc as a single global setting.
+    Converts pbc back to the single global row that predates per-system pbc,
+    for external model integrations that still read state.pbc that way.
+    Remove once those integrations migrate.
 
     Args:
-        state (SimState): The state to validate.
-        feature (str): Name of the feature used in the error message.
+        state (SimState): The state to convert.
 
     Returns:
-        torch.Tensor: The shared pbc row with shape (3,).
+        SimState: A copy sharing all tensors with state, with pbc of shape (3,).
 
     Raises:
-        ValueError: If systems have different pbc.
+        ValueError: If systems have different pbc, which the legacy shape
+            cannot represent.
     """
     pbc = state.pbc
     if (pbc != pbc[0]).any():
         raise ValueError(
-            f"{feature} does not support mixed pbc across systems; "
+            "This model does not support mixed pbc across systems; "
             f"got per-system pbc {pbc.tolist()}. Split the batch by pbc instead."
         )
-    return pbc[0]
+    view = copy.copy(state)
+    object.__setattr__(view, "pbc", pbc[0])
+    return view
 
 
 def require_full_pbc(state: SimState, feature: str) -> None:

From 65872a1be309d01ca45a4c5c975c09bc95f40b38 Mon Sep 17 00:00:00 2001
From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com>
Date: Sat, 25 Jul 2026 19:55:57 +0900
Subject: [PATCH 4/5] refactor: drop dead non-tensor pbc fallbacks now that pbc
 is always a (n_systems, 3) tensor

---
 torch_sim/io.py         | 10 ++--------
 torch_sim/trajectory.py |  8 +-------
 2 files changed, 3 insertions(+), 15 deletions(-)

diff --git a/torch_sim/io.py b/torch_sim/io.py
index 3f504e933..551fbe097 100644
--- a/torch_sim/io.py
+++ b/torch_sim/io.py
@@ -78,11 +78,7 @@ def state_to_atoms(
         if state.system_idx is not None
         else np.zeros(state.positions.shape[0], dtype=np.int64)
     )
-    pbc_np = (
-        state.pbc.detach().cpu().numpy()
-        if torch.is_tensor(state.pbc)
-        else np.array([state.pbc] * 3 if isinstance(state.pbc, bool) else state.pbc)
-    )
+    pbc_np = state.pbc.detach().cpu().numpy()  # Shape: (n_systems, 3)
 
     atoms_list = []
     for sys_idx in np.unique(system_indices):
@@ -94,9 +90,7 @@ def state_to_atoms(
         # Convert atomic numbers to chemical symbols
         symbols = [chemical_symbols[z] for z in system_numbers]
 
-        pbc_for_sys = (
-            tuple(pbc_np[sys_idx].tolist()) if pbc_np.ndim > 1 else tuple(pbc_np.tolist())
-        )
+        pbc_for_sys = tuple(pbc_np[sys_idx].tolist())
         atoms = Atoms(
             symbols=symbols, positions=system_positions, cell=system_cell, pbc=pbc_for_sys
         )
diff --git a/torch_sim/trajectory.py b/torch_sim/trajectory.py
index 279298933..403fbcafc 100644
--- a/torch_sim/trajectory.py
+++ b/torch_sim/trajectory.py
@@ -969,14 +969,8 @@ def write_state(  # noqa: C901
             self.write_arrays({"atomic_numbers": state[0].atomic_numbers}, 0)
 
         if "pbc" not in self.array_registry:
-            pbc_val = sub_states[0].pbc
-            pbc_arr = (
-                pbc_val
-                if torch.is_tensor(pbc_val)
-                else torch.tensor([pbc_val] * 3 if isinstance(pbc_val, bool) else pbc_val)
-            )
             # Trajectory files hold a single system, so store its (3,) pbc row
-            self.write_global_array("pbc", pbc_arr.reshape(-1))
+            self.write_global_array("pbc", sub_states[0].pbc.reshape(-1))
 
         # Write all arrays to file
         self.write_arrays(data, steps)

From dc3b6feedcf4bf8846565c1b5153ffbadfa78b27 Mon Sep 17 00:00:00 2001
From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com>
Date: Sat, 25 Jul 2026 20:24:19 +0900
Subject: [PATCH 5/5] fix: present legacy (3,) pbc row to fairchem integration

---
 torch_sim/models/fairchem.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/torch_sim/models/fairchem.py b/torch_sim/models/fairchem.py
index 7eb341dfc..1929c54bb 100644
--- a/torch_sim/models/fairchem.py
+++ b/torch_sim/models/fairchem.py
@@ -13,7 +13,26 @@
 
 
 try:
-    from fairchem.core.calculate.torchsim_interface import FairChemModel
+    from fairchem.core.calculate.torchsim_interface import (
+        FairChemModel as _FairChemTorchSimModel,
+    )
+
+    from torch_sim.state import SimState, _to_legacy_pbc_state
+
+    class FairChemModel(_FairChemTorchSimModel):
+        """FairChem model wrapper for torch-sim."""
+
+        def forward(self, *args: Any, **kwargs: Any) -> Any:
+            """Run forward pass with the legacy global pbc row.
+
+            The FairChem integration reads state.pbc as a single global setting,
+            so present the shared (3,) row until it supports per-system pbc.
+            """
+            if args and isinstance(args[0], SimState):
+                args = (_to_legacy_pbc_state(args[0]), *args[1:])
+            elif isinstance(kwargs.get("state"), SimState):
+                kwargs["state"] = _to_legacy_pbc_state(kwargs["state"])
+            return super().forward(*args, **kwargs)
 
 except ImportError as exc:
     warnings.warn(f"FairChem import failed: {traceback.format_exc()}", stacklevel=2)