Fix types so pylance's type checker doesn't complain

.gitignore

@@ -34,3 +34,5 @@ coverage.xml
 
 # env
 uv.lock
+
+.vscode/

examples/tutorials/state_tutorial.py

@@ -112,7 +112,7 @@
     f"Multi-state has {multi_state.n_atoms} total atoms across {multi_state.n_batches} batches"
 )
 
-# we can see how the shapes of batchwise, atomwise, and global properties change
+# we can see how the shapes of batchwise, atomwise, and per-graph properties change
 print(f"Positions shape: {multi_state.positions.shape}")
 print(f"Cell shape: {multi_state.cell.shape}")
 print(f"PBC: {multi_state.pbc}")
@@ -258,7 +258,7 @@
 
 print("MDState properties:")
 scope = infer_property_scope(md_state)
-print("Global properties:", scope["global"])
+print("Per-graph properties:", scope["per_graph"])
 print("Per-atom properties:", scope["per_atom"])
 print("Per-batch properties:", scope["per_batch"])
 

tests/test_state.py

@@ -26,7 +26,7 @@
 def test_infer_sim_state_property_scope(si_sim_state: SimState) -> None:
     """Test inference of property scope."""
     scope = infer_property_scope(si_sim_state)
-    assert set(scope["global"]) == {"pbc"}
+    assert set(scope["per_graph"]) == {"pbc"}
     assert set(scope["per_atom"]) == {"positions", "masses", "atomic_numbers", "batch"}
     assert set(scope["per_batch"]) == {"cell"}
 
@@ -40,7 +40,7 @@ def test_infer_md_state_property_scope(si_sim_state: SimState) -> None:
         energy=torch.zeros((1,)),
     )
     scope = infer_property_scope(state)
-    assert set(scope["global"]) == {"pbc"}
+    assert set(scope["per_graph"]) == {"pbc"}
     assert set(scope["per_atom"]) == {
         "positions",
         "masses",

torch_sim/models/lennard_jones.py

@@ -325,7 +325,7 @@ def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]:
 
         # we always return tensors
         # per atom properties are returned as (atoms, ...) tensors
-        # global properties are returned as shape (..., n) tensors
+        # per-graph properties are returned as shape (..., n) tensors
         results = {}
         for key in ("stress", "energy"):
             if key in properties:

torch_sim/models/morse.py

@@ -308,7 +308,7 @@ def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]:
 
         # we always return tensors
         # per atom properties are returned as (atoms, ...) tensors
-        # global properties are returned as shape (..., n) tensors
+        # per-graph properties are returned as shape (..., n) tensors
         results = {}
         for key in ("stress", "energy"):
             if key in properties:

torch_sim/state.py

@@ -8,7 +8,7 @@
 import importlib
 import warnings
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Literal, Self
+from typing import TYPE_CHECKING, Literal, Self, TypedDict, cast
 
 import torch
 
@@ -152,7 +152,7 @@
         return torch.unique(self.batch).shape[0]
 
     @property
-    def volume(self) -> torch.Tensor:
+    def volume(self) -> torch.Tensor | None:
         """Volume of the system."""
         return torch.det(self.cell) if self.pbc else None
 
@@ -226,7 +226,7 @@
         """
         return ts.io.state_to_phonopy(self)
 
-    def split(self) -> list[Self]:
+    def split(self) -> list["SimState"]:
         """Split the SimState into a list of single-batch SimStates.
 
         Divides the current state into separate states, each containing a single batch,
@@ -237,7 +237,9 @@
         """
         return _split_state(self)
 
-    def pop(self, batch_indices: int | list[int] | slice | torch.Tensor) -> list[Self]:
+    def pop(
+        self, batch_indices: int | list[int] | slice | torch.Tensor
+    ) -> list["SimState"]:
         """Pop off states with the specified batch indices.
 
         This method modifies the original state object by removing the specified
@@ -268,7 +270,7 @@
 
     def to(
         self, device: torch.device | None = None, dtype: torch.dtype | None = None
-    ) -> Self:
+    ) -> "SimState":
         """Convert the SimState to a new device and/or data type.
 
         Args:
@@ -282,7 +284,9 @@
         """
         return state_to_device(self, device, dtype)
 
-    def __getitem__(self, batch_indices: int | list[int] | slice | torch.Tensor) -> Self:
+    def __getitem__(
+        self, batch_indices: int | list[int] | slice | torch.Tensor
+    ) -> "SimState":
         """Enable standard Python indexing syntax for slicing batches.
 
         Args:
@@ -387,7 +391,7 @@
 
 def state_to_device(
     state: SimState, device: torch.device | None = None, dtype: torch.dtype | None = None
-) -> Self:
+) -> SimState:
     """Convert the SimState to a new device and dtype.
 
     Creates a new SimState with all tensors moved to the specified device and
@@ -419,10 +423,18 @@
     return type(state)(**attrs)
 
 
+class Scope(TypedDict):
+    """Dictionary mapping each scope category to a list of property names"""
+
+    per_graph: list[str]
+    per_atom: list[str]
+    per_batch: list[str]
+
+
 def infer_property_scope(
     state: SimState,
     ambiguous_handling: Literal["error", "globalize", "globalize_warn"] = "error",
-) -> dict[Literal["global", "per_atom", "per_batch"], list[str]]:
+) -> Scope:
     """Infer whether a property is global, per-atom, or per-batch.
 
     Analyzes the shapes of tensor attributes to determine their scope within
@@ -437,14 +449,13 @@
             - "globalize_warn": Treat ambiguous properties as global with a warning
 
     Returns:
-        dict[Literal["global", "per_atom", "per_batch"], list[str]]: Dictionary mapping
-            each scope category to a list of property names
+        Scope: Dictionary mapping each scope category to a list of property names
 
     Raises:
         ValueError: If n_atoms equals n_batches (making scope inference ambiguous) or
             if ambiguous_handling="error" and an ambiguous property is encountered
     """
-    # TODO: this cannot effectively resolve global properties with
+    # TODO: this cannot effectively resolve per-graph properties with
     # length of n_atoms or n_batches, they will be classified incorrectly,
     # no clear fix
 
@@ -454,25 +465,25 @@
             "which means shapes cannot be inferred unambiguously."
         )
 
-    scope = {
-        "global": [],
+    scope: Scope = {
+        "per_graph": [],
         "per_atom": [],
         "per_batch": [],
     }
 
     # Iterate through all attributes
     for attr_name, attr_value in vars(state).items():
-        # Handle scalar values (global properties)
+        # Handle scalar values (per-graph properties)
         if not isinstance(attr_value, torch.Tensor):
-            scope["global"].append(attr_name)
+            scope["per_graph"].append(attr_name)
             continue
 
         # Handle tensor properties based on shape
         shape = attr_value.shape
 
         # Empty tensor case
         if len(shape) == 0:
-            scope["global"].append(attr_name)
+            scope["per_graph"].append(attr_name)
         # Vector/matrix with first dimension matching number of atoms
         elif shape[0] == state.n_atoms:
             scope["per_atom"].append(attr_name)
@@ -487,7 +498,7 @@
                 f"{state.n_batches} (per-batch), or a scalar (global)."
             )
         elif ambiguous_handling in ("globalize", "globalize_warn"):
-            scope["global"].append(attr_name)
+            scope["per_graph"].append(attr_name)
 
             if ambiguous_handling == "globalize_warn":
                 warnings.warn(
@@ -519,11 +530,11 @@
     """
     scope = infer_property_scope(state, ambiguous_handling=ambiguous_handling)
 
-    attrs = {"global": {}, "per_atom": {}, "per_batch": {}}
+    attrs: dict[str, dict] = {"per_graph": {}, "per_atom": {}, "per_batch": {}}
 
-    # Process global properties
-    for attr_name in scope["global"]:
-        attrs["global"][attr_name] = getattr(state, attr_name)
+    # Process per-graph properties
+    for attr_name in scope["per_graph"]:
+        attrs["per_graph"][attr_name] = getattr(state, attr_name)
 
     # Process per-atom properties
     for attr_name in scope["per_atom"]:
@@ -537,7 +548,7 @@
 
 
 def _filter_attrs_by_mask(
-    attrs: dict[str, dict],
+    attrs: Scope,
     atom_mask: torch.Tensor,
     batch_mask: torch.Tensor,
 ) -> dict:
@@ -546,7 +557,7 @@
     Selects subsets of attributes based on boolean masks for atoms and batches.
 
     Args:
-        attrs (dict[str, dict]): Dictionary with 'global', 'per_atom', and 'per_batch'
+        attrs (Scope): Dictionary with 'per_graph', 'per_atom', and 'per_batch'
             attributes
         atom_mask (torch.Tensor): Boolean mask for atoms to include with shape
             (n_atoms,)
@@ -559,7 +570,7 @@
     filtered_attrs = {}
 
     # Copy global attributes directly
-    filtered_attrs.update(attrs["global"])
+    filtered_attrs.update(attrs["per_graph"])
 
     # Filter per-atom attributes
     for attr_name, attr_value in attrs["per_atom"].items():
@@ -638,7 +649,7 @@
             # Add the split per-batch attributes
             **{attr_name: split_per_batch[attr_name][i] for attr_name in split_per_batch},
             # Add the global attributes
-            **attrs["global"],
+            **attrs["per_graph"],
         }
         states.append(type(state)(**batch_attrs))
 
@@ -753,7 +764,7 @@
     """Concatenate a list of SimStates into a single SimState.
 
     Combines multiple states into a single state with multiple batches.
-    Global properties are taken from the first state, and per-atom and per-batch
+    per-graph properties are taken from the first state, and per-atom and per-batch
     properties are concatenated.
 
     Args:
@@ -783,14 +794,14 @@
     target_device = device or first_state.device
 
     # Get property scopes from the first state to identify
-    # global/per-atom/per-batch properties
+    # per-graph/per-atom/per-batch properties
     first_scope = infer_property_scope(first_state)
-    global_props = set(first_scope["global"])
+    per_graph_props = set(first_scope["per_graph"])
     per_atom_props = set(first_scope["per_atom"])
     per_batch_props = set(first_scope["per_batch"])
 
-    # Initialize result with global properties from first state
-    concatenated = {prop: getattr(first_state, prop) for prop in global_props}
+    # Initialize result with per-graph properties from first state
+    concatenated = {prop: getattr(first_state, prop) for prop in per_graph_props}
 
     # Pre-allocate lists for tensors to concatenate
     per_atom_tensors = {prop: [] for prop in per_atom_props}
@@ -864,6 +875,7 @@
         return state_to_device(system, device, dtype)
 
     if isinstance(system, list) and all(isinstance(s, SimState) for s in system):
+        system = cast(list[SimState], system)
         if not all(state.n_batches == 1 for state in system):
             raise ValueError(
                 "When providing a list of states, to the initialize_state function, "