From 4aff5b36d3506ebf3ddcc2d436da3d791054f768 Mon Sep 17 00:00:00 2001 From: abhijeetgangan Date: Sun, 25 Oct 2026 08:26:14 -0700 Subject: [PATCH] Bump CI actions and lint hooks, plus small correctness fixes --- .github/workflows/docs.yml | 10 +++--- .github/workflows/link-check.yml | 2 +- .github/workflows/lint.yml | 4 +-- .github/workflows/stale.yml | 2 +- .github/workflows/test.yml | 24 +++++++------- .pre-commit-config.yaml | 9 +++-- README.md | 2 ++ docs/user/reproducibility.md | 2 +- pyproject.toml | 3 +- tests/models/test_nequip_framework.py | 2 +- tests/test_autobatching.py | 13 ++++++++ tests/test_correlations.py | 33 +++++++++++++++++++ tests/test_optimizers_vs_ase.py | 10 +++--- tests/test_quantities.py | 20 ++++++++++++ tests/test_state.py | 17 ++++++++++ torch_sim/__init__.py | 18 ---------- torch_sim/_duecredit.py | 16 ++++----- torch_sim/autobatching.py | 12 ++++--- torch_sim/constraints.py | 12 +++---- torch_sim/models/interface.py | 2 +- torch_sim/neighbors/vesin.py | 3 +- torch_sim/optimizers/__init__.py | 7 +--- torch_sim/optimizers/cell_filters.py | 5 +-- torch_sim/properties/correlations.py | 47 ++++++++++++++------------- torch_sim/quantities.py | 18 +++++----- torch_sim/state.py | 9 ++--- torch_sim/trajectory.py | 5 +-- 27 files changed, 186 insertions(+), 121 deletions(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index b1b2f8f13..5adeeeff2 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -17,17 +17,17 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Install pandoc run: sudo apt-get install pandoc - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v6 with: python-version: '3.12' - name: Set up uv - uses: astral-sh/setup-uv@v6 + uses: astral-sh/setup-uv@v8.2.0 - name: Install dependencies run: | @@ -53,7 +53,7 @@ jobs: run: sphinx-build docs docs_build - name: Upload build artifact - uses: actions/upload-pages-artifact@v3 + uses: actions/upload-pages-artifact@v5 with: path: ./docs_build @@ -67,4 +67,4 @@ jobs: steps: - name: Deploy to GitHub Pages id: deployment - uses: actions/deploy-pages@v4 + uses: actions/deploy-pages@v5 diff --git a/.github/workflows/link-check.yml b/.github/workflows/link-check.yml index 791adec3c..ef4d24550 100644 --- a/.github/workflows/link-check.yml +++ b/.github/workflows/link-check.yml @@ -11,7 +11,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Check out repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Discover broken links uses: lycheeverse/lychee-action@v2 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 3c17bc5a8..1694dbe64 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -11,7 +11,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Check out repo - uses: actions/checkout@v5 + uses: actions/checkout@v6 - name: Run prek - uses: j178/prek-action@v1 + uses: j178/prek-action@v2 diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index d80f7654d..d65b3e08f 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -9,7 +9,7 @@ jobs: stale-pull-requests: runs-on: ubuntu-latest steps: - - uses: actions/stale@v9 + - uses: actions/stale@v10 with: repo-token: ${{ secrets.GITHUB_TOKEN }} days-before-issue-stale: -1 # Disable issue staleness, only manage PRs diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3ff79b68c..d64e39891 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -26,15 +26,15 @@ jobs: steps: - name: Check out repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.version.python }} - name: Set up uv - uses: astral-sh/setup-uv@v6 + uses: astral-sh/setup-uv@v8.2.0 - name: Install torch_sim run: | @@ -46,7 +46,7 @@ jobs: pytest --cov=torch_sim --cov-report=xml - name: Upload coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} slug: torchsim/torch-sim @@ -81,10 +81,10 @@ jobs: steps: - name: Check out repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.version.python }} @@ -93,7 +93,7 @@ jobs: run: brew install hdf5 - name: Set up uv - uses: astral-sh/setup-uv@v6 + uses: astral-sh/setup-uv@v8.2.0 - name: Install torch_sim with model dependencies run: | @@ -112,7 +112,7 @@ jobs: pytest -vv -ra -rs --cov=torch_sim --cov-report=xml ${{ matrix.model.test_path }} - name: Upload coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} slug: torchsim/torch-sim @@ -123,7 +123,7 @@ jobs: examples: ${{ steps.set-matrix.outputs.examples }} steps: - name: Check out repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Find example scripts id: set-matrix @@ -140,15 +140,15 @@ jobs: example: ${{fromJson(needs.find-examples.outputs.examples)}} steps: - name: Check out repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: '3.12' - name: Set up uv - uses: astral-sh/setup-uv@v6 + uses: astral-sh/setup-uv@v8.2.0 - name: Run example if: ${{ !contains(matrix.example, 'fairchem') || github.event.pull_request.head.repo.fork == false }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a562c576b..5d47de8b9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,14 +2,13 @@ default_install_hook_types: [pre-commit, commit-msg] repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.4 + rev: v0.17.1 hooks: - id: ruff-check args: [--fix] - id: ruff-format - - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v6.0.0 + - repo: builtin hooks: - id: check-case-conflict - id: check-symlinks @@ -21,13 +20,13 @@ repos: - id: trailing-whitespace - repo: https://github.com/codespell-project/codespell - rev: v2.4.1 + rev: v2.4.3 hooks: - id: codespell stages: [pre-commit, commit-msg] - repo: https://github.com/igorshubovych/markdownlint-cli - rev: v0.47.0 + rev: v0.49.1 hooks: - id: markdownlint # MD013: line length diff --git a/README.md b/README.md index 0a62be3b5..c76e39e3c 100644 --- a/README.md +++ b/README.md @@ -43,10 +43,12 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # easily load the model from mace-mp from mace.calculators.foundations_models import mace_mp from torch_sim.models.mace import MaceModel + mace = mace_mp(model="small", return_raw_model=True) mace_model = MaceModel(model=mace, device=device) from ase.build import bulk + cu_atoms = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2)) many_cu_atoms = [cu_atoms] * 50 trajectory_files = [f"Cu_traj_{i}.h5md" for i in range(len(many_cu_atoms))] diff --git a/docs/user/reproducibility.md b/docs/user/reproducibility.md index 9bee7e0f9..f6b045049 100644 --- a/docs/user/reproducibility.md +++ b/docs/user/reproducibility.md @@ -45,7 +45,7 @@ In PyTorch, `torch.manual_seed()` only seeds the default global generator, i.e. ```python sim_state = ts.initialize_state(atoms, device, dtype) -sim_state.rng = 42 # required for reproducibility — torch.manual_seed() has no effect here +sim_state.rng = 42 # needed for reproducibility: torch.manual_seed() has no effect ``` ### Deterministic vs stochastic integrators in TorchSim diff --git a/pyproject.toml b/pyproject.toml index 29c88e178..b9f1b970c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,7 +89,6 @@ module-name = "torch_sim" module-root = "" [tool.ruff] -target-version = "py312" line-length = 90 output-format = "concise" @@ -101,6 +100,7 @@ ignore = [ "ANN401", # Allow use of type Any for function argument {name} "C408", # Unnecessary {obj_type} call (rewrite as a literal) "COM812", # missing-trailing-comma + "CPY001", # Missing copyright notice at top of file (license lives in LICENSE) "D205", # 1 blank line required between summary line and description "EM101", # Exception must not use a string literal, assign to variable first "EM102", # Exception must not use an f-string literal, assign to variable first @@ -112,6 +112,7 @@ ignore = [ "PLC0415", # import should be at the top-level of a file "PLR0912", # too many branches "PLR0913", # too many function arguments + "PLR0917", # too many positional arguments (same reason as PLR0913) "PLR2004", # Magic value used in comparison, consider replacing {value} with a constant variable "PLW2901", # Outer for loop variable overwritten by inner assignment target "PTH", # flake8-use-pathlib diff --git a/tests/models/test_nequip_framework.py b/tests/models/test_nequip_framework.py index 1a82a6051..7f6a7d504 100644 --- a/tests/models/test_nequip_framework.py +++ b/tests/models/test_nequip_framework.py @@ -45,7 +45,7 @@ def _get_nequip_model_zip() -> Path: zip_path = NEQUIP_CACHE_DIR / NEQUIP_OAM_S_ZIP_NAME if not zip_path.exists(): - urllib.request.urlretrieve(NEQUIP_OAM_S_ZENODO_URL, zip_path) # noqa: S310 + urllib.request.urlretrieve(NEQUIP_OAM_S_ZENODO_URL, zip_path) return zip_path diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index 99c9bb96e..feeff7a31 100644 --- a/tests/test_autobatching.py +++ b/tests/test_autobatching.py @@ -92,6 +92,19 @@ def test_bounds_and_tuples(): ] +@pytest.mark.parametrize( + ("items", "lower_bound", "upper_bound"), [([1, 2], 11, 1), ([1, 2], 5, 5), ([], 5, 5)] +) +def test_to_constant_volume_bins_invalid_bounds( + items: list[int], lower_bound: int, upper_bound: int +) -> None: + """Inverted or empty bounds raise instead of silently applying only lower_bound.""" + with pytest.raises(ValueError, match="lower_bound is greater or equal"): + to_constant_volume_bins( + items, 5, lower_bound=lower_bound, upper_bound=upper_bound + ) + + def test_calculate_scaling_metric(si_sim_state: ts.SimState) -> None: """Test calculation of scaling metrics for a state.""" # Test n_atoms metric diff --git a/tests/test_correlations.py b/tests/test_correlations.py index 5b3cbb0b8..b92ac71d1 100644 --- a/tests/test_correlations.py +++ b/tests/test_correlations.py @@ -14,9 +14,11 @@ import torch_sim as ts from tests.conftest import DEVICE +from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.properties.correlations import ( CircularBuffer, CorrelationCalculator, + HeatFluxAutoCorrelation, VelocityAutoCorrelation, ) @@ -475,3 +477,34 @@ def test_velocity_autocorrelation_with_trajectory_reporter( assert props[0]["vacf"].shape == (1,) reporter.close() + + +def test_heat_flux_autocorrelation( + lj_model: LennardJonesModel, ar_double_sim_state: ts.SimState +) -> None: + """HFACF runs one model forward per update and yields a normalized HFACF.""" + window_size = 4 + hfacf_calc = HeatFluxAutoCorrelation( + model=lj_model, window_size=window_size, device=DEVICE + ) + state = ts.nve_init(ar_double_sim_state, lj_model, kT=0.01) + heat_flux = hfacf_calc.corr_calc.properties["heat_flux"](state) + assert heat_flux.shape == (state.n_systems, 3) + assert heat_flux.abs().sum() > 0 + + n_forward_calls = 0 + + def count_forward(*_args: object) -> None: + nonlocal n_forward_calls + n_forward_calls += 1 + + lj_model.register_forward_hook(count_forward) + for _ in range(window_size): + state.momenta = torch.randn_like(state.momenta) + n_windows = hfacf_calc(state) + + assert n_forward_calls == window_size + assert n_windows.tolist() == [1] + assert hfacf_calc.hfacf.shape == (window_size,) + assert torch.isfinite(hfacf_calc.hfacf).all() + assert hfacf_calc.hfacf[0] == pytest.approx(1.0) diff --git a/tests/test_optimizers_vs_ase.py b/tests/test_optimizers_vs_ase.py index 8e6f67b77..f59653ee9 100644 --- a/tests/test_optimizers_vs_ase.py +++ b/tests/test_optimizers_vs_ase.py @@ -112,7 +112,7 @@ def _run_and_compare_optimizers( ase_mace_mpa: "MACECalculator", fire_type: ts.Optimizer, cell_filter: ts.CellFilter, - ase_filter_cls: FrechetCellFilter | UnitCellFilter, + ase_filter_cls: type[FrechetCellFilter | UnitCellFilter], checkpoints: list[int], force_tol: float, tolerances: dict[str, float], @@ -129,7 +129,7 @@ def _run_and_compare_optimizers( initial_sim_state_fixture.clone().to(dtype=DTYPE, device=device) )[0] ase_atoms.calc = ase_mace_mpa - filtered_ase_atoms = ase_filter_cls(ase_atoms) # type: ignore[call-non-callable] + filtered_ase_atoms = ase_filter_cls(ase_atoms) ase_optimizer = FIRE(filtered_ase_atoms, logfile=None) last_checkpoint_step_count = 0 @@ -424,7 +424,7 @@ def test_optimizer_vs_ase_parametrized( sim_state_fixture_name: str, fire_type: ts.Optimizer, cell_filter: ts.CellFilter, - ase_filter_cls: FrechetCellFilter | UnitCellFilter, + ase_filter_cls: type[FrechetCellFilter | UnitCellFilter], checkpoints: list[int], force_tol: float, tolerances: dict[str, float], @@ -534,7 +534,7 @@ def test_fire_frechet_matches_ase_strict( def test_bfgs_vs_ase_parametrized( sim_state_fixture_name: str, cell_filter: ts.CellFilter, - ase_filter_cls: type, + ase_filter_cls: type[FrechetCellFilter | UnitCellFilter], checkpoints: list[int], force_tol: float, tolerances: dict[str, float], @@ -610,7 +610,7 @@ def test_bfgs_vs_ase_parametrized( def test_lbfgs_vs_ase_parametrized( sim_state_fixture_name: str, cell_filter: ts.CellFilter, - ase_filter_cls: type, + ase_filter_cls: type[FrechetCellFilter | UnitCellFilter], checkpoints: list[int], force_tol: float, tolerances: dict[str, float], diff --git a/tests/test_quantities.py b/tests/test_quantities.py index 619524594..2f5a46630 100644 --- a/tests/test_quantities.py +++ b/tests/test_quantities.py @@ -1,5 +1,7 @@ """Tests for quantities module functions.""" +from collections.abc import Callable + import pytest import torch from numpy.testing import assert_allclose @@ -222,6 +224,24 @@ def test_calc_kinetic_energy_batched_system( assert torch.allclose(ke_mom, batched_system_data["ke"]) +@pytest.mark.parametrize("calc_fn", [calc_kinetic_energy, calc_kT]) +def test_per_system_sums_with_unsorted_system_idx(calc_fn: Callable[..., Tensor]) -> None: + """Per-system reductions must not assume atoms are grouped by system.""" + torch.manual_seed(0) + system_idx = torch.tensor([1, 0, 1, 2, 0, 2, 1], device=DEVICE) + masses = torch.rand(7, device=DEVICE, dtype=DTYPE) + 0.5 + velocities = torch.randn(7, 3, device=DEVICE, dtype=DTYPE) + + result = calc_fn(masses=masses, velocities=velocities, system_idx=system_idx) + expected = torch.stack( + [ + calc_fn(masses=masses[mask], velocities=velocities[mask]) + for mask in (system_idx == sys_id for sys_id in range(3)) + ] + ) + assert torch.allclose(result, expected) + + def test_calc_kinetic_energy_errors(single_system_data: dict[str, Tensor]) -> None: with pytest.raises(ValueError, match="Must pass either one of momenta or velocities"): calc_kinetic_energy( diff --git a/tests/test_state.py b/tests/test_state.py index ffe484c89..775361839 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -275,6 +275,23 @@ def test_concatenate_double_si_and_fe_states( assert torch.allclose(fe_slice.positions, fe_supercell_sim_state.positions) +@pytest.mark.parametrize( + "system_idx", + [ + [1] * 8 + [0] * 8, # descending + [0, 1] * 8, # interleaved + [0] * 8 + [2] * 8, # gap + [1] * 16, # not starting at 0 + ], +) +def test_invalid_system_idx_raises( + si_double_sim_state: SimState, system_idx: list[int] +) -> None: + """Atoms must be grouped by system in ascending order starting at 0.""" + with pytest.raises(ValueError, match="unique consecutive integers"): + si_double_sim_state.system_idx = torch.tensor(system_idx, device=DEVICE) + + def test_split_state(si_double_sim_state: SimState) -> None: """Test splitting a state into a list of states.""" states = si_double_sim_state.split() diff --git a/torch_sim/__init__.py b/torch_sim/__init__.py index 9a848f9f2..4886ddb13 100644 --- a/torch_sim/__init__.py +++ b/torch_sim/__init__.py @@ -112,12 +112,9 @@ __all__ = [ "CELL_FILTER_REGISTRY", "INTEGRATOR_REGISTRY", - "INTEGRATOR_REGISTRY", - "OPTIM_REGISTRY", "OPTIM_REGISTRY", "BFGSState", "BinningAutoBatcher", - "BinningAutoBatcher", "CellBFGSState", "CellFilter", "CellFireState", @@ -126,20 +123,16 @@ "CorrelationCalculator", "FireState", "InFlightAutoBatcher", - "InFlightAutoBatcher", - "Integrator", "Integrator", "LBFGSState", "NPTLangevinAnisotropicState", "NPTLangevinIsotropicState", "NPTNoseHooverIsotropicState", "NVTNoseHooverState", - "NVTNoseHooverState", "OptimState", "Optimizer", "SimState", "SwapMCState", - "SwapMCState", "TorchSimTrajectory", "TrajectoryReporter", "autobatching", @@ -180,22 +173,13 @@ "npt_nose_hoover_isotropic_invariant", "npt_nose_hoover_isotropic_step", "nve_init", - "nve_init", - "nve_step", "nve_step", "nvt_langevin_init", - "nvt_langevin_init", - "nvt_langevin_step", "nvt_langevin_step", "nvt_nose_hoover_init", - "nvt_nose_hoover_init", "nvt_nose_hoover_invariant", - "nvt_nose_hoover_invariant", - "nvt_nose_hoover_step", "nvt_nose_hoover_step", "nvt_vrescale_init", - "nvt_vrescale_init", - "nvt_vrescale_step", "nvt_vrescale_step", "optimize", "optimizers", @@ -204,8 +188,6 @@ "state", "static", "swap_mc_init", - "swap_mc_init", - "swap_mc_step", "swap_mc_step", "system_wise_max_force", "trajectory", diff --git a/torch_sim/_duecredit.py b/torch_sim/_duecredit.py index e20421f7f..3fb13a1f4 100644 --- a/torch_sim/_duecredit.py +++ b/torch_sim/_duecredit.py @@ -37,19 +37,15 @@ def _donothing_func(*_args: Any, **_kwargs: Any) -> Any: return None -def _disable_duecredit(exc: Exception) -> None: - import logging - - logging.getLogger("duecredit").exception( - "Failed to import duecredit despite being installed: %s", exc - ) - - try: from duecredit import BibTeX, Doi, Text, Url, due -except Exception as e: # noqa: BLE001 +except Exception as e: if not isinstance(e, ImportError): - _disable_duecredit(e) + import logging + + logging.getLogger("duecredit").exception( + "Failed to import duecredit despite being installed" + ) due = InactiveDueCreditCollector() # ty: ignore[invalid-assignment] BibTeX = Doi = Url = Text = _donothing_func # ty: ignore[invalid-assignment] diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 4b2c7c861..92797ecfa 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -21,7 +21,7 @@ """ import logging -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Iterable, Iterator, Sequence from itertools import chain from typing import Any, get_args @@ -98,6 +98,8 @@ def to_constant_volume_bins( # noqa: C901 if not hasattr(items, "__len__"): raise TypeError("items must be iterable") + if lower_bound is not None and upper_bound is not None and lower_bound >= upper_bound: + raise ValueError("lower_bound is greater or equal to upper_bound") if len(items) == 0: return [] @@ -624,7 +626,7 @@ def load_states(self, states: T | Sequence[T]) -> float: will be restarted when this method is called. """ batched = ( - states if isinstance(states, SimState) else ts.concatenate_states(states) + ts.concatenate_states(states) if isinstance(states, Sequence) else states ) self.memory_scalers = calculate_memory_scalers( batched, self.memory_scales_with, self.cutoff @@ -930,9 +932,9 @@ def load_states(self, states: Sequence[T] | Iterator[T] | T) -> float | None: This method resets the current state indices and completed state tracking, so any ongoing processing will be restarted when this method is called. """ - if isinstance(states, SimState): - states = states.split() - self.states_iterator = iter(states) + self.states_iterator = iter( + states if isinstance(states, Iterable) else states.split() + ) self.current_scalers = [] self.current_idx = [] diff --git a/torch_sim/constraints.py b/torch_sim/constraints.py index adc7b5c27..acf0661a6 100644 --- a/torch_sim/constraints.py +++ b/torch_sim/constraints.py @@ -108,7 +108,7 @@ def adjust_cell( # noqa: B027 @abstractmethod def select_constraint( self, atom_mask: torch.Tensor, system_mask: torch.Tensor - ) -> None | Self: + ) -> Self | None: """Update the constraint to account for atom and system masks. Args: @@ -117,7 +117,7 @@ def select_constraint( """ @abstractmethod - def select_sub_constraint(self, atom_idx: torch.Tensor, sys_idx: int) -> None | Self: + def select_sub_constraint(self, atom_idx: torch.Tensor, sys_idx: int) -> Self | None: """Select a constraint for a given atom and system index. Args: @@ -237,7 +237,7 @@ def select_constraint( self, atom_mask: torch.Tensor, system_mask: torch.Tensor, # noqa: ARG002 - ) -> None | Self: + ) -> Self | None: """Update the constraint to account for atom and system masks. Args: @@ -254,7 +254,7 @@ def select_sub_constraint( self, atom_idx: torch.Tensor, sys_idx: int, # noqa: ARG002 - ) -> None | Self: + ) -> Self | None: """Select a constraint for a given atom and system index. Args: @@ -347,7 +347,7 @@ def select_constraint( self, atom_mask: torch.Tensor, # noqa: ARG002 system_mask: torch.Tensor, - ) -> None | Self: + ) -> Self | None: """Update the constraint to account for atom and system masks. Args: @@ -364,7 +364,7 @@ def select_sub_constraint( self, atom_idx: torch.Tensor, # noqa: ARG002 sys_idx: int, - ) -> None | Self: + ) -> Self | None: """Select a constraint for a given atom and system index. Args: diff --git a/torch_sim/models/interface.py b/torch_sim/models/interface.py index 5a7a68347..2f728b39d 100644 --- a/torch_sim/models/interface.py +++ b/torch_sim/models/interface.py @@ -255,7 +255,7 @@ def __init__(self, *models: ModelInterface) -> None: def _children(self) -> list[ModelInterface]: """Return child models with proper typing for static analysis.""" - return list(self.models.children()) # type: ignore[return-value] + return [m for m in self.models.children() if isinstance(m, ModelInterface)] @ModelInterface.compute_stress.setter def compute_stress(self, value: bool) -> None: # noqa: FBT001 diff --git a/torch_sim/neighbors/vesin.py b/torch_sim/neighbors/vesin.py index f726f1452..064ab5e87 100644 --- a/torch_sim/neighbors/vesin.py +++ b/torch_sim/neighbors/vesin.py @@ -107,6 +107,7 @@ def vesin_nl( cell_sys = cell[sys_idx] # Calculate neighbor list for this system + # vesin>=0.6 makes cutoff keyword-only neighbor_list_fn = VesinNeighborList( cutoff=float(cutoff), full_list=True, sorted=False ) @@ -250,7 +251,7 @@ def vesin_nl_ts( continue # Calculate neighbor list for this system - neighbor_list_fn = VesinNeighborListTorch(cutoff, full_list=True) + neighbor_list_fn = VesinNeighborListTorch(cutoff=cutoff, full_list=True) # Get the cell for this system cell_sys = cell[sys_idx] diff --git a/torch_sim/optimizers/__init__.py b/torch_sim/optimizers/__init__.py index 7223e84cb..8a21c7666 100644 --- a/torch_sim/optimizers/__init__.py +++ b/torch_sim/optimizers/__init__.py @@ -23,12 +23,7 @@ gradient_descent_step, ) from torch_sim.optimizers.lbfgs import lbfgs_init, lbfgs_step -from torch_sim.optimizers.state import ( # noqa: F401 - BFGSState, - FireState, - LBFGSState, - OptimState, -) +from torch_sim.optimizers.state import BFGSState, FireState, LBFGSState, OptimState # noqa: F401 FireFlavor = Literal["vv_fire", "ase_fire"] diff --git a/torch_sim/optimizers/cell_filters.py b/torch_sim/optimizers/cell_filters.py index cc6a8eea2..66baa77d1 100644 --- a/torch_sim/optimizers/cell_filters.py +++ b/torch_sim/optimizers/cell_filters.py @@ -96,8 +96,9 @@ def _setup_pressure( def _compute_cell_masses(state: SimState) -> torch.Tensor: """Compute cell masses by summing atomic masses per system.""" - system_counts = torch.bincount(state.system_idx) - cell_masses = torch.segment_reduce(state.masses, reduce="sum", lengths=system_counts) + cell_masses = state.masses.new_zeros(state.n_systems).index_add( + 0, state.system_idx, state.masses + ) return cell_masses.unsqueeze(-1).expand(-1, 3) diff --git a/torch_sim/properties/correlations.py b/torch_sim/properties/correlations.py index 3614254b2..8e816e16d 100644 --- a/torch_sim/properties/correlations.py +++ b/torch_sim/properties/correlations.py @@ -24,6 +24,7 @@ import torch from torch_sim.elastic import full_3x3_to_voigt_6_stress +from torch_sim.integrators.md import MDState from torch_sim.quantities import calc_heat_flux from torch_sim.state import SimState @@ -97,7 +98,7 @@ def is_full(self) -> bool: return self.count == self.size -class CorrelationCalculator: +class CorrelationCalculator[S: SimState]: """Efficient on-the-fly correlation function calculator. Manage the calculation of time correlation functions during @@ -136,7 +137,7 @@ def __init__( self, *, window_size: int, - properties: dict[str, Callable[[SimState], torch.Tensor]], + properties: dict[str, Callable[[S], torch.Tensor]], device: torch.device, normalize: bool = True, ) -> None: @@ -161,9 +162,7 @@ def __init__( self.correlations: dict[str, torch.Tensor] = {} self.cross_correlations: dict[tuple[str, str], torch.Tensor] = {} - def add_property( - self, name: str, calculator: Callable[[SimState], torch.Tensor] - ) -> None: + def add_property(self, name: str, calculator: Callable[[S], torch.Tensor]) -> None: """Track a new simulation property. Args: @@ -176,7 +175,7 @@ def add_property( self.properties[name] = calculator self.buffers[name] = CircularBuffer(self.window_size, device=self.device) - def update(self, state: SimState) -> None: + def update(self, state: S) -> None: """Update correlation calculations with new state data. Args: @@ -340,7 +339,7 @@ def reset(self) -> None: self.correlations = {} self.cross_correlations = {} - def to(self, device: torch.device) -> "CorrelationCalculator": + def to(self, device: torch.device) -> "CorrelationCalculator[S]": """Move calculator to specified device. Args: @@ -431,7 +430,7 @@ def __init__( use_running_average: Whether to compute running average across windows normalize: Whether to normalize correlation functions to [0,1] """ - self.corr_calc = CorrelationCalculator( + self.corr_calc = CorrelationCalculator[MDState]( window_size=window_size, properties={"velocity": lambda s: s.velocities}, device=device, @@ -441,7 +440,7 @@ def __init__( self._window_count = 0 self._avg = torch.zeros(window_size, device=device) - def __call__(self, state: SimState, _: Any = None) -> torch.Tensor: + def __call__(self, state: MDState, _: Any = None) -> torch.Tensor: """Update VACF with new state. Args: @@ -526,20 +525,22 @@ def __init__( self.model.per_atom_stresses = True # ty: ignore[unresolved-attribute] self.model.per_atom_energies = True # ty: ignore[unresolved-attribute] - self.corr_calc = CorrelationCalculator( + def heat_flux(state: MDState) -> torch.Tensor: + output = self.model(state) # single forward pass for energies and stresses + return calc_heat_flux( + momenta=state.momenta, + masses=state.masses, + velocities=None, + energies=output["energies"], + stresses=full_3x3_to_voigt_6_stress(output["stresses"]), + batch=state.system_idx, + is_centroid_stress=False, + is_virial_only=False, + ) + + self.corr_calc = CorrelationCalculator[MDState]( window_size=window_size, - properties={ - "heat_flux": lambda s: calc_heat_flux( - momenta=s.momenta, - masses=s.masses, - velocities=None, - energies=self.model(s)["energies"], - stresses=full_3x3_to_voigt_6_stress(self.model(s)["stresses"]), - batch=s.system_idx, - is_centroid_stress=False, - is_virial_only=False, - ) - }, + properties={"heat_flux": heat_flux}, device=device, normalize=normalize, ) @@ -547,7 +548,7 @@ def __init__( self._window_count = 0 self._avg = torch.zeros(window_size, device=device) - def __call__(self, state: SimState, _: Any = None) -> torch.Tensor: + def __call__(self, state: MDState, _: Any = None) -> torch.Tensor: """Update HFACF with new state. Args: diff --git a/torch_sim/quantities.py b/torch_sim/quantities.py index ab6d04bff..88a10bf96 100644 --- a/torch_sim/quantities.py +++ b/torch_sim/quantities.py @@ -76,8 +76,8 @@ def calc_kT( # noqa: N802 dof_per_system = system_sizes * squared_term.shape[-1] # multiply by n_dimensions # Calculate temperature per system - system_sums = torch.segment_reduce( - flattened_squared, reduce="sum", lengths=system_sizes + system_sums = flattened_squared.new_zeros(len(system_sizes)).index_add( + 0, system_idx, flattened_squared ) return system_sums / dof_per_system @@ -150,8 +150,9 @@ def calc_kinetic_energy( if system_idx is None: return 0.5 * torch.sum(squared_term) flattened_squared = torch.sum(squared_term, dim=-1) - return 0.5 * torch.segment_reduce( - flattened_squared, reduce="sum", lengths=torch.bincount(system_idx) + n_systems = int(system_idx.max()) + 1 + return 0.5 * flattened_squared.new_zeros(n_systems).index_add( + 0, system_idx, flattened_squared ) @@ -266,14 +267,13 @@ def calc_heat_flux( Returns: Heat flux vector of shape (3,) or (n_systems, 3) """ - if momenta is not None and velocities is not None: - raise ValueError("Must pass either momenta or velocities, not both") - if momenta is None and velocities is None: - raise ValueError("Must pass either momenta or velocities") - # Deduce velocities if velocities is None: + if momenta is None: + raise ValueError("Must pass either momenta or velocities") velocities = momenta / masses.unsqueeze(-1) + elif momenta is not None: + raise ValueError("Must pass either momenta or velocities, not both") convective_flux = energies.unsqueeze(-1) * velocities diff --git a/torch_sim/state.py b/torch_sim/state.py index 4c8387d84..2f6319de0 100644 --- a/torch_sim/state.py +++ b/torch_sim/state.py @@ -247,9 +247,10 @@ def __setattr__(self, name: str, value: object) -> None: # noqa: C901 value = torch.zeros( self.n_atoms, device=self.device, dtype=torch.int64 ) - elif isinstance(value, torch.Tensor): - _, counts = torch.unique_consecutive(value, return_counts=True) - if not torch.all(counts == torch.bincount(value)): + elif isinstance(value, torch.Tensor) and value.numel() > 0: + # atoms must be grouped by system in ascending order: 0, ..., 0, 1, ... + steps = value.diff() + if value[0] != 0 or ((steps < 0) | (steps > 1)).any(): raise ValueError("System indices must be unique consecutive integers") super().__setattr__(name, value) @@ -628,7 +629,7 @@ def from_state(cls, state: "SimState", **additional_attrs: Any) -> Self: ... momenta=torch.zeros_like(sim_state.positions), ... ) """ - attrs = {} + attrs: dict[str, Any] = {} for attr_name, attr_value in state.attributes.items(): if attr_name in cls._get_all_attributes(): attrs[attr_name] = cls._clone_attr(attr_value) diff --git a/torch_sim/trajectory.py b/torch_sim/trajectory.py index ec4a11418..ce9377fe4 100644 --- a/torch_sim/trajectory.py +++ b/torch_sim/trajectory.py @@ -226,8 +226,9 @@ def truncate_to_step(self, step: int) -> None: """ if step <= 0: raise ValueError(f"Step must be greater than 0. Got step={step}.") - last_steps = self.last_written_steps - if any(s is None for s in last_steps): + all_last_steps = self.last_written_steps + last_steps = [s for s in all_last_steps if s is not None] + if len(last_steps) < len(all_last_steps): raise ValueError("Cannot truncate: one or more trajectories are empty.") if step > min(last_steps): raise ValueError(