Bump CI actions and lint hooks, plus small correctness fixes

.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

.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

.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

.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

.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 }}

.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

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))]

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

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

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
 

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

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)

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],

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(