From 7944aeb71bc6912bfec5cfae2dc6eabd6178082e Mon Sep 17 00:00:00 2001 From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com> Date: Sat, 18 Jul 2026 11:33:58 +0900 Subject: [PATCH 1/5] fix(neighbors): require nvalchemi-toolkit-ops 0.4.0 and drop the naive pbc workaround --- pyproject.toml | 2 +- torch_sim/neighbors/alchemiops.py | 10 ---------- 2 files changed, 1 insertion(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1586e41e..36bcaa31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "h5py>=3.15.1", "numpy>=1.26,<3; python_version < '3.13'", "numpy>=2.3.2,<3; python_version >= '3.13'", - "nvalchemi-toolkit-ops[torch]>=0.3.1", + "nvalchemi-toolkit-ops[torch]>=0.4.0", "tables>=3.11.1", "torch>=2", "tqdm>=4.67", diff --git a/torch_sim/neighbors/alchemiops.py b/torch_sim/neighbors/alchemiops.py index 6c8914d9..6cc2710c 100644 --- a/torch_sim/neighbors/alchemiops.py +++ b/torch_sim/neighbors/alchemiops.py @@ -72,16 +72,6 @@ def alchemiops_nl_n2( if _batch_naive_neighbor_list is None: raise RuntimeError("nvalchemiops neighbor list is unavailable") - # Temporary fix: the naive kernel wraps on every axis ignoring pbc, so a - # non-zero cell breaks non-periodic and partial-pbc systems. Zeroing the cell - # makes the wrap a no-op, fixing fully non-periodic systems (e.g. molecules). - # Slabs / partial pbc are not covered yet and need the upstream fix. - # See https://github.com/TorchSim/torch-sim/issues/575 - non_periodic = ~pbc.any(dim=1) # [n_systems] - if non_periodic.any(): - cell = cell.clone() # avoid modifying the original - cell[non_periodic] = 0.0 - res = _batch_naive_neighbor_list( positions=positions, cutoff=cutoff, From f56b47082fbabf92ef02eca75182b22ff4367c6c Mon Sep 17 00:00:00 2001 From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com> Date: Sat, 18 Jul 2026 11:37:20 +0900 Subject: [PATCH 2/5] fix(neighbors): pass pbc per axis to vesin and fix the import for vesin 0.6 --- pyproject.toml | 2 +- torch_sim/neighbors/vesin.py | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 36bcaa31..29d2cb6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ test = [ "pytest>=8", "spglib>=2.6", ] -vesin = ["vesin[torch]>=0.5.3"] +vesin = ["vesin[torch]>=0.5.8"] io = ["ase>=3.26", "phonopy>=3.0.0", "pymatgen>=2026.3.23"] symmetry = ["moyopy>=0.7.8"] mace = ["mace-torch>=0.3.16"] diff --git a/torch_sim/neighbors/vesin.py b/torch_sim/neighbors/vesin.py index 733a25b3..4fb2cbb4 100644 --- a/torch_sim/neighbors/vesin.py +++ b/torch_sim/neighbors/vesin.py @@ -18,10 +18,14 @@ VesinNeighborList = None +# vesin 0.6.0 renamed the TorchScript module from vesin.torch to vesin_torch try: - from vesin.torch import NeighborList as VesinNeighborListTorch + from vesin_torch import NeighborList as VesinNeighborListTorch except ImportError: - VesinNeighborListTorch = None + try: + from vesin.torch import NeighborList as VesinNeighborListTorch + except ImportError: + VesinNeighborListTorch = None # ty: ignore[conflicting-declarations] VESIN_AVAILABLE = VesinNeighborList is not None VESIN_TORCHSCRIPT_AVAILABLE = VesinNeighborListTorch is not None @@ -111,13 +115,12 @@ def vesin_nl( positions_cpu = wrapped[system_mask].detach().cpu().to(dtype=torch.float64) cell_cpu = cell_sys.detach().cpu().to(dtype=torch.float64) periodic_cpu = pbc[sys_idx].detach().to(dtype=torch.bool).cpu() - periodic_bool = bool(torch.all(periodic_cpu).item()) # Only works on CPU and returns numpy arrays i, j, S = neighbor_list_fn.compute( points=positions_cpu, box=cell_cpu, - periodic=periodic_bool, + periodic=periodic_cpu, quantities="ijS", ) i, j = ( From e2224ca7fdd575c103439fb0079730004efacff1 Mon Sep 17 00:00:00 2001 From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com> Date: Sat, 18 Jul 2026 11:47:27 +0900 Subject: [PATCH 3/5] test(neighbors): add partial pbc slab coverage and refactor dummy-cell molecules --- tests/test_neighbors.py | 46 +++++++++++++++++++++++++++-------------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py index 9e682ce4..6f947d25 100644 --- a/tests/test_neighbors.py +++ b/tests/test_neighbors.py @@ -6,7 +6,7 @@ import pytest import torch from ase import Atoms -from ase.build import bulk, molecule +from ase.build import bulk, fcc111, molecule from ase.neighborlist import neighbor_list from tests.conftest import DEVICE, DTYPE @@ -135,9 +135,32 @@ def periodic_atoms_set(): @pytest.fixture def molecule_atoms_set() -> list: - return [ + """Non-periodic molecules, each duplicated with a dummy cell set. + + A dummy cell arises e.g. from ase.Atoms.get_cell(complete=True). + """ + molecules = [ *map(molecule, ("CH3CH2NH2", "H2O", "methylenecyclopropane", "OCHCHO", "C3H9C")), ] + dummy_cell_molecules = [atoms.copy() for atoms in molecules] + for atoms in dummy_cell_molecules: + atoms.set_cell([1.0, 1.0, 1.0]) + return molecules + dummy_cell_molecules + + +@pytest.fixture +def partial_pbc_atoms_set() -> list[Atoms]: + """Slab structures to test partially periodic systems. + + To increase coverage, they include atoms drifted across the cell boundary + along the non-periodic axis. + """ + cu111 = fcc111("Cu", size=(2, 2, 3), a=3.6, vacuum=10.0) + cu111.translate([0.0, 0.0, -cu111.cell[2, 2] / 2]) + cu_slab = bulk("Cu", "fcc", a=3.6, cubic=True).repeat((2, 2, 2)) + cu_slab.set_pbc([True, True, False]) + cu_slab.translate([0.0, 0.0, -3.6]) + return [cu111, cu_slab] def _sorted_mic_distances( @@ -286,31 +309,22 @@ def test_neighbor_list_invariant_under_lattice_image_shifts( _nl_backends_x_cutoffs(), ) @pytest.mark.parametrize("self_interaction", [True, False]) -@pytest.mark.parametrize("molecule_dummy_cell", [0.0, 1.0]) def test_neighbor_list_implementations( *, nl_implementation: Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]], cutoff: float, self_interaction: bool, - molecule_dummy_cell: float, molecule_atoms_set: list[Atoms], periodic_atoms_set: list[Atoms], + partial_pbc_atoms_set: list[Atoms], ) -> None: """Check that neighbor list implementations give the same results as ASE. - Tests all implementations in batched mode with mixed periodic and non-periodic - systems, comparing sorted distances against ASE reference values. - - With molecule_dummy_cell != 0 the non-periodic molecules carry a non-zero - (dummy) cell, e.g. from ase.Atoms.get_cell(complete=True). Since pbc stays False - the neighbor list must remain cell-independent and still match ASE values. + Tests all implementations in batched mode with mixed periodic, partially + periodic, and non-periodic systems, comparing sorted distances against ASE + reference values. """ - molecules = molecule_atoms_set - if molecule_dummy_cell: - molecules = [atoms.copy() for atoms in molecule_atoms_set] - for atoms in molecules: - atoms.set_cell([molecule_dummy_cell] * 3) # dummy cell; pbc stays False - atoms_list = molecules + periodic_atoms_set + atoms_list = molecule_atoms_set + periodic_atoms_set + partial_pbc_atoms_set is_alchemiops = "alchemiops" in nl_implementation.__name__ if is_alchemiops and cutoff >= 3: atoms_list = [a for a in atoms_list if not a.info.get("very_skewed")] From 5d590790ce01cac3c7509528a642d5a2822b4bbb Mon Sep 17 00:00:00 2001 From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com> Date: Sat, 18 Jul 2026 12:16:05 +0900 Subject: [PATCH 4/5] refactor(neighbors): import nvalchemiops from its public path and drop the fallback --- torch_sim/neighbors/alchemiops.py | 25 ++++++------------------- 1 file changed, 6 insertions(+), 19 deletions(-) diff --git a/torch_sim/neighbors/alchemiops.py b/torch_sim/neighbors/alchemiops.py index 6cc2710c..cf8961ee 100644 --- a/torch_sim/neighbors/alchemiops.py +++ b/torch_sim/neighbors/alchemiops.py @@ -1,9 +1,7 @@ """Alchemiops-based neighbor list implementations. -This module provides neighbor lists via nvalchemiops: prefer the PyTorch subtree -(``nvalchemiops.torch.neighbors``), typical for CUDA builds, and fall back to -``nvalchemiops.neighborlist`` when that import path is missing (CPU-oriented API -with the same call surface). Supports naive N^2 and cell-list algorithms. +This module provides neighbor lists via nvalchemiops +(``nvalchemiops.torch.neighbors``). Supports naive N^2 and cell-list algorithms. nvalchemiops is available at: https://github.com/NVIDIA/nvalchemiops """ @@ -18,23 +16,12 @@ def _import_nvalchemiops_batch_neighbors() -> tuple[object, object] | None: - """Return ``(batch_cell_list, batch_naive_neighbor_list)`` if a layout is importable. - - Tries ``nvalchemiops.torch.neighbors`` first (PyTorch tensors; usual GPU wheel). - On ``ImportError``, tries ``nvalchemiops.neighborlist`` — same API, CPU fallback - when the ``torch.neighbors`` subtree is absent. - """ + """Return ``(batch_cell_list, batch_naive_neighbor_list)`` if importable.""" try: - from nvalchemiops.torch.neighbors.batch_cell_list import batch_cell_list as bcl - from nvalchemiops.torch.neighbors.batch_naive import ( - batch_naive_neighbor_list as bnl, - ) + from nvalchemiops.torch.neighbors import batch_cell_list as bcl + from nvalchemiops.torch.neighbors import batch_naive_neighbor_list as bnl except (ImportError, RuntimeError): - try: - from nvalchemiops.neighborlist import batch_cell_list as bcl - from nvalchemiops.neighborlist import batch_naive_neighbor_list as bnl - except (ImportError, RuntimeError): - return None + return None return bcl, bnl From 74413a6fadfeca0ed8d9a7b6b88dea8c8b4590fa Mon Sep 17 00:00:00 2001 From: Ryuhei Okuno <69704776+lil-lon@users.noreply.github.com> Date: Sat, 18 Jul 2026 12:58:58 +0900 Subject: [PATCH 5/5] fix(neighbors): pass cutoff to VesinNeighborList as a keyword for vesin 0.6 --- torch_sim/neighbors/vesin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_sim/neighbors/vesin.py b/torch_sim/neighbors/vesin.py index 4fb2cbb4..f726f145 100644 --- a/torch_sim/neighbors/vesin.py +++ b/torch_sim/neighbors/vesin.py @@ -108,7 +108,7 @@ def vesin_nl( # Calculate neighbor list for this system neighbor_list_fn = VesinNeighborList( - (float(cutoff)), full_list=True, sorted=False + cutoff=float(cutoff), full_list=True, sorted=False ) # Convert tensors to CPU and float64 without gradients