cleanup state

Author: curtischong

examples/scripts/7_Others/7.3_Batched_neighbor_list.py

@@ -18,6 +18,7 @@
 cutoff = torch.tensor(4.0, dtype=pos.dtype)
 self_interaction = False
 
+# Ensure pbc has the correct shape [n_systems, 3]
 pbc_tensor = torch.tensor(pbc).repeat(state.n_systems, 1)
 
 mapping, mapping_system, shifts_idx = torch_nl_linked_cell(

torch_sim/state.py

@@ -109,16 +109,6 @@ def pbc(self) -> torch.Tensor:
 
     def __post_init__(self) -> None:
         """Initialize the SimState and validate the arguments."""
-        if isinstance(self.pbc, bool):
-            self.pbc = [self.pbc] * 3
-        if not isinstance(self.pbc, torch.Tensor):
-            self.pbc = torch.tensor(
-                self.pbc, dtype=torch.bool, device=self.positions.device
-            )
-
-        # Validate and process the state after initialization.
-        # data validation and fill system_idx
-        # should make pbc a tensor here
         # if devices aren't all the same, raise an error, in a clean way
         devices = {
             attr: getattr(self, attr).device
@@ -139,6 +129,13 @@ def __post_init__(self) -> None:
                 f"masses {shapes[1]}, atomic_numbers {shapes[2]}"
             )
 
+        if isinstance(self.pbc, bool):
+            self.pbc = [self.pbc] * 3
+        if not isinstance(self.pbc, torch.Tensor):
+            self.pbc = torch.tensor(
+                self.pbc, dtype=torch.bool, device=self.positions.device
+            )
+
         initial_system_idx = self.system_idx
         if initial_system_idx is None:
             self.system_idx = torch.zeros(