From 58ddea3de64f7ec67f067a82be0d85aadaaf3689 Mon Sep 17 00:00:00 2001 From: James Kermode Date: Tue, 23 Jun 2026 12:52:17 +0100 Subject: [PATCH 1/4] Add experimental ASE neighbour-list plugin (host + device) Register Vesin under the `ase.plugins` entry point (ASE v4 plugin API): - vesin._ase_plugin.neighbor_list: host NeighborListFunction backend, a thin wrapper over vesin.ase_neighbor_list (rejects self_interaction=True). - vesin._ase_device.VesinDeviceNeighborList: optional device-resident capability implementing ASE's experimental DeviceNeighborList protocol over Vesin's CUDA cell list via its CuPy interface (build_device returns device-resident CuPy arrays; needs_rebuild is an on-device max-displacement reduction). COO only, so the padded path is unsupported; differentiable=False. Registration is guarded, so installing alongside an older ASE (no v4 plugin API) registers nothing rather than breaking discovery. Adds an `ase` optional extra and tests/test_ase_plugin.py (host always; device skipped without CuPy/GPU/ASE-v4). Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 7 ++ python/vesin/pyproject.toml | 5 + python/vesin/tests/test_ase_plugin.py | 146 +++++++++++++++++++++++ python/vesin/vesin/_ase_device.py | 164 ++++++++++++++++++++++++++ python/vesin/vesin/_ase_plugin.py | 65 ++++++++++ 5 files changed, 387 insertions(+) create mode 100644 python/vesin/tests/test_ase_plugin.py create mode 100644 python/vesin/vesin/_ase_device.py create mode 100644 python/vesin/vesin/_ase_plugin.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fdcd5e3b..15aa0f93 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,13 @@ changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows ### Removed --> +### Added + +- Experimental ASE neighbour-list plugin (`vesin._ase_plugin`): registers Vesin + under the `ase.plugins` entry point as a host `NeighborListFunction` backend, + plus an optional device-resident `DeviceNeighborList` capability (CUDA via + CuPy) for ASE's experimental device neighbour-list protocol. + ### Changed - Importing from `vesin.torch` is deprecated, users should now import from diff --git a/python/vesin/pyproject.toml b/python/vesin/pyproject.toml index 7f9f5882..b4fcb6ab 100644 --- a/python/vesin/pyproject.toml +++ b/python/vesin/pyproject.toml @@ -32,6 +32,11 @@ classifiers = [ [project.optional-dependencies] torch = ["vesin-torch"] +ase = ["ase"] + +# ASE v4 neighbour-list plugin discovery (registers nothing on older ASE). +[project.entry-points."ase.plugins"] +vesin = "vesin._ase_plugin" [project.urls] homepage = "https://github.com/Luthaf/vesin/" diff --git a/python/vesin/tests/test_ase_plugin.py b/python/vesin/tests/test_ase_plugin.py new file mode 100644 index 00000000..c8383b84 --- /dev/null +++ b/python/vesin/tests/test_ase_plugin.py @@ -0,0 +1,146 @@ +import ase.build +import ase.neighborlist +import numpy as np +import pytest + +from vesin._ase_plugin import device_neighbor_list, neighbor_list + + +def _ijS_set(i, j, S): + i = np.asarray(i, dtype=np.int64) + j = np.asarray(j, dtype=np.int64) + S = np.asarray(S, dtype=np.int64) + return set( + zip( + i.tolist(), + j.tolist(), + *[S[:, k].tolist() for k in range(3)], + strict=True, + ) + ) + + +def _system(): + atoms = ase.build.bulk("Cu", "fcc", a=3.6, cubic=True).repeat((2, 2, 2)) + return atoms + + +# --------------------------------------------------------------------------- # +# Host plugin (no GPU) +# --------------------------------------------------------------------------- # +def test_host_matches_ase(): + atoms = _system() + cutoff = 4.0 + vi, vj, vd, vD, vS = neighbor_list("ijdDS", atoms, cutoff) + ai, aj, ad, aD, aS = ase.neighborlist.neighbor_list("ijdDS", atoms, cutoff) + assert len(vi) == len(ai) + assert _ijS_set(vi, vj, vS) == _ijS_set(ai, aj, aS) + + +def test_host_self_interaction_rejected(): + atoms = _system() + with pytest.raises(NotImplementedError): + neighbor_list("ij", atoms, 4.0, self_interaction=True) + + +def test_plugin_registration(): + # Exposes __ase_plugins__; the "vesin" plugin is present iff ASE has the v4 + # plugin API (older ASE -> empty set, by design). + from vesin import _ase_plugin + + try: + from ase._4.plugins.neighborlist import NeighborListPlugin # noqa: F401 + except ImportError: + assert _ase_plugin.__ase_plugins__ == set() + else: + names = {p.name for p in _ase_plugin.__ase_plugins__} + assert "vesin" in names + + +# --------------------------------------------------------------------------- # +# Device backend (CUDA via CuPy + ASE v4 device protocol) +# --------------------------------------------------------------------------- # +def _device_skip_reason(): + try: + import cupy as cp + except ImportError as exc: + return f"cupy not available: {exc}" + try: + cp.cuda.Device(0).compute_capability + except Exception as exc: # pragma: no cover - env dependent + return f"CUDA not available: {exc}" + try: + import ase._4.plugins.neighborlist_device # noqa: F401 + except ImportError: + return "ASE has no experimental device neighbour-list protocol" + return None + + +_DEVICE_SKIP = _device_skip_reason() +device = pytest.mark.skipif(_DEVICE_SKIP is not None, reason=str(_DEVICE_SKIP)) + + +@device +def test_device_satisfies_protocol(): + from ase._4.plugins.neighborlist_device import DeviceNeighborList + + be = device_neighbor_list() + assert isinstance(be, DeviceNeighborList) + assert be.differentiable is False + assert be.device[0] == 2 # CUDA + + +@device +def test_device_equivalence_vs_ase(): + import cupy as cp + + atoms = _system() + cutoff = 4.0 + res = device_neighbor_list().build_device( + cp.asarray(atoms.positions), + np.asarray(atoms.cell[:]), + tuple(bool(b) for b in atoms.pbc), + cutoff, + "ijSD", + ) + vi = cp.asnumpy(res.get("i")) + vj = cp.asnumpy(res.get("j")) + vS = cp.asnumpy(res.get("S")) + vD = cp.asnumpy(res.get("D")) + assert res.get("i").__dlpack_device__()[0] == 2 # device-resident + + ai, aj, aS, aD = ase.neighborlist.neighbor_list("ijSD", atoms, cutoff) + assert _ijS_set(vi, vj, vS) == _ijS_set(ai, aj, aS) + + +@device +def test_device_needs_rebuild_scalar(): + import cupy as cp + + atoms = _system() + base = atoms.positions + be = device_neighbor_list() + be.build_device( + cp.asarray(base), np.asarray(atoms.cell[:]), + tuple(bool(b) for b in atoms.pbc), 4.0, "ijS", + ) + skin = 0.4 + nr = be.needs_rebuild(cp.asarray(base), skin=skin) + assert nr.__dlpack_device__()[0] == 2 + assert bool(nr) is False + moved = base.copy() + moved[0, 0] += 1.0 + assert bool(be.needs_rebuild(cp.asarray(moved), skin=skin)) is True + + +@device +def test_device_padded_unsupported(): + import cupy as cp + + atoms = _system() + be = device_neighbor_list() + with pytest.raises(NotImplementedError): + be.build_device( + cp.asarray(atoms.positions), np.asarray(atoms.cell[:]), + tuple(bool(b) for b in atoms.pbc), 4.0, max_capacity=64, + ) diff --git a/python/vesin/vesin/_ase_device.py b/python/vesin/vesin/_ase_device.py new file mode 100644 index 00000000..3365bc80 --- /dev/null +++ b/python/vesin/vesin/_ase_device.py @@ -0,0 +1,164 @@ +"""Experimental: device-resident neighbour-list adapter for ASE. + +Adapts Vesin's CUDA cell list (reached through its CuPy array interface) to ASE's +experimental device neighbour-list protocol +(:mod:`ase._4.plugins.neighborlist_device`): :class:`DeviceNeighborList` and +:class:`DeviceNeighborResult`. A device-resident calculator discovers it via +``isinstance(backend, DeviceNeighborList)`` and exchanges edge data on-device via +DLPack, with no host round-trip per step. + +GPU-only and requires CuPy: passing CuPy device positions to +:meth:`vesin.NeighborList.compute` dispatches to the CUDA kernel and returns +device-resident CuPy arrays. The host path is unchanged (``vesin.ase_neighbor_list``). + +Notes / limitations +------------------- +* ``differentiable = False`` (the CuPy path is not an autograd graph; Vesin's + *torch* binding is autograd-differentiable but computes on CPU, so it does not + satisfy this device protocol). +* No fixed-capacity / dense output -- Vesin returns COO only, so the padded + (``max_capacity``) path is unsupported and raises. +* ``needs_rebuild`` is a CuPy max-squared-displacement reduction (device-resident + 0-d bool, no host sync inside the call). Vesin's ``NeighborList`` also has a + native ``skin=`` Verlet reuse, but that does not expose an on-device rebuild + scalar, so the protocol uses the explicit reduction here. +* Scalar cutoff only; ``self_interaction=True`` rejected -- both raise. +""" + +import numpy as np + + +def _cupy(): + import cupy as cp # imported lazily; this adapter is GPU-only + + return cp + + +def _to_device_cell(cell): + """Return the (3, 3) cell as a CuPy array (Vesin needs box and points to be + the same array type). The cell is tiny, so any host->device copy here is not + the residency concern (that is the O(n_atoms) positions / O(n_edges) edges).""" + cp = _cupy() + if isinstance(cell, cp.ndarray): + return cell + to_dlpack_device = getattr(cell, "__dlpack_device__", None) + if to_dlpack_device is not None and to_dlpack_device()[0] != 1: # device array + return cp.from_dlpack(cell) + return cp.asarray(np.ascontiguousarray(np.asarray(cell), dtype=float)) + + +class VesinDeviceResult: + """On-device neighbour data (CuPy); satisfies ``DeviceNeighborResult``. + + Vesin returns COO arrays only, so this result is never padded. + """ + + def __init__(self, arrays, *, n_edges): + self._arrays = arrays # name -> CuPy device array (DLPack-exporting) + self._n_edges = n_edges + + @property + def n_edges(self): + return int(self._n_edges) + + @property + def did_overflow(self): + return False + + @property + def padded(self): + return False + + def get(self, quantity): + try: + return self._arrays[quantity] + except KeyError: + available = ", ".join(sorted(self._arrays)) or "(none)" + raise KeyError( + f"quantity {quantity!r} not available; have {available}. " + "It was not among the requested quantities." + ) from None + + def mask(self): + return None + + +class VesinDeviceNeighborList: + """Vesin device backend; satisfies ``DeviceNeighborList``.""" + + differentiable = False + + def __init__(self, device_id=0): + self._device_type = 2 # kDLCUDA (Vesin's GPU backend is CUDA) + self._device_id = int(device_id) + self._ref = None # cached build-time positions (device) + + @property + def device(self): + return (self._device_type, self._device_id) + + def build_device( + self, + positions, + cell, + pbc, + cutoff, + quantities="ijS", + *, + self_interaction=False, + max_capacity=None, + stream=None, + ): + if self_interaction: + raise NotImplementedError( + "the Vesin backend does not support self_interaction=True" + ) + if isinstance(cutoff, dict) or np.ndim(cutoff) != 0: + raise NotImplementedError( + "the Vesin device backend supports only a scalar cutoff" + ) + if max_capacity is not None: + raise NotImplementedError( + "Vesin has no fixed-capacity (dense) output; use the tight path " + "(max_capacity=None) for COO i/j/S/D" + ) + invalid = set(quantities) - set("ijdDS") + if invalid or not quantities: + raise ValueError( + f'quantities must be a non-empty subset of "ijdDS"; ' + f"got {quantities!r}." + ) + + from ._neighbors import NeighborList + + cp = _cupy() + pos = cp.from_dlpack(positions) # device array (zero-copy view) + self._device_id = int(pos.device.id) + # Snapshot build-time positions on device for needs_rebuild (the consumer + # may overwrite its positions buffer in place each step). + self._ref = pos.copy() + box = _to_device_cell(cell) + pbc_t = tuple(bool(b) for b in pbc) + + calculator = NeighborList(cutoff=float(cutoff), full_list=True) + out = calculator.compute(pos, box, pbc_t, quantities=quantities) + if not isinstance(out, (list, tuple)): # single-quantity -> bare array + out = (out,) + # Keyed by NAME (Vesin returns in requested order; we never rely on it). + arrays = {name: arr for name, arr in zip(quantities, out)} + n_edges = int(arrays[quantities[0]].shape[0]) + return VesinDeviceResult(arrays, n_edges=n_edges) + + def needs_rebuild(self, positions, *, skin, stream=None): + """On-device 0-d bool: max squared displacement > skin**2. + + CuPy reduction (device-resident); the result stays on device -- no host + sync inside the call. The eager ``bool()`` (CuPy's native ``__bool__``) + is the single documented sync point; a compiled consumer adopts the + scalar via ``from_dlpack`` and branches in-graph. + """ + if self._ref is None: + raise RuntimeError("call build_device(...) before needs_rebuild(...)") + cp = _cupy() + cur = cp.from_dlpack(positions) + return cp.max(((cur - self._ref) ** 2).sum(axis=1)) > float(skin) ** 2 diff --git a/python/vesin/vesin/_ase_plugin.py b/python/vesin/vesin/_ase_plugin.py new file mode 100644 index 00000000..bf4c1f41 --- /dev/null +++ b/python/vesin/vesin/_ase_plugin.py @@ -0,0 +1,65 @@ +"""ASE neighbour-list plugin for Vesin (host + experimental device capability). + +Registers Vesin as an ``ase.plugins`` neighbour-list backend so it can be +selected via :func:`ase.neighborlist.get_neighbor_list` (ASE v4 plugin API). +Selection is never automatic -- this only makes the backend *available* under the +name ``"vesin"``. + +This module is the ``ase.plugins`` entry point (it exposes ``__ase_plugins__``). +Registration is guarded so installing Vesin alongside an ASE that predates the v4 +plugin API registers nothing rather than breaking plugin discovery. The heavy +imports (ase, CuPy) happen lazily, inside the adapters. +""" +from __future__ import annotations + + +def neighbor_list(quantities, atoms, cutoff, *, self_interaction=False): + """Vesin neighbour list adapted to ASE's ``NeighborListFunction`` protocol. + + Thin wrapper over :func:`vesin.ase_neighbor_list`, returning the same + ``(i, j, d, D, S)`` flat arrays and quantity letters as + :func:`ase.neighborlist.neighbor_list`. Vesin has no ``self_interaction`` + option (it never returns pure self-pairs); per the plugin contract a request + it cannot honour is rejected rather than silently differing. + """ + if self_interaction: + raise NotImplementedError( + "the vesin backend does not support self_interaction=True" + ) + from ._ase import ase_neighbor_list + + return ase_neighbor_list(quantities, atoms, cutoff) + + +def device_neighbor_list(device_id=0): + """Return the *experimental* device-resident Vesin backend. + + The result satisfies ASE's experimental ``DeviceNeighborList`` protocol + (:mod:`ase._4.plugins.neighborlist_device`) and builds edge data on-device + (exchanged via DLPack). GPU-only (CUDA) and requires CuPy. + + The device capability is separate from the host ``neighbor_list`` + registration: the host ``NeighborListPlugin`` carries no device slot, so a + consumer obtains the device backend through this factory and checks + ``isinstance(backend, DeviceNeighborList)``. + """ + from ._ase_device import VesinDeviceNeighborList + + return VesinDeviceNeighborList(device_id=device_id) + + +try: + from ase._4.plugins.neighborlist import NeighborListPlugin +except ImportError: + # ASE without the v4 plugin API: register nothing, but do not break the + # discovery of other plugins. + __ase_plugins__: set = set() +else: + __ase_plugins__ = { + NeighborListPlugin( + "vesin", + long_name="Vesin neighbour list", + citation="https://github.com/Luthaf/vesin", + implementation="vesin._ase_plugin.neighbor_list", + ), + } From 2018b1f79240c5d5999c5a1e021e808ebbe00297 Mon Sep 17 00:00:00 2001 From: James Kermode Date: Tue, 23 Jun 2026 14:18:26 +0100 Subject: [PATCH 2/4] Fix docstring per review: vesin-torch is GPU-capable; clarify needs_rebuild - Correct the differentiable note: differentiable=False is only because this adapter uses the CuPy path (not an autograd framework), not a vesin-torch limitation -- vesin-torch runs on GPU and is autograd-differentiable (would make a natural differentiable=True device backend). - Note dense/fixed-capacity output is in progress upstream. - Clarify needs_rebuild: it returns an on-device boolean *scalar* for a compiled consumer to branch on; vesin already rebuilds on device via skin= but doesn't expose that decision as a queryable scalar. Co-Authored-By: Claude Opus 4.8 --- python/vesin/vesin/_ase_device.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/python/vesin/vesin/_ase_device.py b/python/vesin/vesin/_ase_device.py index 3365bc80..481c4a03 100644 --- a/python/vesin/vesin/_ase_device.py +++ b/python/vesin/vesin/_ase_device.py @@ -13,15 +13,20 @@ Notes / limitations ------------------- -* ``differentiable = False`` (the CuPy path is not an autograd graph; Vesin's - *torch* binding is autograd-differentiable but computes on CPU, so it does not - satisfy this device protocol). -* No fixed-capacity / dense output -- Vesin returns COO only, so the padded - (``max_capacity``) path is unsupported and raises. -* ``needs_rebuild`` is a CuPy max-squared-displacement reduction (device-resident - 0-d bool, no host sync inside the call). Vesin's ``NeighborList`` also has a - native ``skin=`` Verlet reuse, but that does not expose an on-device rebuild - scalar, so the protocol uses the explicit reduction here. +* ``differentiable = False`` here only because this adapter goes through the + **CuPy** path, which is not an autograd framework -- not a Vesin limitation. + Vesin's torch binding (``vesin-torch``) runs on GPU and is autograd- + differentiable, so it would make a natural ``differentiable = True`` device + backend for this protocol (a useful follow-up). +* No fixed-capacity / dense output -- Vesin currently returns COO only, so the + padded (``max_capacity``) path is unsupported and raises. (Dense output is in + progress upstream; the padded path can map onto it once available.) +* ``needs_rebuild`` returns an on-device boolean scalar (a CuPy max-squared- + displacement reduction; no host sync inside the call) so a compiled consumer + can branch on rebuild-vs-reuse without leaving the device. Vesin's + ``NeighborList`` already rebuilds on device as needed via its ``skin=`` Verlet + reuse, but does not expose that decision as a queryable scalar; if it did, this + adapter would use it directly instead of the explicit reduction. * Scalar cutoff only; ``self_interaction=True`` rejected -- both raise. """ From c53e0c09bb18436894d1b2da146a50cd816947f0 Mon Sep 17 00:00:00 2001 From: James Kermode Date: Tue, 23 Jun 2026 14:37:32 +0100 Subject: [PATCH 3/4] Fix lint: ruff format + strict= on zip + drop unused var Address the CI `lint` env (ruff): apply ruff format to the new files, add strict=True to the zip in build_device (B905), and remove an unused variable in the device equivalence test (F841). Tests unchanged (7 passed on GPU). Co-Authored-By: Claude Opus 4.8 --- python/vesin/tests/test_ase_plugin.py | 15 ++++++++++----- python/vesin/vesin/_ase_device.py | 5 ++--- python/vesin/vesin/_ase_plugin.py | 1 + 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/python/vesin/tests/test_ase_plugin.py b/python/vesin/tests/test_ase_plugin.py index c8383b84..0afec5be 100644 --- a/python/vesin/tests/test_ase_plugin.py +++ b/python/vesin/tests/test_ase_plugin.py @@ -106,7 +106,6 @@ def test_device_equivalence_vs_ase(): vi = cp.asnumpy(res.get("i")) vj = cp.asnumpy(res.get("j")) vS = cp.asnumpy(res.get("S")) - vD = cp.asnumpy(res.get("D")) assert res.get("i").__dlpack_device__()[0] == 2 # device-resident ai, aj, aS, aD = ase.neighborlist.neighbor_list("ijSD", atoms, cutoff) @@ -121,8 +120,11 @@ def test_device_needs_rebuild_scalar(): base = atoms.positions be = device_neighbor_list() be.build_device( - cp.asarray(base), np.asarray(atoms.cell[:]), - tuple(bool(b) for b in atoms.pbc), 4.0, "ijS", + cp.asarray(base), + np.asarray(atoms.cell[:]), + tuple(bool(b) for b in atoms.pbc), + 4.0, + "ijS", ) skin = 0.4 nr = be.needs_rebuild(cp.asarray(base), skin=skin) @@ -141,6 +143,9 @@ def test_device_padded_unsupported(): be = device_neighbor_list() with pytest.raises(NotImplementedError): be.build_device( - cp.asarray(atoms.positions), np.asarray(atoms.cell[:]), - tuple(bool(b) for b in atoms.pbc), 4.0, max_capacity=64, + cp.asarray(atoms.positions), + np.asarray(atoms.cell[:]), + tuple(bool(b) for b in atoms.pbc), + 4.0, + max_capacity=64, ) diff --git a/python/vesin/vesin/_ase_device.py b/python/vesin/vesin/_ase_device.py index 481c4a03..4aaba7c9 100644 --- a/python/vesin/vesin/_ase_device.py +++ b/python/vesin/vesin/_ase_device.py @@ -130,8 +130,7 @@ def build_device( invalid = set(quantities) - set("ijdDS") if invalid or not quantities: raise ValueError( - f'quantities must be a non-empty subset of "ijdDS"; ' - f"got {quantities!r}." + f'quantities must be a non-empty subset of "ijdDS"; got {quantities!r}.' ) from ._neighbors import NeighborList @@ -150,7 +149,7 @@ def build_device( if not isinstance(out, (list, tuple)): # single-quantity -> bare array out = (out,) # Keyed by NAME (Vesin returns in requested order; we never rely on it). - arrays = {name: arr for name, arr in zip(quantities, out)} + arrays = {name: arr for name, arr in zip(quantities, out, strict=True)} n_edges = int(arrays[quantities[0]].shape[0]) return VesinDeviceResult(arrays, n_edges=n_edges) diff --git a/python/vesin/vesin/_ase_plugin.py b/python/vesin/vesin/_ase_plugin.py index bf4c1f41..aaaaf4b5 100644 --- a/python/vesin/vesin/_ase_plugin.py +++ b/python/vesin/vesin/_ase_plugin.py @@ -10,6 +10,7 @@ plugin API registers nothing rather than breaking plugin discovery. The heavy imports (ase, CuPy) happen lazily, inside the adapters. """ + from __future__ import annotations From 212d0889dd2a9ec08d0f3d048cee6be960ff350e Mon Sep 17 00:00:00 2001 From: James Kermode Date: Tue, 23 Jun 2026 16:32:46 +0100 Subject: [PATCH 4/4] Device adapter: delegate reuse to a persistent NeighborList(skin=) Make the device backend stateful so Vesin's own Verlet skin reuse spans steps instead of reimplementing a displacement check in the adapter: - VesinDeviceNeighborList(skin=) holds a persistent NeighborList, rebuilt only when the cutoff changes; build_device calls .compute on it so Vesin decides rebuild-vs-reuse internally (and gains NPT box-change reuse once #172 lands). - needs_rebuild now returns a device-resident True (a no-op signal): the consumer always calls build_device and Vesin makes reuse cheap. - device_neighbor_list(skin=) plumbs the skin through the factory. Tests updated: needs_rebuild delegation + a skin>0 reuse-correctness check (reused list still matches a fresh build, since Vesin filters to the true cutoff on reuse). Co-Authored-By: Claude Opus 4.8 --- python/vesin/tests/test_ase_plugin.py | 43 ++++++++++++++----- python/vesin/vesin/_ase_device.py | 59 ++++++++++++++++----------- python/vesin/vesin/_ase_plugin.py | 7 +++- 3 files changed, 74 insertions(+), 35 deletions(-) diff --git a/python/vesin/tests/test_ase_plugin.py b/python/vesin/tests/test_ase_plugin.py index 0afec5be..468b1c07 100644 --- a/python/vesin/tests/test_ase_plugin.py +++ b/python/vesin/tests/test_ase_plugin.py @@ -113,26 +113,49 @@ def test_device_equivalence_vs_ase(): @device -def test_device_needs_rebuild_scalar(): +def test_device_needs_rebuild_delegates(): + # Reuse is delegated to Vesin's persistent NeighborList(skin=); needs_rebuild + # is a no-op device-resident True (the consumer always calls build_device). import cupy as cp atoms = _system() - base = atoms.positions be = device_neighbor_list() be.build_device( - cp.asarray(base), + cp.asarray(atoms.positions), np.asarray(atoms.cell[:]), tuple(bool(b) for b in atoms.pbc), 4.0, "ijS", ) - skin = 0.4 - nr = be.needs_rebuild(cp.asarray(base), skin=skin) - assert nr.__dlpack_device__()[0] == 2 - assert bool(nr) is False - moved = base.copy() - moved[0, 0] += 1.0 - assert bool(be.needs_rebuild(cp.asarray(moved), skin=skin)) is True + nr = be.needs_rebuild(cp.asarray(atoms.positions), skin=0.4) + assert nr.__dlpack_device__()[0] == 2 # device-resident + assert bool(nr) is True + + +@device +def test_device_skin_reuse_correct(): + # With skin>0 the persistent calculator reuses the list across a sub-skin + # move; the result must still match a fresh build (Vesin filters to the true + # cutoff on reuse). + import cupy as cp + + from vesin._ase_device import VesinDeviceNeighborList + + atoms = _system() + cell = np.asarray(atoms.cell[:]) + pbc = tuple(bool(b) for b in atoms.pbc) + be = VesinDeviceNeighborList(skin=0.5) + be.build_device(cp.asarray(atoms.positions), cell, pbc, 4.0, "ijS") # seed cache + moved = atoms.positions + 0.05 # sub-skin displacement -> reuse path + res = be.build_device(cp.asarray(moved), cell, pbc, 4.0, "ijS") + vi = cp.asnumpy(res.get("i")) + vj = cp.asnumpy(res.get("j")) + vS = cp.asnumpy(res.get("S")) + + a2 = atoms.copy() + a2.positions = moved + ai, aj, aS = ase.neighborlist.neighbor_list("ijS", a2, 4.0) + assert _ijS_set(vi, vj, vS) == _ijS_set(ai, aj, aS) @device diff --git a/python/vesin/vesin/_ase_device.py b/python/vesin/vesin/_ase_device.py index 4aaba7c9..a49b8342 100644 --- a/python/vesin/vesin/_ase_device.py +++ b/python/vesin/vesin/_ase_device.py @@ -21,12 +21,16 @@ * No fixed-capacity / dense output -- Vesin currently returns COO only, so the padded (``max_capacity``) path is unsupported and raises. (Dense output is in progress upstream; the padded path can map onto it once available.) -* ``needs_rebuild`` returns an on-device boolean scalar (a CuPy max-squared- - displacement reduction; no host sync inside the call) so a compiled consumer - can branch on rebuild-vs-reuse without leaving the device. Vesin's - ``NeighborList`` already rebuilds on device as needed via its ``skin=`` Verlet - reuse, but does not expose that decision as a queryable scalar; if it did, this - adapter would use it directly instead of the explicit reduction. +* Reuse is delegated to Vesin: the adapter holds a **persistent** + ``NeighborList(skin=...)`` and ``build_device`` calls ``.compute`` on it, so + Vesin's own Verlet logic reuses the list across steps (displacement, and box + changes once Luthaf/vesin#172 lands -- NPT). ``needs_rebuild`` therefore just + returns a device-resident ``True`` (a no-op signal): the consumer always calls + ``build_device`` and Vesin makes the reuse cheap internally. The Verlet skin is + set at construction (``VesinDeviceNeighborList(skin=...)``), matching Vesin's + API; ``needs_rebuild``'s ``skin`` argument is unused. (A compiled consumer that + wants to *skip* the build itself would instead need an explicit device-scalar + displacement check -- an open ASE-protocol design point.) * Scalar cutoff only; ``self_interaction=True`` rejected -- both raise. """ @@ -93,10 +97,14 @@ class VesinDeviceNeighborList: differentiable = False - def __init__(self, device_id=0): + def __init__(self, device_id=0, skin=0.0): self._device_type = 2 # kDLCUDA (Vesin's GPU backend is CUDA) self._device_id = int(device_id) - self._ref = None # cached build-time positions (device) + # Persistent calculator so Vesin's own Verlet ``skin`` reuse kicks in + # across steps. skin=0 means rebuild every call (no reuse). + self._skin = float(skin) + self._nl = None + self._nl_cutoff = None @property def device(self): @@ -138,14 +146,18 @@ def build_device( cp = _cupy() pos = cp.from_dlpack(positions) # device array (zero-copy view) self._device_id = int(pos.device.id) - # Snapshot build-time positions on device for needs_rebuild (the consumer - # may overwrite its positions buffer in place each step). - self._ref = pos.copy() box = _to_device_cell(cell) pbc_t = tuple(bool(b) for b in pbc) - calculator = NeighborList(cutoff=float(cutoff), full_list=True) - out = calculator.compute(pos, box, pbc_t, quantities=quantities) + # Reuse one persistent NeighborList so Vesin's Verlet ``skin`` reuse spans + # calls (rebuilt only if the cutoff changes). Vesin decides internally + # whether to rebuild or reuse on each ``.compute``. + if self._nl is None or self._nl_cutoff != float(cutoff): + self._nl = NeighborList( + cutoff=float(cutoff), full_list=True, skin=self._skin + ) + self._nl_cutoff = float(cutoff) + out = self._nl.compute(pos, box, pbc_t, quantities=quantities) if not isinstance(out, (list, tuple)): # single-quantity -> bare array out = (out,) # Keyed by NAME (Vesin returns in requested order; we never rely on it). @@ -154,15 +166,16 @@ def build_device( return VesinDeviceResult(arrays, n_edges=n_edges) def needs_rebuild(self, positions, *, skin, stream=None): - """On-device 0-d bool: max squared displacement > skin**2. - - CuPy reduction (device-resident); the result stays on device -- no host - sync inside the call. The eager ``bool()`` (CuPy's native ``__bool__``) - is the single documented sync point; a compiled consumer adopts the - scalar via ``from_dlpack`` and branches in-graph. + """Delegate reuse to Vesin: return a device-resident ``True`` scalar. + + Vesin's persistent ``NeighborList(skin=...)`` owns the rebuild-vs-reuse + decision internally (displacement, and box changes once Luthaf/vesin#172 + lands), so the adapter always signals "build" and lets ``build_device``'s + persistent calculator reuse the list cheaply. The ``skin`` argument is not + used here -- the Verlet skin is set at construction + (``VesinDeviceNeighborList(skin=...)``), matching Vesin's API. (A compiled + consumer that needs to *skip* the build itself would instead want an + explicit device-scalar displacement check; see the ASE protocol notes.) """ - if self._ref is None: - raise RuntimeError("call build_device(...) before needs_rebuild(...)") cp = _cupy() - cur = cp.from_dlpack(positions) - return cp.max(((cur - self._ref) ** 2).sum(axis=1)) > float(skin) ** 2 + return cp.asarray(True) diff --git a/python/vesin/vesin/_ase_plugin.py b/python/vesin/vesin/_ase_plugin.py index aaaaf4b5..26086741 100644 --- a/python/vesin/vesin/_ase_plugin.py +++ b/python/vesin/vesin/_ase_plugin.py @@ -32,13 +32,16 @@ def neighbor_list(quantities, atoms, cutoff, *, self_interaction=False): return ase_neighbor_list(quantities, atoms, cutoff) -def device_neighbor_list(device_id=0): +def device_neighbor_list(device_id=0, skin=0.0): """Return the *experimental* device-resident Vesin backend. The result satisfies ASE's experimental ``DeviceNeighborList`` protocol (:mod:`ase._4.plugins.neighborlist_device`) and builds edge data on-device (exchanged via DLPack). GPU-only (CUDA) and requires CuPy. + ``skin`` sets the Verlet skin of the backend's persistent ``NeighborList`` so + it reuses the list across steps (``skin=0`` rebuilds every call). + The device capability is separate from the host ``neighbor_list`` registration: the host ``NeighborListPlugin`` carries no device slot, so a consumer obtains the device backend through this factory and checks @@ -46,7 +49,7 @@ def device_neighbor_list(device_id=0): """ from ._ase_device import VesinDeviceNeighborList - return VesinDeviceNeighborList(device_id=device_id) + return VesinDeviceNeighborList(device_id=device_id, skin=skin) try: