rm classvar annotation from all simstate

Author: curtischong

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py

@@ -8,7 +8,6 @@
 # ///
 
 from dataclasses import dataclass
-from typing import ClassVar
 
 import torch
 from mace.calculators.foundations_models import mace_mp
@@ -77,7 +76,7 @@ class HybridSwapMCState(MDState):
     """
 
     last_permutation: torch.Tensor
-    _atom_attributes: ClassVar[set[str]] = (
+    _atom_attributes = (
         MDState._atom_attributes | {"last_permutation"}  # noqa: SLF001
     )
 

examples/tutorials/hybrid_swap_tutorial.py

@@ -106,7 +106,7 @@ class HybridSwapMCState(ts.integrators.MDState):
     """
 
     last_permutation: torch.Tensor
-    _atom_attributes: ClassVar[set[str]] = (
+    _atom_attributes = (
         MDState._atom_attributes | {"last_permutation"}  # noqa: SLF001
     )
 

tests/test_state.py

@@ -66,7 +66,7 @@ class ChildState(SimState):
             duplicated_attribute: bool
 
             _system_attributes: typing.ClassVar[set[str]] = (
-                SimState._atom_attributes | {"duplicated_attribute"}  # noqa: SLF001
+                SimState._system_attributes | {"duplicated_attribute"}  # noqa: SLF001
             )
             _global_attributes: typing.ClassVar[set[str]] = (
                 SimState._global_attributes | {"duplicated_attribute"}  # noqa: SLF001

torch_sim/integrators/md.py

@@ -2,7 +2,6 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import ClassVar
 
 import torch
 
@@ -42,10 +41,10 @@ class MDState(SimState):
     energy: torch.Tensor
     forces: torch.Tensor
 
-    _atom_attributes: ClassVar[set[str]] = (
+    _atom_attributes = (
         SimState._atom_attributes | {"momenta", "forces"}  # noqa: SLF001
     )
-    _system_attributes: ClassVar[set[str]] = (
+    _system_attributes = (
         SimState._system_attributes | {"energy"}  # noqa: SLF001
     )
 

torch_sim/integrators/npt.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, ClassVar
+from typing import Any
 
 import torch
 
@@ -67,10 +67,10 @@ class NPTLangevinState(SimState):
     cell_velocities: torch.Tensor
     cell_masses: torch.Tensor
 
-    _atom_attributes: ClassVar[set[str]] = (
+    _atom_attributes = (
         SimState._atom_attributes | {"forces", "velocities"}  # noqa: SLF001
     )
-    _system_attributes: ClassVar[set[str]] = SimState._system_attributes | {  # noqa: SLF001
+    _system_attributes = SimState._system_attributes | {  # noqa: SLF001
         "stress",
         "cell_positions",
         "cell_velocities",
@@ -879,7 +879,7 @@ class NPTNoseHooverState(MDState):
     barostat: NoseHooverChain
     barostat_fns: NoseHooverChainFns
 
-    _system_attributes: ClassVar[set[str]] = (
+    _system_attributes = (
         MDState._system_attributes  # noqa: SLF001
         | {
             "reference_cell",
@@ -888,7 +888,7 @@ class NPTNoseHooverState(MDState):
             "cell_mass",
         }
     )
-    _global_attributes: ClassVar[set[str]] = (
+    _global_attributes = (
         MDState._global_attributes  # noqa: SLF001
         | {
             "thermostat",

torch_sim/integrators/nvt.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, ClassVar
+from typing import Any
 
 import torch
 
@@ -266,7 +266,7 @@ class NVTNoseHooverState(MDState):
     chain: NoseHooverChain
     _chain_fns: NoseHooverChainFns
 
-    _global_attributes: ClassVar[set[str]] = (
+    _global_attributes = (
         MDState._global_attributes | {"chain", "_chain_fns"}  # noqa: SLF001
     )
 

torch_sim/monte_carlo.py

@@ -12,7 +12,6 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import ClassVar
 
 import torch
 
@@ -37,12 +36,8 @@ class SwapMCState(SimState):
     energy: torch.Tensor
     last_permutation: torch.Tensor
 
-    _atom_attributes: ClassVar[set[str]] = (
-        SimState._atom_attributes | {"last_permutation"}  # noqa: SLF001
-    )
-    _system_attributes: ClassVar[set[str]] = (
-        SimState._system_attributes | {"energy"}  # noqa: SLF001
-    )
+    _atom_attributes = SimState._atom_attributes | {"last_permutation"}  # noqa: SLF001
+    _system_attributes = SimState._system_attributes | {"energy"}  # noqa: SLF001
 
 
 def generate_swaps(

torch_sim/optimizers.py

@@ -20,7 +20,7 @@
 import functools
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, ClassVar, Literal, get_args
+from typing import Any, Literal, get_args
 
 import torch
 
@@ -221,7 +221,7 @@ class UnitCellGDState(GDState, DeformGradMixin):
     cell_forces: torch.Tensor
     cell_masses: torch.Tensor
 
-    _system_attributes: ClassVar[set[str]] = (
+    _system_attributes = (
         GDState._system_attributes  # noqa: SLF001
         | DeformGradMixin._system_attributes  # noqa: SLF001
         | {
@@ -233,7 +233,7 @@ class UnitCellGDState(GDState, DeformGradMixin):
             "cell_masses",
         }
     )
-    _global_attributes: ClassVar[set[str]] = (
+    _global_attributes = (
         GDState._global_attributes | {"hydrostatic_strain", "constant_volume"}  # noqa: SLF001
     )
 
@@ -524,8 +524,8 @@ class FireState(SimState):
     alpha: torch.Tensor
     n_pos: torch.Tensor
 
-    _atom_attributes: ClassVar[set[str]] = md_atom_attributes
-    _system_attributes: ClassVar[set[str]] = (
+    _atom_attributes = md_atom_attributes
+    _system_attributes = (
         SimState._system_attributes  # noqa: SLF001
         | {
             "energy",
@@ -745,9 +745,9 @@ class UnitCellFireState(SimState, DeformGradMixin):
     alpha: torch.Tensor
     n_pos: torch.Tensor
 
-    _atom_attributes: ClassVar[set[str]] = md_atom_attributes
-    _system_attributes: ClassVar[set[str]] = _fire_system_attributes
-    _global_attributes: ClassVar[set[str]] = _fire_global_attributes
+    _atom_attributes = md_atom_attributes
+    _system_attributes = _fire_system_attributes
+    _global_attributes = _fire_global_attributes
 
 
 def unit_cell_fire(
@@ -1037,9 +1037,9 @@ class FrechetCellFIREState(SimState, DeformGradMixin):
     alpha: torch.Tensor
     n_pos: torch.Tensor
 
-    _atom_attributes: ClassVar[set[str]] = md_atom_attributes
-    _system_attributes: ClassVar[set[str]] = _fire_system_attributes
-    _global_attributes: ClassVar[set[str]] = _fire_global_attributes
+    _atom_attributes = md_atom_attributes
+    _system_attributes = _fire_system_attributes
+    _global_attributes = _fire_global_attributes
 
 
 def frechet_cell_fire(

torch_sim/runners.py

@@ -9,7 +9,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from itertools import chain
-from typing import Any, ClassVar
+from typing import Any
 
 import torch
 from tqdm import tqdm
@@ -541,10 +541,10 @@ class StaticState(type(state)):
         forces: torch.Tensor
         stress: torch.Tensor
 
-        _atom_attributes: ClassVar[set[str]] = (
+        _atom_attributes = (
             state._atom_attributes | {"forces"}  # noqa: SLF001
         )
-        _system_attributes: ClassVar[set[str]] = (
+        _system_attributes = (
             state._system_attributes | {"energy", "stress"}  # noqa: SLF001
         )
 

torch_sim/state.py

@@ -477,13 +477,13 @@ def _assert_all_attributes_have_defined_scope(cls) -> None:
                 vars(cls).get(attr_name), property
             )
             is_method = hasattr(cls, attr_name) and callable(getattr(cls, attr_name))
-            is_scope_list = attr_name in [
-                "_atom_attributes",
-                "_system_attributes",
-                "_global_attributes",
-            ]
+            is_class_variable = (
+                # Note: _atom_attributes, _system_attributes, and _global_attributes
+                # are all class variables
+                typing.get_origin(all_annotations.get(attr_name)) is typing.ClassVar
+            )
 
-            if is_special_attribute or is_property or is_method or is_scope_list:
+            if is_special_attribute or is_property or is_method or is_class_variable:
                 continue
 
             if attr_name not in all_defined_attributes: