Make system_idx non-optional in SimState [1/2] (#231)

Author: curtischong

torch_sim/integrators/nvt.py

@@ -389,6 +389,7 @@ def nvt_nose_hoover_init(
             cell=state.cell,
             pbc=state.pbc,
             atomic_numbers=atomic_numbers,
+            system_idx=state.system_idx,
             chain=chain_fns.initialize(total_dof, KE, kT),
             _chain_fns=chain_fns,  # Store the chain functions
         )

torch_sim/state.py

@@ -7,7 +7,7 @@
 import copy
 import importlib
 import warnings
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Literal, Self
 
 import torch
@@ -22,7 +22,7 @@
     from pymatgen.core import Structure
 
 
-@dataclass
+@dataclass(init=False)
 class SimState:
     """State representation for atomistic systems with batched operations support.
 
@@ -47,9 +47,8 @@ class SimState:
             used by ASE.
         pbc (bool): Boolean indicating whether to use periodic boundary conditions
         atomic_numbers (torch.Tensor): Atomic numbers with shape (n_atoms,)
-        system_idx (torch.Tensor, optional): Maps each atom index to its system index.
-            Has shape (n_atoms,), defaults to None, must be unique consecutive
-            integers starting from 0
+        system_idx (torch.Tensor): Maps each atom index to its system index.
+            Has shape (n_atoms,), must be unique consecutive integers starting from 0.
 
     Properties:
         wrap_positions (torch.Tensor): Positions wrapped according to periodic boundary
@@ -81,10 +80,35 @@ class SimState:
     cell: torch.Tensor
     pbc: bool  # TODO: do all calculators support mixed pbc?
     atomic_numbers: torch.Tensor
-    system_idx: torch.Tensor | None = field(default=None, kw_only=True)
+    system_idx: torch.Tensor
+
+    def __init__(
+        self,
+        positions: torch.Tensor,
+        masses: torch.Tensor,
+        cell: torch.Tensor,
+        pbc: bool,  # noqa: FBT001 # TODO(curtis): maybe make the constructor be keyword-only (it can be easy to confuse positions vs masses, etc.)
+        atomic_numbers: torch.Tensor,
+        system_idx: torch.Tensor | None = None,
+    ) -> None:
+        """Initialize the SimState and validate the arguments.
+
+        Args:
+            positions (torch.Tensor): Atomic positions with shape (n_atoms, 3)
+            masses (torch.Tensor): Atomic masses with shape (n_atoms,)
+            cell (torch.Tensor): Unit cell vectors with shape (n_systems, 3, 3).
+            pbc (bool): Boolean indicating whether to use periodic boundary conditions
+            atomic_numbers (torch.Tensor): Atomic numbers with shape (n_atoms,)
+            system_idx (torch.Tensor | None): Maps each atom index to its system index.
+                Has shape (n_atoms,), must be unique consecutive integers starting from 0.
+                If not provided, it is initialized to zeros.
+        """
+        self.positions = positions
+        self.masses = masses
+        self.cell = cell
+        self.pbc = pbc
+        self.atomic_numbers = atomic_numbers
 
-    def __post_init__(self) -> None:
-        """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
@@ -107,24 +131,25 @@ def __post_init__(self) -> None:
                 f"masses {shapes[1]}, atomic_numbers {shapes[2]}"
             )
 
-        if self.cell.ndim != 3 and self.system_idx is None:
-            self.cell = self.cell.unsqueeze(0)
-
-        if self.cell.shape[-2:] != (3, 3):
-            raise ValueError("Cell must have shape (n_systems, 3, 3)")
-
-        if self.system_idx is None:
+        if system_idx is None:
             self.system_idx = torch.zeros(
                 self.n_atoms, device=self.device, dtype=torch.int64
             )
         else:
+            self.system_idx = system_idx
             # assert that system indices are unique consecutive integers
             # TODO(curtis): I feel like this logic is not reliable.
             # I'll come up with something better later.
             _, counts = torch.unique_consecutive(self.system_idx, return_counts=True)
             if not torch.all(counts == torch.bincount(self.system_idx)):
                 raise ValueError("System indices must be unique consecutive integers")
 
+        if self.cell.ndim != 3 and system_idx is None:
+            self.cell = self.cell.unsqueeze(0)
+
+        if self.cell.shape[-2:] != (3, 3):
+            raise ValueError("Cell must have shape (n_systems, 3, 3)")
+
         if self.cell.shape[0] != self.n_systems:
             raise ValueError(
                 f"Cell must have shape (n_systems, 3, 3), got {self.cell.shape}"