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..468b1c07 --- /dev/null +++ b/python/vesin/tests/test_ase_plugin.py @@ -0,0 +1,174 @@ +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")) + 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_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() + be = device_neighbor_list() + be.build_device( + cp.asarray(atoms.positions), + np.asarray(atoms.cell[:]), + tuple(bool(b) for b in atoms.pbc), + 4.0, + "ijS", + ) + 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 +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..a49b8342 --- /dev/null +++ b/python/vesin/vesin/_ase_device.py @@ -0,0 +1,181 @@ +"""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`` 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.) +* 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. +""" + +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, skin=0.0): + self._device_type = 2 # kDLCUDA (Vesin's GPU backend is CUDA) + self._device_id = int(device_id) + # 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): + 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"; 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) + box = _to_device_cell(cell) + pbc_t = tuple(bool(b) for b in pbc) + + # 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). + 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) + + def needs_rebuild(self, positions, *, skin, stream=None): + """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.) + """ + cp = _cupy() + return cp.asarray(True) diff --git a/python/vesin/vesin/_ase_plugin.py b/python/vesin/vesin/_ase_plugin.py new file mode 100644 index 00000000..26086741 --- /dev/null +++ b/python/vesin/vesin/_ase_plugin.py @@ -0,0 +1,69 @@ +"""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, 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 + ``isinstance(backend, DeviceNeighborList)``. + """ + from ._ase_device import VesinDeviceNeighborList + + return VesinDeviceNeighborList(device_id=device_id, skin=skin) + + +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", + ), + }