Add integration tests checking the physical observables resulting from different the integrators (#521)

Author: thomasloux

.gitignore

@@ -21,6 +21,7 @@ docs/reference/torch_sim.*
 
 # ipynb files
 *.ipynb
+!docs/tutorials/integrator_tests_analysis.ipynb
 
 # ignore trajectory files
 *.h5
@@ -38,6 +39,9 @@ coverage.xml
 # test cache (compiled models, etc.)
 tests/.cache/
 
+# physical validation data and plots
+tests/physical_validation_data/
+
 # env
 uv.lock
 

docs/tutorials/index.rst

@@ -20,3 +20,4 @@ versions of the tutorials can also be found in the `torch-sim /examples/tutorial
     hybrid_swap_tutorial
     using_graphpes_tutorial
     metatomic_tutorial
+    integrator_tests_analysis

docs/tutorials/integrator_tests_analysis.ipynb

@@ -50,7 +50,7 @@ sim_state.rng = 42  # required for reproducibility — torch.manual_seed() has n
 
 ### Deterministic vs stochastic integrators in TorchSim
 
-- `ts.Integrator.nvt_langevin` and `ts.Integrator.npt_langevin` include stochastic
+- `ts.Integrator.nvt_langevin` and `ts.Integrator.npt_langevin_anisotropic` include stochastic
   terms by design. When seeded via `state.rng`, they produce identical trajectories.
   The `rng` generator controls **both** the initial momenta sampling **and** all per-step stochastic noise (Langevin OU noise, V-Rescale draws, C-Rescale barostat noise, etc.). It is stored on the state and automatically advances on every step, so running the same seed twice produces identical trajectories.
 - `ts.Integrator.nvt_nose_hoover` and `ts.Integrator.nve` are deterministic at the

docs/user/reproducibility.md

@@ -199,7 +199,7 @@ def run_ase_optimization(
         atoms = adaptor.get_atoms(Structure.from_dict(struct_dict))
         atoms.calc = calculator
         system: Any = cell_filter_cls(atoms) if cell_filter_cls is not None else atoms
-        opt = ase_optimizer_cls(system, logfile=os.devnull)  # type: ignore[arg-type]
+        opt = ase_optimizer_cls(system, logfile=os.devnull)
         opt.run(fmax=f_max, steps=max_steps)
         if opt.get_number_of_steps() < max_steps:
             converged += 1

examples/benchmarking/opt-throughput.py

@@ -323,7 +323,7 @@
 target_pressure = torch.tensor(0.0 * Units.pressure, device=device, dtype=dtype)  # 0 bar
 
 # Initialize NPT with NVT equilibration
-state = ts.npt_nose_hoover_init(
+state = ts.npt_nose_hoover_isotropic_init(
     state=state, model=mace_model_stress, kT=kT, dt=torch.tensor(dt)
 )
 
@@ -337,12 +337,14 @@
             / Units.temperature
         )
         invariant = float(
-            ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure)
+            ts.npt_nose_hoover_isotropic_invariant(
+                state, kT=kT, external_pressure=target_pressure
+            )
         )
         log.info(
             f"Step {step}: Temperature: {temp.item():.4f} K, Invariant: {invariant:.4f}"
         )
-    state = ts.npt_nose_hoover_step(
+    state = ts.npt_nose_hoover_isotropic_step(
         state=state,
         model=mace_model_stress,
         dt=torch.tensor(dt),
@@ -351,7 +353,7 @@
     )
 
 # Reinitialize for NPT phase
-state = ts.npt_nose_hoover_init(
+state = ts.npt_nose_hoover_isotropic_init(
     state=state, model=mace_model_stress, kT=kT, dt=torch.tensor(dt)
 )
 
@@ -365,7 +367,9 @@
             / Units.temperature
         )
         invariant = float(
-            ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure)
+            ts.npt_nose_hoover_isotropic_invariant(
+                state, kT=kT, external_pressure=target_pressure
+            )
         )
         stress = mace_model_stress(state)["stress"]
         volume = torch.det(state.current_cell)
@@ -379,7 +383,7 @@
             f"Pressure: {pressure:.4f} eV/ų, "
             f"Cell: [{xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}]"
         )
-    state = ts.npt_nose_hoover_step(
+    state = ts.npt_nose_hoover_isotropic_step(
         state=state,
         model=mace_model_stress,
         dt=torch.tensor(dt),

examples/scripts/3_dynamics.py

@@ -39,6 +39,7 @@ dependencies = [
 [project.optional-dependencies]
 test = [
     "torch-sim-atomistic[io,symmetry,vesin]",
+    "physical-validation>=1.0.5",
     "platformdirs>=4.0.0",
     "psutil>=7.0.0",
     "pymatgen>=2025.6.14",
@@ -127,20 +128,25 @@ isort.lines-after-imports = 2
 pep8-naming.ignore-names = ["get_kT", "kT"]
 
 [tool.ruff.lint.per-file-ignores]
-"**/tests/*" = ["ANN201", "D", "INP001", "S101"]
-"examples/**/*" = ["B018", "T201"]
+"**/tests/*" = ["ANN001", "ANN201", "ANN202", "D", "INP001", "S101", "T201"]
+"examples/**/*" = ["ANN001", "ANN201", "ANN202", "B018", "T201"]
 "examples/tutorials/**/*" = ["ALL"]
+"docs/tutorials/*.ipynb" = ["ANN001", "ANN201", "B905", "BLE001", "E501", "F841", "N816", "PLR1714", "RUF001", "T201"]
 
 [tool.ruff.format]
 docstring-code-format = true
 
 [tool.codespell]
 check-filenames = true
 ignore-words-list = ["convertor"]  # codespell:ignore convertor
+skip = "docs/tutorials/integrator_tests_analysis.ipynb"
 
 [tool.pytest.ini_options]
-addopts = ["-p no:warnings"]
+addopts = ["-p no:warnings", "-m not physical_validation"]
 testpaths = ["tests"]
+markers = [
+    "physical_validation: long-running physical validation tests (run with: pytest -m physical_validation)",
+]
 
 [tool.uv]
 # make these dependencies mutually exclusive since they use incompatible e3nn versions
@@ -220,18 +226,32 @@ include = [
 [tool.ty.overrides.rules]
 unresolved-import = "ignore"
 
+[[tool.ty.overrides]]
+include = [
+    "torch_sim/models/dispersion.py",
+    "torch_sim/neighbors/vesin.py",
+]
+[tool.ty.overrides.rules]
+invalid-argument-type = "ignore"
+invalid-assignment = "ignore"
+
 [[tool.ty.overrides]]
 include = ["tests/**/*.py"]
 
 [tool.ty.overrides.rules]
 invalid-argument-type = "ignore"
+invalid-assignment = "ignore"
 no-matching-overload = "ignore"
 unresolved-attribute = "ignore"
 unresolved-import = "ignore"
 
 [[tool.ty.overrides]]
-include = ["docs/**/*.py", "examples/**/*.py"]
+include = ["docs/**/*.py", "docs/**/*.ipynb", "examples/**/*.py"]
 [tool.ty.overrides.rules]
+invalid-argument-type = "ignore"
+not-iterable = "ignore"
+not-subscriptable = "ignore"
+unresolved-attribute = "ignore"
 unresolved-import = "ignore"
 
 

pyproject.toml

@@ -15,6 +15,22 @@
 
 torch.set_num_threads(4)
 
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--validation-plots",
+        action="store_true",
+        default=False,
+        help="Save physical validation plots to tests/physical_validation_data/plots/",
+    )
+    parser.addoption(
+        "--clean-validation-data",
+        action="store_true",
+        default=False,
+        help="Delete saved physical validation data before running tests",
+    )
+
+
 DEVICE = torch.device("cpu")
 DTYPE = torch.float64
 

tests/conftest.py

@@ -555,8 +555,8 @@ def test_constraint_validation_errors(
     [
         ("nve", FixAtoms(atom_idx=[0, 1]), 100),
         ("nvt_nose_hoover", FixCom([0]), 200),
-        ("npt_langevin", FixAtoms(atom_idx=[0, 3]), 200),
-        ("npt_nose_hoover", FixCom([0]), 200),
+        ("npt_langevin_anisotropic", FixAtoms(atom_idx=[0, 3]), 200),
+        ("npt_nose_hoover_isotropic", FixCom([0]), 200),
     ],
 )
 def test_integrators_with_constraints(
@@ -591,20 +591,20 @@ def test_integrators_with_constraints(
         state = ts.nvt_nose_hoover_init(cu_sim_state, lj_model, kT=kT, dt=dt)
         for _ in range(n_steps):
             state = ts.nvt_nose_hoover_step(state, lj_model, dt=dt, kT=kT)
-    elif integrator == "npt_langevin":
-        state = ts.npt_langevin_init(cu_sim_state, lj_model, kT=kT, dt=dt)
+    elif integrator == "npt_langevin_anisotropic":
+        state = ts.npt_langevin_anisotropic_init(cu_sim_state, lj_model, kT=kT, dt=dt)
         for _ in range(n_steps):
-            state = ts.npt_langevin_step(
+            state = ts.npt_langevin_anisotropic_step(
                 state,
                 lj_model,
                 dt=dt,
                 kT=kT,
                 external_pressure=torch.tensor(0.0, dtype=DTYPE),
             )
-    else:  # npt_nose_hoover
-        state = ts.npt_nose_hoover_init(cu_sim_state, lj_model, kT=kT, dt=dt)
+    else:  # npt_nose_hoover_isotropic
+        state = ts.npt_nose_hoover_isotropic_init(cu_sim_state, lj_model, kT=kT, dt=dt)
         for _ in range(n_steps):
-            state = ts.npt_nose_hoover_step(
+            state = ts.npt_nose_hoover_isotropic_step(
                 state,
                 lj_model,
                 dt=torch.tensor(0.001, dtype=DTYPE),

tests/test_constraints.py

@@ -539,9 +539,9 @@ def test_build_symmetry_map_chunked_matches_vectorized() -> None:
 
     old_threshold = sym_mod._SYMM_MAP_CHUNK_THRESHOLD  # noqa: SLF001
     try:
-        sym_mod._SYMM_MAP_CHUNK_THRESHOLD = len(state.positions) + 1  # noqa: SLF001  # ty: ignore[invalid-assignment]
+        sym_mod._SYMM_MAP_CHUNK_THRESHOLD = len(state.positions) + 1  # noqa: SLF001
         vectorized = build_symmetry_map(rotations, translations, frac)
-        sym_mod._SYMM_MAP_CHUNK_THRESHOLD = 0  # noqa: SLF001  # ty: ignore[invalid-assignment]
+        sym_mod._SYMM_MAP_CHUNK_THRESHOLD = 0  # noqa: SLF001
         chunked = build_symmetry_map(rotations, translations, frac)
     finally:
         sym_mod._SYMM_MAP_CHUNK_THRESHOLD = old_threshold  # noqa: SLF001