Define attribute scopes in SimStates (#228)

Co-authored-by: Rhys Goodall <[email protected]>
Co-authored-by: Orion Cohen <[email protected]>

Author: 3 people

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py

@@ -76,6 +76,9 @@ class HybridSwapMCState(MDState):
     """
 
     last_permutation: torch.Tensor
+    _atom_attributes = (
+        MDState._atom_attributes | {"last_permutation"}  # noqa: SLF001
+    )
 
 
 nvt_init, nvt_step = nvt_langevin(model=model, dt=0.002, kT=kT)

examples/tutorials/hybrid_swap_tutorial.py

@@ -34,9 +34,11 @@
 """
 
 # %%
+from typing import ClassVar
 import torch
 import torch_sim as ts
 from mace.calculators.foundations_models import mace_mp
+from torch_sim.integrators.md import MDState
 from torch_sim.models.mace import MaceModel
 
 # Initialize the mace model
@@ -104,6 +106,9 @@ class HybridSwapMCState(ts.integrators.MDState):
     """
 
     last_permutation: torch.Tensor
+    _atom_attributes = (
+        MDState._atom_attributes | {"last_permutation"}  # noqa: SLF001
+    )
 
 
 # %% [markdown]

examples/tutorials/state_tutorial.py

@@ -79,16 +79,28 @@
   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]
 """
@@ -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]

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,52 @@
     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 | {"attribute_specified_in_scopes"}  # noqa: SLF001
+            )
+
+    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._system_attributes | {"duplicated_attribute"}  # noqa: SLF001
+            _global_attributes = SimState._global_attributes | {"duplicated_attribute"}  # noqa: SLF001
+
+    assert "are declared multiple times" in str(excinfo.value)
+    assert "duplicated_attribute" in str(excinfo.value)
 
 
 def test_slice_substate(
@@ -497,6 +511,11 @@ def test_column_vector_cell(si_sim_state: ts.SimState) -> None:
 class DeformState(SimState, DeformGradMixin):
     """Test class that combines SimState with DeformGradMixin."""
 
+    _system_attributes = (
+        SimState._system_attributes  # noqa: SLF001
+        | DeformGradMixin._system_attributes  # noqa: SLF001
+    )
+
     def __init__(
         self,
         *args,

torch_sim/integrators/md.py

@@ -41,6 +41,13 @@ 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

@@ -67,6 +67,18 @@ class NPTLangevinState(SimState):
     cell_velocities: torch.Tensor
     cell_masses: torch.Tensor
 
+    _atom_attributes = (
+        SimState._atom_attributes | {"forces", "velocities"}  # noqa: SLF001
+    )
+    _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."""
@@ -867,6 +879,25 @@ 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

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

torch_sim/monte_carlo.py

@@ -36,6 +36,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