Introduce more extensible scheme for extra properties.

.github/workflows/docs.yml

@@ -34,7 +34,7 @@ jobs:
           uv pip install ".[test,docs]" --system
 
       - name: Install extras for tutorial generation
-        run: uv pip install ".[mace,metatomic]" --system
+        run: uv pip install ".[mace]" --system
 
       - name: Copy tutorials
         run: |

pyproject.toml

@@ -3,6 +3,7 @@ name = "torch-sim-atomistic"
 version = "0.5.2"
 description = "A pytorch toolkit for calculating material properties using MLIPs"
 authors = [
+    { name = "TorchSim Maintainers", email = "torchsimatomistic@gmail.com" },
     { name = "Abhijeet Gangan", email = "abhijeetgangan@g.ucla.edu" },
     { name = "Janosh Riebesell", email = "janosh.riebesell@gmail.com" },
     { name = "Orion Cohen", email = "orioncohen@berkeley.edu" },

tests/models/conftest.py

@@ -1,5 +1,7 @@
 """Pytest fixtures and test factories for model testing."""
 
+from __future__ import annotations
+
 import typing
 
 import pytest
@@ -10,7 +12,10 @@
 
 
 if typing.TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
     from torch_sim.models.interface import ModelInterface
+    from torch_sim.state import SimState
 
 
 def make_model_calculator_consistency_test(
@@ -81,22 +86,39 @@ def make_validate_model_outputs_test(
     dtype: torch.dtype = DTYPE,
     *,
     check_detached: bool = True,
+    state_modifiers: Sequence[Callable[[SimState], SimState]] = (),
 ):
     """Factory function to create model output validation tests.
 
+    Runs ``validate_model_outputs`` once with no modifier (baseline), then
+    once more for each entry in *state_modifiers* so that every modifier
+    gets a full, independent validation pass.
+
     Args:
         model_fixture_name: Name of the model fixture to validate
         device: Device to run validation on
         dtype: Data type to use for validation
         check_detached: Whether to assert output tensors are detached from the
             autograd graph (skipped for models with ``retain_graph=True``).
+        state_modifiers: Each callable receives a ``SimState`` and returns a
+            (possibly new) ``SimState``.  The full validation suite is run
+            once per modifier so that different input edge-cases are
+            exercised independently.
     """
     from torch_sim.models.interface import validate_model_outputs
 
     def test_model_output_validation(request: pytest.FixtureRequest) -> None:
         """Test that a model implementation follows the ModelInterface contract."""
         model: ModelInterface = request.getfixturevalue(model_fixture_name)
-        validate_model_outputs(model, device, dtype, check_detached=check_detached)
+        modifiers = state_modifiers or [None]
+        for modifier in modifiers:
+            validate_model_outputs(
+                model,
+                device,
+                dtype,
+                check_detached=check_detached,
+                state_modifier=modifier,
+            )
 
     test_model_output_validation.__name__ = f"test_{model_fixture_name}_output_validation"
     return test_model_output_validation

tests/models/test_fairchem.py

@@ -1,23 +1,21 @@
 import traceback
+from collections.abc import Callable
 
 import pytest
 import torch
+from ase.build import bulk, fcc100, molecule
 
 import torch_sim as ts
 from tests.conftest import DEVICE, DTYPE
 from tests.models.conftest import make_validate_model_outputs_test
 
 
 try:
-    from collections.abc import Callable
-
-    from ase.build import bulk, fcc100, molecule
     from fairchem.core.calculate.pretrained_mlip import (
         pretrained_checkpoint_path_from_name,
     )
     from huggingface_hub.utils._auth import get_token
 
-    import torch_sim as ts
     from torch_sim.models.fairchem import FairChemModel
 
 except (ImportError, OSError, RuntimeError, AttributeError, ValueError):
@@ -263,10 +261,12 @@ def test_fairchem_charge_spin(charge: float, spin: float) -> None:
     mol.info["charge"] = charge
     mol.info["spin"] = spin
 
-    # Convert to SimState (should extract charge/spin)
-    state = ts.io.atoms_to_state([mol], device=DEVICE, dtype=DTYPE)
-
-    # Verify charge/spin were extracted correctly
+    state = ts.io.atoms_to_state(
+        [mol],
+        device=DEVICE,
+        dtype=DTYPE,
+        system_extras_map={"charge": "charge", "spin": "spin"},
+    )
     assert state.charge is not None
     assert state.spin is not None
     assert state.charge[0].item() == charge

tests/test_elastic.py

@@ -221,7 +221,7 @@ def test_get_elementary_deformations_strain_consistency(
         n_deform=n_deform,
         max_strain_normal=max_strain_normal,
         max_strain_shear=max_strain_shear,
-        bravais_type=BravaisType.triclinic,  # Test all axes
+        bravais_type=BravaisType.TRICLINIC,  # Test all axes
     )
 
     # Should generate deformations for all 6 axes (triclinic)
@@ -271,12 +271,12 @@ def mace_model() -> MaceModel:
 @pytest.mark.parametrize(
     ("sim_state_name", "expected_bravais_type", "atol"),
     [
-        ("cu_sim_state", BravaisType.cubic, 2e-1),
-        ("mg_sim_state", BravaisType.hexagonal, 5e-1),
-        ("sb_sim_state", BravaisType.trigonal, 5e-1),
-        ("tio2_sim_state", BravaisType.tetragonal, 5e-1),
-        ("ga_sim_state", BravaisType.orthorhombic, 5e-1),
-        ("niti_sim_state", BravaisType.monoclinic, 5e-1),
+        ("cu_sim_state", BravaisType.CUBIC, 2e-1),
+        ("mg_sim_state", BravaisType.HEXAGONAL, 5e-1),
+        ("sb_sim_state", BravaisType.TRIGONAL, 5e-1),
+        ("tio2_sim_state", BravaisType.TETRAGONAL, 5e-1),
+        ("ga_sim_state", BravaisType.ORTHORHOMBIC, 5e-1),
+        ("niti_sim_state", BravaisType.MONOCLINIC, 5e-1),
     ],
 )
 def test_elastic_tensor_symmetries(
@@ -340,7 +340,7 @@ def test_elastic_tensor_symmetries(
     )
     C_triclinic = (
         calculate_elastic_tensor(
-            state=state, model=model, bravais_type=BravaisType.triclinic
+            state=state, model=model, bravais_type=BravaisType.TRICLINIC
         )
         * UnitConversion.eV_per_Ang3_to_GPa
     )