use hardcoded attr names

add test for the new pre-defined attributes

wip make fire simstate attributes predefined

rename features to attributes

define attribute scope for fire state

consolidate attribute definitions and add it to npt

rm infer_property_scope import

__init_subclass__ to enforce that all attributes are specified

find callable attributes

do not check some user attributes

filter for @properties

__init_subclass__ doesn't catch for static state since it's a grandchild state

added documentation for get_attrs_for_scope

more examples for the other scopes

cleaner documentation

rename batch to system

more docs cleanup

test for running the init subclass

more tests

define more scope for the integrators

split state handles none properties a bit better

Author: curtischong

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py

@@ -76,6 +76,7 @@ class HybridSwapMCState(MDState):
     """
 
     last_permutation: torch.Tensor
+    _system_attributes = (*MDState._system_attributes, "last_permutation")  # noqa: SLF001
 
 
 nvt_init, nvt_step = nvt_langevin(model=model, dt=0.002, kT=kT)

examples/tutorials/state_tutorial.py

@@ -71,24 +71,36 @@
 
 # %% [markdown]
 """
-SimState attributes fall into three categories: atomwise, batchwise, and global.
+SimState attributes fall into three categories: atomwise, systemwise, and global.
 
 * Atomwise attributes are tensors with shape (n_atoms, ...), these are `positions`,
-  `masses`, `atomic_numbers`, and `batch`. Names are plural.
-* Batchwise attributes are tensors with shape (n_systems, ...), this is just `cell` for
+  `masses`, `atomic_numbers`, and `system_idx`. Names are plural.
+* Systemwise attributes are tensors with shape (n_systems, ...), this is just `cell` for
   the base SimState. Names are singular.
 * Global attributes have any other shape or type, just `pbc` here. Names are singular.
 
-You can use the `infer_property_scope` function to analyze a state's properties. This
+For TorchSim to know which attributes are atomwise, systemwise, and global, each attribute's
+name is explicitly defined in the `_atom_attributes`, `_system_attributes`, and `_global_attributes`:
+
+_atom_attributes = ("positions", "masses", "atomic_numbers", "system_idx")
+_system_attributes = ("cell",)
+_global_attributes = ("pbc",)
+
+You can use the `get_attrs_for_scope` generator function to analyze a state's properties. This
 is mostly used internally but can be useful for debugging.
 """
 
 # %%
-from torch_sim.state import infer_property_scope
+from torch_sim.state import get_attrs_for_scope
 
-scope = infer_property_scope(si_state)
-print(scope)
+# loop through each attribute:
+for attr_name, attr_value in get_attrs_for_scope(si_state, "per-atom"):
+    print(f"per-atom attribute: {attr_name}")
+    print(f"value: {attr_value}")
 
+# or access the attributes via a dict:
+print("Per-system attributes:", dict(get_attrs_for_scope(si_state, "per-system")))  # noqa: E501
+print("Global attributes:", dict(get_attrs_for_scope(si_state, "global")))
 
 # %% [markdown]
 """
@@ -112,7 +124,7 @@
     f"Multi-state has {multi_state.n_atoms} total atoms across {multi_state.n_systems} systems"
 )
 
-# we can see how the shapes of batchwise, atomwise, and global properties change
+# we can see how the shapes of atomwise, systemwise, and global properties change
 print(f"Positions shape: {multi_state.positions.shape}")
 print(f"Cell shape: {multi_state.cell.shape}")
 print(f"PBC: {multi_state.pbc}")
@@ -142,7 +154,7 @@
 
 SimState supports many convenience operations for manipulating batched states. Slicing
 is supported through fancy indexing, e.g. `state[[0, 1, 2]]` will return a new state
-containing only the first three batches. The other operations are available through the
+containing only the first three systems. The other operations are available through the
 `pop`, `split`, `clone`, and `to` methods.
 """
 
@@ -182,19 +194,19 @@
 # %% [markdown]
 """
 
-You can extract specific batches from a batched state using Python's slicing syntax.
+You can extract specific systems from a batched state using Python's slicing syntax.
 This is extremely useful for analyzing specific systems or for implementing complex
 workflows where different systems need separate processing:
 
 The slicing interface follows Python's standard indexing conventions, making it
 intuitive to use. Behind the scenes, TorchSim is creating a new SimState with only the
-selected batches, maintaining all the necessary properties and relationships.
+selected systems, maintaining all the necessary properties and relationships.
 
 Note the difference between these operations:
-- `split()` returns all batches as separate states but doesn't modify the original
-- `pop()` removes specified batches from the original state and returns them as
+- `split()` returns all systems as separate states but doesn't modify the original
+- `pop()` removes specified systems from the original state and returns them as
 separate states
-- `__getitem__` (slicing) creates a new state with specified batches without modifying
+- `__getitem__` (slicing) creates a new state with specified systems without modifying
 the original
 
 This flexibility allows you to structure your simulation workflows in the most
@@ -203,7 +215,7 @@
 ### Splitting and Popping Batches
 
 SimState provides methods to split a batched state into separate states or to remove
-specific batches:
+specific systems:
 """
 
 # %% [markdown]
@@ -257,10 +269,9 @@
 )
 
 print("MDState properties:")
-scope = infer_property_scope(md_state)
-print("Global properties:", scope["global"])
-print("Per-atom properties:", scope["per_atom"])
-print("Per-system properties:", scope["per_system"])
+print("Per-atom attributes:", dict(get_attrs_for_scope(si_state, "per-atom")))
+print("Per-system attributes:", dict(get_attrs_for_scope(si_state, "per-system")))
+print("Global attributes:", dict(get_attrs_for_scope(si_state, "global")))
 
 
 # %% [markdown]

pyproject.toml

@@ -44,6 +44,7 @@ test = [
     "pymatgen>=2024.11.3",
     "pytest-cov>=6",
     "pytest>=8",
+    "pytest-xdist>=3.8.0",
 ]
 io = ["ase>=3.24", "phonopy>=2.37.0", "pymatgen>=2024.11.3"]
 mace = ["mace-torch>=0.3.12"]

tests/test_state.py

@@ -13,7 +13,7 @@
     _pop_states,
     _slice_state,
     concatenate_states,
-    infer_property_scope,
+    get_attrs_for_scope,
     initialize_state,
 )
 
@@ -24,38 +24,59 @@
     from pymatgen.core import Structure
 
 
-def test_infer_sim_state_property_scope(si_sim_state: ts.SimState) -> None:
-    """Test inference of property scope."""
-    scope = infer_property_scope(si_sim_state)
-    assert set(scope["global"]) == {"pbc"}
-    assert set(scope["per_atom"]) == {
+def test_get_attrs_for_scope(si_sim_state: ts.SimState) -> None:
+    """Test getting attributes for a scope."""
+    per_atom_attrs = dict(get_attrs_for_scope(si_sim_state, "per-atom"))
+    assert set(per_atom_attrs.keys()) == {
         "positions",
         "masses",
         "atomic_numbers",
         "system_idx",
     }
-    assert set(scope["per_system"]) == {"cell"}
+    per_system_attrs = dict(get_attrs_for_scope(si_sim_state, "per-system"))
+    assert set(per_system_attrs.keys()) == {"cell"}
+    global_attrs = dict(get_attrs_for_scope(si_sim_state, "global"))
+    assert set(global_attrs.keys()) == {"pbc"}
 
 
-def test_infer_md_state_property_scope(si_sim_state: ts.SimState) -> None:
-    """Test inference of property scope."""
-    state = MDState(
-        **asdict(si_sim_state),
-        momenta=torch.randn_like(si_sim_state.positions),
-        forces=torch.randn_like(si_sim_state.positions),
-        energy=torch.zeros((1,)),
-    )
-    scope = infer_property_scope(state)
-    assert set(scope["global"]) == {"pbc"}
-    assert set(scope["per_atom"]) == {
-        "positions",
-        "masses",
-        "atomic_numbers",
-        "system_idx",
-        "forces",
-        "momenta",
-    }
-    assert set(scope["per_system"]) == {"cell", "energy"}
+def test_all_attributes_must_be_specified_in_scopes() -> None:
+    """Test that an error is raised when we forget to specify the scope
+    for an attribute in a child SimState class."""
+    with pytest.raises(TypeError) as excinfo:
+
+        class ChildState(SimState):
+            attribute_specified_in_scopes: bool
+            attribute_not_specified_in_scopes: bool
+
+            _atom_attributes = (
+                *SimState._atom_attributes,  # noqa: SLF001
+                "attribute_specified_in_scopes",
+            )
+
+    assert "attribute_not_specified_in_scopes" in str(excinfo.value)
+    assert "attribute_specified_in_scopes" not in str(excinfo.value)
+
+
+def test_no_duplicate_attributes_in_scopes() -> None:
+    """Test that no attributes are specified in multiple scopes."""
+
+    # Capture the exception information using "as excinfo"
+    with pytest.raises(TypeError) as excinfo:
+
+        class ChildState(SimState):
+            duplicated_attribute: bool
+
+            _system_attributes = (
+                *SimState._atom_attributes,  # noqa: SLF001
+                "duplicated_attribute",
+            )
+            _global_attributes = (
+                *SimState._global_attributes,  # noqa: SLF001
+                "duplicated_attribute",
+            )
+
+    assert "are declared multiple times" in str(excinfo.value)
+    assert "duplicated_attribute" in str(excinfo.value)
 
 
 def test_slice_substate(

torch_sim/integrators/md.py

@@ -40,6 +40,9 @@ class MDState(SimState):
     energy: torch.Tensor
     forces: torch.Tensor
 
+    _atom_attributes = (*SimState._atom_attributes, "momenta", "forces")  # noqa: SLF001
+    _system_attributes = (*SimState._system_attributes, "energy")  # noqa: SLF001
+
     @property
     def velocities(self) -> torch.Tensor:
         """Velocities calculated from momenta and masses with shape

torch_sim/integrators/npt.py

@@ -14,6 +14,7 @@
     calculate_momenta,
     construct_nose_hoover_chain,
 )
+from torch_sim.optimizers import md_atom_attributes
 from torch_sim.quantities import calc_kinetic_energy
 from torch_sim.state import SimState
 from torch_sim.typing import StateDict
@@ -66,6 +67,17 @@ class NPTLangevinState(SimState):
     cell_velocities: torch.Tensor
     cell_masses: torch.Tensor
 
+    _atom_attributes = md_atom_attributes
+    _system_attributes = (
+        *SimState._system_attributes,  # noqa: SLF001
+        "stress",
+        "cell_positions",
+        "cell_velocities",
+        "cell_masses",
+        "reference_cell",
+        "energy",
+    )
+
     @property
     def momenta(self) -> torch.Tensor:
         """Calculate momenta from velocities and masses."""
@@ -866,6 +878,21 @@ class NPTNoseHooverState(MDState):
     barostat: NoseHooverChain
     barostat_fns: NoseHooverChainFns
 
+    _system_attributes = (
+        *MDState._system_attributes,  # noqa: SLF001
+        "reference_cell",
+        "cell_position",
+        "cell_momentum",
+        "cell_mass",
+    )
+    _global_attributes = (
+        *MDState._global_attributes,  # noqa: SLF001
+        "thermostat",
+        "barostat",
+        "thermostat_fns",
+        "barostat_fns",
+    )
+
     @property
     def velocities(self) -> torch.Tensor:
         """Calculate particle velocities from momenta and masses.

torch_sim/integrators/nvt.py

@@ -265,6 +265,12 @@ class NVTNoseHooverState(MDState):
     chain: NoseHooverChain
     _chain_fns: NoseHooverChainFns
 
+    _global_attributes = (
+        *MDState._global_attributes,  # noqa: SLF001
+        "chain",
+        "_chain_fns",
+    )
+
     @property
     def velocities(self) -> torch.Tensor:
         """Velocities calculated from momenta and masses with shape

torch_sim/monte_carlo.py

@@ -35,6 +35,9 @@ class SwapMCState(SimState):
     energy: torch.Tensor
     last_permutation: torch.Tensor
 
+    _atom_attributes = (*SimState._atom_attributes, "last_permutation")  # noqa: SLF001
+    _system_attributes = (*SimState._system_attributes, "energy")  # noqa: SLF001
+
 
 def generate_swaps(
     state: SimState, generator: torch.Generator | None = None

torch_sim/optimizers.py

@@ -32,6 +32,28 @@
 MdFlavor = Literal["vv_fire", "ase_fire"]
 vv_fire_key, ase_fire_key = get_args(MdFlavor)
 
+md_atom_attributes = (*SimState._atom_attributes, "forces", "velocities")  # noqa: SLF001
+_fire_system_attributes = (
+    *SimState._system_attributes,  # noqa: SLF001
+    "energy",
+    "stress",
+    "cell_positions",
+    "cell_velocities",
+    "cell_forces",
+    "cell_masses",
+    "reference_cell",
+    "cell_factor",
+    "pressure",
+    "dt",
+    "alpha",
+    "n_pos",
+)
+_fire_global_attributes = (
+    *SimState._global_attributes,  # noqa: SLF001
+    "hydrostatic_strain",
+    "constant_volume",
+)
+
 
 @dataclass
 class GDState(SimState):
@@ -55,6 +77,9 @@ class GDState(SimState):
     forces: torch.Tensor
     energy: torch.Tensor
 
+    _atom_attributes = (*SimState._atom_attributes, "forces")  # noqa: SLF001
+    _system_attributes = (*SimState._system_attributes, "energy")  # noqa: SLF001
+
 
 def gradient_descent(
     model: torch.nn.Module, *, lr: torch.Tensor | float = 0.01
@@ -194,6 +219,22 @@ class UnitCellGDState(GDState, DeformGradMixin):
     cell_forces: torch.Tensor
     cell_masses: torch.Tensor
 
+    _system_attributes = (
+        *GDState._system_attributes,  # noqa: SLF001
+        "cell_forces",
+        "pressure",
+        "stress",
+        "cell_positions",
+        "cell_factor",
+        "cell_masses",
+    )
+    _global_attributes = (
+        *GDState._global_attributes,  # noqa: SLF001
+        "reference_cell",
+        "hydrostatic_strain",
+        "constant_volume",
+    )
+
 
 def unit_cell_gradient_descent(  # noqa: PLR0915, C901
     model: torch.nn.Module,
@@ -481,6 +522,9 @@ class FireState(SimState):
     alpha: torch.Tensor
     n_pos: torch.Tensor
 
+    _atom_attributes = md_atom_attributes
+    _system_attributes = (*SimState._system_attributes, "energy", "dt", "alpha", "n_pos")  # noqa: SLF001
+
 
 def fire(
     model: torch.nn.Module,
@@ -692,6 +736,10 @@ class UnitCellFireState(SimState, DeformGradMixin):
     alpha: torch.Tensor
     n_pos: torch.Tensor
 
+    _atom_attributes = md_atom_attributes
+    _system_attributes = _fire_system_attributes
+    _global_attributes = _fire_global_attributes
+
 
 def unit_cell_fire(
     model: torch.nn.Module,
@@ -980,6 +1028,10 @@ class FrechetCellFIREState(SimState, DeformGradMixin):
     alpha: torch.Tensor
     n_pos: torch.Tensor
 
+    _atom_attributes = md_atom_attributes
+    _system_attributes = _fire_system_attributes
+    _global_attributes = _fire_global_attributes
+
 
 def frechet_cell_fire(
     model: torch.nn.Module,