add type checking and fix pbc type in dataclass def

Author: curtischong

torch_sim/state.py

@@ -82,7 +82,7 @@ class SimState:
     positions: torch.Tensor
     masses: torch.Tensor
     cell: torch.Tensor
-    pbc: torch.Tensor
+    pbc: torch.Tensor | list[bool] | bool
     atomic_numbers: torch.Tensor
     system_idx: torch.Tensor | None = field(default=None)
 
@@ -93,6 +93,11 @@ def system_idx(self) -> torch.Tensor:
             """A getter for system_idx that tells type checkers it's always defined."""
             return self.system_idx
 
+        @property
+        def pbc(self) -> torch.Tensor:
+            """A getter for pbc that tells type checkers it's always defined."""
+            return self.pbc
+
     _atom_attributes: ClassVar[set[str]] = {
         "positions",
         "masses",