trajectory runs

Author: curtischong

tests/test_trajectory.py

@@ -439,7 +439,9 @@ def test_get_atoms(trajectory: TorchSimTrajectory, random_state: MDState) -> Non
     np.testing.assert_allclose(
         atoms.get_atomic_numbers(), random_state.atomic_numbers.numpy()
     )
-    assert atoms.pbc.all() == random_state.pbc
+    np.testing.assert_array_equal(
+        atoms.pbc, random_state.pbc.detach().cpu().numpy()
+    )
 
 
 def test_get_state(trajectory: TorchSimTrajectory, random_state: MDState) -> None:
@@ -509,7 +511,9 @@ def test_write_ase_trajectory(
         np.testing.assert_allclose(
             atoms.get_atomic_numbers(), random_state.atomic_numbers.numpy()
         )
-        assert atoms.pbc.all() == random_state.pbc
+        np.testing.assert_array_equal(
+            atoms.pbc, random_state.pbc.detach().cpu().numpy()
+        )
 
     # Clean up
     ase_traj.close()

torch_sim/integrators/md.py

@@ -48,6 +48,17 @@ class MDState(SimState):
         SimState._system_attributes | {"energy"}  # noqa: SLF001
     )
 
+    def __post_init__(self) -> None:
+        """Ensure SimState initialization logic runs for MDState."""
+        super().__init__(
+            positions=self.positions,
+            masses=self.masses,
+            cell=self.cell,
+            pbc=self.pbc,
+            atomic_numbers=self.atomic_numbers,
+            system_idx=self.system_idx,
+        )
+
     @property
     def velocities(self) -> torch.Tensor:
         """Velocities calculated from momenta and masses with shape

torch_sim/trajectory.py

@@ -742,6 +742,10 @@ def write_state(  # noqa: C901
 
         if len(sub_states) != len(steps):
             raise ValueError(f"{len(sub_states)=} must match the {len(steps)=}")
+
+        # Use the selected states for data serialization
+        state = sub_states
+
         # Initialize data dictionary with required arrays
         data = {
             "positions": torch.stack([s.positions for s in state]),
@@ -781,8 +785,7 @@ def write_state(  # noqa: C901
             # Save atomic numbers only for first frame
             self.write_arrays({"atomic_numbers": state[0].atomic_numbers}, 0)
 
-        if "pbc" not in self.array_registry:
-            self.write_arrays({"pbc": state[0].pbc}, [0])
+        data["pbc"] = torch.stack([s.pbc.reshape(-1) for s in state])
 
         # Write all arrays to file
         self.write_arrays(data, steps)
@@ -833,7 +836,7 @@ def return_prop(self: Self, prop: str, frame: int) -> np.ndarray:
         arrays["cell"] = np.expand_dims(return_prop(self, "cell", frame), axis=0)
         arrays["atomic_numbers"] = return_prop(self, "atomic_numbers", frame)
         arrays["masses"] = return_prop(self, "masses", frame)
-        arrays["pbc"] = np.expand_dims(return_prop(self, "pbc", frame), axis=0)
+        arrays["pbc"] = return_prop(self, "pbc", frame)
 
         return arrays
 
@@ -897,7 +900,7 @@ def get_atoms(self, frame: int = -1) -> "Atoms":
             numbers=np.ascontiguousarray(arrays["atomic_numbers"]),
             positions=np.ascontiguousarray(arrays["positions"]),
             cell=np.ascontiguousarray(arrays["cell"])[0],
-            pbc=np.ascontiguousarray(arrays["pbc"])[0],
+            pbc=np.ascontiguousarray(arrays["pbc"]),
         )
 
     def get_state(