feat(state): make pbc a per-system attribute

examples/scripts/7_others.py

@@ -56,17 +56,14 @@
 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)")
 
 # 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)

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"),
     [

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:

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)

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

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"),
     [

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,
 )
@@ -37,9 +38,96 @@ 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_full_pbc(
+    ar_supercell_sim_state: SimState, benzene_sim_state: SimState
+) -> None:
+    """require_full_pbc validates per-system pbc."""
+    ts.require_full_pbc(ar_supercell_sim_state, "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:
+    """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:

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:

torch_sim/__init__.py

@@ -98,6 +98,7 @@
     concatenate_states,
     detach_state_graph,
     initialize_state,
+    require_full_pbc,
 )
 from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter
 
@@ -200,6 +201,7 @@
     "optimize",
     "optimizers",
     "quantities",
+    "require_full_pbc",
     "runners",
     "state",
     "static",