fix(neighbors): partial pbc and imports

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

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

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
 
 
@@ -72,16 +59,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,

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
@@ -104,20 +108,19 @@ 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
             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 = (