diff --git a/README.md b/README.md index e3a83a688..37199f2da 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ To see Adaptive in action, try the [example notebook on Binder](https://mybinder - [:floppy_disk: Exporting Data](#floppy_disk-exporting-data) - [:test_tube: Implemented Algorithms](#test_tube-implemented-algorithms) - [:package: Installation](#package-installation) + - [Faster triangulation (optional)](#faster-triangulation-optional) - [:wrench: Development](#wrench-development) - [:books: Citing](#books-citing) - [:page_facing_up: Draft Paper](#page_facing_up-draft-paper) @@ -151,6 +152,17 @@ pip install "adaptive[notebook]" The `[notebook]` above will also install the optional dependencies for running `adaptive` inside a Jupyter notebook. +### Faster triangulation (optional) + +Installing the optional [adaptive-triangulation](https://github.com/python-adaptive/adaptive-triangulation) package makes `adaptive.LearnerND` significantly faster by replacing the pure-Python triangulation with a Rust implementation: + +```bash +pip install "adaptive[rust]" +``` + +No code changes are needed — the Rust backend is detected and used automatically. +To control the selection, pass `LearnerND(..., triangulation_backend="python" | "rust" | "auto")` per learner, or set the environment variable `ADAPTIVE_TRIANGULATION_BACKEND=python` to force the pure-Python implementation globally (or to `rust` to make a missing Rust backend an error instead of a silent fallback). + To use Adaptive in Jupyterlab, you need to install the following labextensions. ```bash diff --git a/adaptive/learner/learner1D.py b/adaptive/learner/learner1D.py index 685834222..289ced0d0 100644 --- a/adaptive/learner/learner1D.py +++ b/adaptive/learner/learner1D.py @@ -14,7 +14,7 @@ from adaptive.learner.base_learner import BaseLearner, uses_nth_neighbors from adaptive.learner.learnerND import volume -from adaptive.learner.triangulation import simplex_volume_in_embedding +from adaptive.learner.triangulation_backend import simplex_volume_in_embedding from adaptive.notebook_integration import ensure_holoviews from adaptive.types import Float, Int, Real from adaptive.utils import ( diff --git a/adaptive/learner/learner2D.py b/adaptive/learner/learner2D.py index 125fc055f..60107e146 100644 --- a/adaptive/learner/learner2D.py +++ b/adaptive/learner/learner2D.py @@ -13,7 +13,7 @@ from scipy.interpolate import CloughTocher2DInterpolator, LinearNDInterpolator from adaptive.learner.base_learner import BaseLearner -from adaptive.learner.triangulation import simplex_volume_in_embedding +from adaptive.learner.triangulation_backend import simplex_volume_in_embedding from adaptive.notebook_integration import ensure_holoviews from adaptive.types import Bool, Float, Real from adaptive.utils import ( diff --git a/adaptive/learner/learnerND.py b/adaptive/learner/learnerND.py index eeb94656c..30a3fb2b5 100644 --- a/adaptive/learner/learnerND.py +++ b/adaptive/learner/learnerND.py @@ -14,11 +14,12 @@ from sortedcontainers import SortedKeyList from adaptive.learner.base_learner import BaseLearner, uses_nth_neighbors -from adaptive.learner.triangulation import ( +from adaptive.learner.triangulation import fast_det +from adaptive.learner.triangulation_backend import ( Triangulation, circumsphere, - fast_det, point_in_simplex, + resolve_triangulation_class, simplex_volume_in_embedding, ) from adaptive.notebook_integration import ensure_holoviews, ensure_plotly @@ -274,6 +275,15 @@ class LearnerND(BaseLearner): If not provided, then a default is used, which uses the deviation from a linear estimate, as well as triangle area, to determine the loss. + anisotropic : bool, optional + If True, the triangulation is stretched along the local gradient + when choosing new points. Only works with scalar output. + triangulation_backend : str or type, optional + Which triangulation implementation to use: ``"auto"`` (default, + prefers the optional Rust-accelerated `adaptive-triangulation + `_ + package when it is installed), ``"python"``, ``"rust"``, or a + `~adaptive.learner.triangulation.Triangulation`-compatible class. Attributes @@ -308,7 +318,20 @@ class LearnerND(BaseLearner): children based on volume. """ - def __init__(self, func, bounds, loss_per_simplex=None, *, anisotropic=False): + # Class-level fallback so that learners restored from pickles made + # before `triangulation_backend` existed keep working. + _triangulation_class = Triangulation + + def __init__( + self, + func, + bounds, + loss_per_simplex=None, + *, + anisotropic=False, + triangulation_backend="auto", + ): + self._triangulation_class = resolve_triangulation_class(triangulation_backend) self._vdim = None self.loss_per_simplex = loss_per_simplex or default_loss @@ -385,6 +408,7 @@ def new(self) -> LearnerND: self.bounds, self.loss_per_simplex, anisotropic=self.anisotropic, + triangulation_backend=self._triangulation_class, ) @property @@ -513,7 +537,7 @@ def tri(self): return self._tri try: - self._tri = Triangulation(self.points) + self._tri = self._triangulation_class(self.points) except ValueError: # A ValueError is raised if we do not have enough points or # the provided points are coplanar, so we need more points to @@ -649,7 +673,7 @@ def _try_adding_pending_point_to_simplex(self, point, simplex): if simplex not in self._subtriangulations: vertices = self.tri.get_vertices(simplex) - self._subtriangulations[simplex] = Triangulation(vertices) + self._subtriangulations[simplex] = self._triangulation_class(vertices) self._pending_to_simplex[point] = simplex return self._subtriangulations[simplex].add_point(point) @@ -713,7 +737,8 @@ def _pop_highest_existing_simplex(self): ): return abs(loss), simplex, subsimplex if ( - simplex in self._subtriangulations + subsimplex is not None + and simplex in self._subtriangulations and simplex in self.tri.simplices and subsimplex in self._subtriangulations[simplex].simplices ): diff --git a/adaptive/learner/triangulation_backend.py b/adaptive/learner/triangulation_backend.py new file mode 100644 index 000000000..e9246be41 --- /dev/null +++ b/adaptive/learner/triangulation_backend.py @@ -0,0 +1,147 @@ +"""Select the triangulation backend used by the learners. + +If the optional Rust-accelerated `adaptive-triangulation +`_ package is +installed (``pip install "adaptive[rust]"``), it is used automatically as a +drop-in replacement for the pure-Python implementation in +`adaptive.learner.triangulation`, which makes `~adaptive.LearnerND` +significantly faster. + +The selection can be overridden with the ``ADAPTIVE_TRIANGULATION_BACKEND`` +environment variable: + +- ``auto`` (default): use the Rust backend if available, else pure Python +- ``python``: always use the pure-Python implementation +- ``rust``: require the Rust backend, raising `ImportError` if it is missing + +The active backend is exposed as the string ``TRIANGULATION_BACKEND`` +(``"python"`` or ``"rust"``). + +Note that the pure-Python implementation in `adaptive.learner.triangulation` +is always importable under its own name, regardless of the selected backend, +so pickles that reference it keep working. +""" + +from __future__ import annotations + +import os + +# Minimal version that is a complete drop-in for the learners +# (incl. ``get_opposing_vertices`` and pickle/deepcopy support). +_MIN_RUST_VERSION = (0, 2, 1) + + +def _rust_version() -> tuple[int, ...] | None: + """Return the installed ``adaptive_triangulation`` version, or None.""" + try: + import adaptive_triangulation + except ImportError: + return None + version = adaptive_triangulation.__version__ + return tuple(int(part) for part in version.split(".")[:3] if part.isdigit()) + + +def _import_rust_triangulation(): + """Import the Rust `Triangulation`, raising a helpful `ImportError`.""" + version = _rust_version() + if version is None: + raise ImportError( + "The 'rust' triangulation backend was requested but the " + "'adaptive-triangulation' package is not installed. " + 'Install it with: pip install "adaptive[rust]"' + ) + if version < _MIN_RUST_VERSION: + raise ImportError( + "The 'rust' triangulation backend requires " + f"adaptive-triangulation >= {'.'.join(map(str, _MIN_RUST_VERSION))}, " + f"found {'.'.join(map(str, version))}. Upgrade it with: " + 'pip install -U "adaptive[rust]"' + ) + from adaptive_triangulation import Triangulation + + return Triangulation + + +def _select_backend() -> str: + backend = os.environ.get("ADAPTIVE_TRIANGULATION_BACKEND", "auto").lower() + if backend not in ("auto", "python", "rust"): + raise ValueError( + f"ADAPTIVE_TRIANGULATION_BACKEND={backend!r} is invalid, " + "use 'auto', 'python', or 'rust'." + ) + if backend == "auto": + version = _rust_version() + return ( + "rust" if version is not None and version >= _MIN_RUST_VERSION else "python" + ) + if backend == "rust": + _import_rust_triangulation() # raise with guidance if unusable + return backend + + +def resolve_triangulation_class(backend="auto"): + """Return the `Triangulation` class to use for *backend*. + + Parameters + ---------- + backend : str or type + ``"auto"`` (the module-level default backend, which prefers the Rust + implementation when available), ``"python"``, ``"rust"``, or a + `Triangulation`-compatible class. + """ + if isinstance(backend, type): + return backend + if backend == "auto": + return Triangulation + if backend == "python": + from adaptive.learner.triangulation import Triangulation as tri_class + + return tri_class + if backend == "rust": + return _import_rust_triangulation() + raise ValueError( + f"Invalid triangulation backend {backend!r}, use 'auto', 'python', " + "'rust', or a Triangulation-compatible class." + ) + + +TRIANGULATION_BACKEND: str = _select_backend() + +if TRIANGULATION_BACKEND == "rust": + from adaptive_triangulation import ( + Triangulation, + circumsphere, + fast_2d_circumcircle, + fast_2d_point_in_simplex, + fast_3d_circumcircle, + fast_norm, + orientation, + point_in_simplex, + simplex_volume_in_embedding, + ) +else: + from adaptive.learner.triangulation import ( + Triangulation, + circumsphere, + fast_2d_circumcircle, + fast_2d_point_in_simplex, + fast_3d_circumcircle, + fast_norm, + orientation, + point_in_simplex, + simplex_volume_in_embedding, + ) + +__all__ = [ + "TRIANGULATION_BACKEND", + "Triangulation", + "resolve_triangulation_class", + "circumsphere", + "fast_2d_circumcircle", + "fast_2d_point_in_simplex", + "fast_3d_circumcircle", + "fast_norm", + "orientation", + "point_in_simplex", + "simplex_volume_in_embedding", +] diff --git a/adaptive/tests/unit/test_triangulation_backend.py b/adaptive/tests/unit/test_triangulation_backend.py new file mode 100644 index 000000000..eb42ad4fe --- /dev/null +++ b/adaptive/tests/unit/test_triangulation_backend.py @@ -0,0 +1,124 @@ +"""Tests for the automatic triangulation backend selection.""" + +import os +import subprocess +import sys + +import pytest + +from adaptive.learner import triangulation as python_triangulation +from adaptive.learner import triangulation_backend as backend +from adaptive.learner.triangulation_backend import ( + _MIN_RUST_VERSION, + _rust_version, +) + + +def _run(code, backend_env): + env = {**os.environ, "ADAPTIVE_TRIANGULATION_BACKEND": backend_env} + return subprocess.run( + [sys.executable, "-c", code], env=env, capture_output=True, text=True + ) + + +def rust_is_usable(): + version = _rust_version() + return version is not None and version >= _MIN_RUST_VERSION + + +def test_backend_matches_installation(): + if rust_is_usable(): + import adaptive_triangulation + + assert backend.TRIANGULATION_BACKEND == "rust" + assert backend.Triangulation is adaptive_triangulation.Triangulation + else: + assert backend.TRIANGULATION_BACKEND == "python" + assert backend.Triangulation is python_triangulation.Triangulation + + +def test_python_triangulation_is_never_shadowed(): + # Old pickles reference adaptive.learner.triangulation.Triangulation by + # qualified name, so the pure-Python class must stay importable as itself. + assert python_triangulation.Triangulation.__module__ == ( + "adaptive.learner.triangulation" + ) + + +def test_force_python_backend(): + code = ( + "from adaptive.learner import triangulation, triangulation_backend;" + "assert triangulation_backend.TRIANGULATION_BACKEND == 'python';" + "assert triangulation_backend.Triangulation is triangulation.Triangulation" + ) + result = _run(code, "python") + assert result.returncode == 0, result.stderr + + +def test_force_rust_backend(): + code = ( + "from adaptive.learner import triangulation_backend;" + "assert triangulation_backend.TRIANGULATION_BACKEND == 'rust'" + ) + result = _run(code, "rust") + if rust_is_usable(): + assert result.returncode == 0, result.stderr + else: + # Forcing the Rust backend without (a recent enough version of) + # adaptive-triangulation must raise a helpful ImportError. + assert result.returncode != 0 + assert "ImportError" in result.stderr + assert "adaptive-triangulation" in result.stderr + + +def test_invalid_backend_raises(): + result = _run("import adaptive.learner.triangulation_backend", "bogus") + assert result.returncode != 0 + assert "ValueError" in result.stderr + + +def test_resolve_triangulation_class(): + resolve = backend.resolve_triangulation_class + assert resolve("auto") is backend.Triangulation + assert resolve("python") is python_triangulation.Triangulation + if rust_is_usable(): + import adaptive_triangulation + + assert resolve("rust") is adaptive_triangulation.Triangulation + else: + with pytest.raises(ImportError, match="adaptive-triangulation"): + resolve("rust") + + class MyTriangulation(python_triangulation.Triangulation): + pass + + assert resolve(MyTriangulation) is MyTriangulation + with pytest.raises(ValueError, match="Invalid triangulation backend"): + resolve("bogus") + + +def test_learnernd_triangulation_backend_argument(): + from adaptive import LearnerND + + learner = LearnerND( + lambda xy: sum(xy) ** 2, + bounds=[(-1, 1), (-1, 1)], + triangulation_backend="python", + ) + assert learner._triangulation_class is python_triangulation.Triangulation + assert learner.new()._triangulation_class is python_triangulation.Triangulation + + +@pytest.mark.skipif(not rust_is_usable(), reason="needs adaptive-triangulation") +def test_learnernd_uses_rust_backend(): + import adaptive_triangulation + + from adaptive import LearnerND + + learner = LearnerND(lambda xy: sum(xy) ** 2, bounds=[(-1, 1), (-1, 1)]) + for _ in range(50): + points, _ = learner.ask(1) + for point in points: + learner.tell(point, learner.function(point)) + assert isinstance(learner.tri, adaptive_triangulation.Triangulation) + assert learner.npoints >= 50 diff --git a/docs/source/reference/adaptive.learner.triangulation_backend.md b/docs/source/reference/adaptive.learner.triangulation_backend.md new file mode 100644 index 000000000..7e3407d83 --- /dev/null +++ b/docs/source/reference/adaptive.learner.triangulation_backend.md @@ -0,0 +1,8 @@ +# adaptive.learner.triangulation_backend module + +```{eval-rst} +.. automodule:: adaptive.learner.triangulation_backend + :members: + :undoc-members: + :show-inheritance: +``` diff --git a/docs/source/reference/adaptive.md b/docs/source/reference/adaptive.md index fc0958a94..fb26c5837 100644 --- a/docs/source/reference/adaptive.md +++ b/docs/source/reference/adaptive.md @@ -29,6 +29,8 @@ adaptive.runner.extras ## Other ```{toctree} +adaptive.learner.triangulation +adaptive.learner.triangulation_backend adaptive.utils adaptive.notebook_integration ``` diff --git a/pyproject.toml b/pyproject.toml index e3f4d28f5..fb5366a6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,9 @@ dependencies = [ ] [project.optional-dependencies] +rust = [ + "adaptive-triangulation>=0.2.1", # Rust-accelerated triangulation backend +] other = [ "dill", "distributed",