Adopt adaptive-triangulation 0.3.1: batched tell_pending, Rust default_loss, CI job (#499)

adaptive-triangulation 0.3.1 fixes the degenerate-simplex divergence noted
when the backend was introduced (curvature losses crashed on collinear
input) and 0.3.0 added two APIs that move LearnerND's remaining Python hot
loops into Rust. This bumps _MIN_RUST_VERSION to (0, 3, 1) and adopts both:

- tell_pending asks the triangulation for all simplices containing the
  pending point in one simplices_containing call (with the known containing
  simplex as a hint) instead of looping point_in_simplex over every
  neighbouring simplex; the pure-Python backend keeps the original loop.
- LearnerND's default loss_per_simplex is the Rust default_loss when the
  Rust backend is active (same embedded-simplex-volume computation).

Together: 1.10x (2D ring_of_fire, 3000 pts) to 1.42x (3D, 1500 pts) faster
end-to-end on top of the existing Rust backend, sampling identical points
(asserted by a new test).

Also adds the Rust-backend CI job deferred when the backend landed: a nox
session running the suite with ADAPTIVE_TRIANGULATION_BACKEND=rust, now
that the full suite passes with it (294 passed; previously 7 known
curvature-loss failures).

Author: basnijholt

.github/workflows/nox.yml

@@ -26,3 +26,16 @@ jobs:
       run: uv run --group nox nox -e "pytest_min_deps-${{ matrix.python-version }}"
     - name: Test with nox with all dependencies
       run: uv run --group nox nox -e "pytest_all_deps-${{ matrix.python-version }}"
+
+  test-rust-backend:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.13"
+    - name: Install uv
+      uses: astral-sh/setup-uv@v6
+    - name: Test with nox with the Rust triangulation backend
+      run: uv run --group nox nox -e pytest_rust_backend

adaptive/learner/learnerND.py

@@ -20,6 +20,7 @@
     circumsphere,
     point_in_simplex,
     resolve_triangulation_class,
+    rust_default_loss,
     simplex_volume_in_embedding,
 )
 from adaptive.notebook_integration import ensure_holoviews, ensure_plotly
@@ -333,7 +334,9 @@ def __init__(
     ):
         self._triangulation_class = resolve_triangulation_class(triangulation_backend)
         self._vdim = None
-        self.loss_per_simplex = loss_per_simplex or default_loss
+        # Prefer the Rust implementation of the default loss when the Rust
+        # backend is active; it computes the same embedded simplex volume.
+        self.loss_per_simplex = loss_per_simplex or rust_default_loss or default_loss
 
         if hasattr(self.loss_per_simplex, "nth_neighbors"):
             if self.loss_per_simplex.nth_neighbors > 1:
@@ -649,35 +652,46 @@ def tell_pending(self, point, *, simplex=None):
         if self.tri is None:
             return
 
+        for simpl in self._simplices_containing_point(point, simplex):
+            _, to_add = self._add_pending_point_to_simplex(point, simpl)
+            if to_add is None:
+                continue
+            self._update_subsimplex_losses(simpl, to_add)
+
+    def _simplices_containing_point(self, point, simplex=None):
+        """All simplices of the triangulation containing `point`, found from
+        the `simplex` hint when given."""
+        if hasattr(self.tri, "simplices_containing"):
+            # Rust backend: one call instead of a point_in_simplex loop.
+            return self.tri.simplices_containing(point, simplex=simplex)
+
         simplex = tuple(simplex or self.tri.locate_point(point))
         if not simplex:
-            return
-            # Simplex is None if pending point is outside the triangulation,
-            # then you do not have subtriangles
+            return []
+            # Simplex is empty if the pending point is outside the
+            # triangulation, then you do not have subtriangles
 
-        simplex = tuple(simplex)
         simplices = [self.tri.vertex_to_simplices[i] for i in simplex]
         neighbors = set.union(*simplices)
         # Neighbours also includes the simplex itself
+        return [s for s in neighbors if self.tri.point_in_simplex(point, s)]
 
-        for simpl in neighbors:
-            _, to_add = self._try_adding_pending_point_to_simplex(point, simpl)
-            if to_add is None:
-                continue
-            self._update_subsimplex_losses(simpl, to_add)
-
-    def _try_adding_pending_point_to_simplex(self, point, simplex):
-        # try to insert it
-        if not self.tri.point_in_simplex(point, simplex):
-            return None, None
-
+    def _add_pending_point_to_simplex(self, point, simplex):
+        """Insert `point` into the subtriangulation of `simplex`, which must
+        contain the point."""
         if simplex not in self._subtriangulations:
             vertices = self.tri.get_vertices(simplex)
             self._subtriangulations[simplex] = self._triangulation_class(vertices)
 
         self._pending_to_simplex[point] = simplex
         return self._subtriangulations[simplex].add_point(point)
 
+    def _try_adding_pending_point_to_simplex(self, point, simplex):
+        # try to insert it
+        if not self.tri.point_in_simplex(point, simplex):
+            return None, None
+        return self._add_pending_point_to_simplex(point, simplex)
+
     def _update_subsimplex_losses(self, simplex, new_subsimplices):
         loss = self._losses[simplex]
 

adaptive/learner/triangulation_backend.py

@@ -26,9 +26,11 @@
 
 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)
+# Minimal version that is a complete drop-in for the learners: includes the
+# degenerate-simplex fix for curvature losses, plus the batched
+# ``simplices_containing`` query and Rust ``default_loss`` that `LearnerND`
+# uses when this backend is active.
+_MIN_RUST_VERSION = (0, 3, 1)
 
 
 def _rust_version() -> tuple[int, ...] | None:
@@ -119,6 +121,12 @@ def resolve_triangulation_class(backend="auto"):
         point_in_simplex,
         simplex_volume_in_embedding,
     )
+
+    # The Rust implementation of `adaptive.learner.learnerND.default_loss`,
+    # which `LearnerND` prefers when no loss is given. Defined here (rather
+    # than re-exporting the Python one) to avoid a circular import with
+    # `learnerND`; ``None`` means "use the pure-Python default".
+    from adaptive_triangulation import default_loss as rust_default_loss
 else:
     from adaptive.learner.triangulation import (
         Triangulation,
@@ -132,6 +140,8 @@ def resolve_triangulation_class(backend="auto"):
         simplex_volume_in_embedding,
     )
 
+    rust_default_loss = None
+
 __all__ = [
     "TRIANGULATION_BACKEND",
     "Triangulation",
@@ -143,5 +153,6 @@ def resolve_triangulation_class(backend="auto"):
     "fast_norm",
     "orientation",
     "point_in_simplex",
+    "rust_default_loss",
     "simplex_volume_in_embedding",
 ]

adaptive/tests/unit/test_triangulation_backend.py

@@ -122,3 +122,40 @@ def test_learnernd_uses_rust_backend():
             learner.tell(point, learner.function(point))
     assert isinstance(learner.tri, adaptive_triangulation.Triangulation)
     assert learner.npoints >= 50
+
+
+def test_rust_default_loss_matches_backend():
+    if backend.TRIANGULATION_BACKEND == "rust":
+        import adaptive_triangulation
+
+        assert backend.rust_default_loss is adaptive_triangulation.default_loss
+    else:
+        assert backend.rust_default_loss is None
+
+
+def _ring_of_fire(xy):
+    import numpy as np
+
+    x, y = xy
+    a, d = 0.2, 0.5
+    return x + np.exp(-((x**2 + y**2 - d**2) ** 2) / a**4)
+
+
+@pytest.mark.skipif(not rust_is_usable(), reason="needs adaptive-triangulation")
+def test_rust_backend_samples_identical_points():
+    # The batched tell_pending path and the Rust default loss must not change
+    # which points the learner chooses.
+    from adaptive import LearnerND
+
+    learners = {
+        which: LearnerND(
+            _ring_of_fire, bounds=[(-1, 1), (-1, 1)], triangulation_backend=which
+        )
+        for which in ("python", "rust")
+    }
+    for learner in learners.values():
+        for _ in range(200):
+            points, _ = learner.ask(1)
+            for point in points:
+                learner.tell(point, learner.function(point))
+    assert sorted(learners["python"].data) == sorted(learners["rust"].data)

noxfile.py

@@ -27,6 +27,14 @@ def pytest_all_deps(session: nox.Session) -> None:
     session.run("pytest", *xdist)
 
 
+@nox.session(python="3.13")
+def pytest_rust_backend(session: nox.Session) -> None:
+    """Run the test suite with the Rust triangulation backend required."""
+    session.install(".[test,other,rust]")
+    session.run("coverage", "erase")
+    session.run("pytest", *xdist, env={"ADAPTIVE_TRIANGULATION_BACKEND": "rust"})
+
+
 @nox.session(python="3.13")
 def pytest_typeguard(session: nox.Session) -> None:
     """Run pytest with typeguard."""

pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
 
 [project.optional-dependencies]
 rust = [
-    "adaptive-triangulation>=0.2.1", # Rust-accelerated triangulation backend
+    "adaptive-triangulation>=0.3.1", # Rust-accelerated triangulation backend
 ]
 other = [
     "dill",