Add Rust-powered Learner1D with PyO3 bindings

.pre-commit-config.yaml

@@ -20,7 +20,7 @@ repos:
       - id: mypy
         additional_dependencies: [numpy]
         args: [--strict]
-        exclude: '(\.pyi$|tests/|examples/)'
+        exclude: '(\.pyi$|tests/|examples/|benches/)'
   - repo: local
     hooks:
       - id: cargo-fmt

Cargo.toml

@@ -17,11 +17,12 @@ crate-type = ["cdylib", "rlib"]
 criterion = "0.5"
 
 [dependencies]
-pyo3 = { version = "0.23", features = ["extension-module", "abi3-py310"] }
+pyo3 = { version = "0.23", features = ["abi3-py310"] }
 numpy = "0.23"
 nalgebra = "0.33"
 thiserror = "2"
 rustc-hash = "2"
+ordered-float = "4"
 
 [profile.release]
 lto = true

benches/bench_learner1d.py

@@ -0,0 +1,89 @@
+"""Benchmarks for the Rust-powered Learner1D."""
+
+from __future__ import annotations
+
+import math
+import time
+
+from adaptive_triangulation import Learner1D
+
+
+def bench(name: str, fn, *, n_iter: int = 10):
+    """Run a benchmark and print results."""
+    times = []
+    for _ in range(n_iter):
+        t0 = time.perf_counter()
+        fn()
+        times.append(time.perf_counter() - t0)
+    times.sort()
+    median = times[len(times) // 2]
+    print(f"  {name}: {median * 1000:.2f} ms (median of {n_iter})")
+    return median
+
+
+def bench_tell_single():
+    print("=== tell_single ===")
+    for n in [1_000, 10_000, 100_000]:
+        step = 2.0 / n
+
+        def run(n=n, step=step):
+            l = Learner1D(bounds=(-1.0, 1.0))
+            for i in range(n):
+                x = -1.0 + step * i
+                l.tell(x, math.sin(x * 10))
+
+        bench(f"tell {n:,} points", run, n_iter=5 if n >= 100_000 else 10)
+
+
+def bench_tell_many_batch():
+    print("\n=== tell_many (force rebuild) ===")
+    for n in [1_000, 10_000]:
+        step = 2.0 / n
+        xs = [-1.0 + step * i for i in range(n)]
+        ys = [math.sin(x * 10) for x in xs]
+
+        def run(xs=xs, ys=ys):
+            l = Learner1D(bounds=(-1.0, 1.0))
+            l.tell_many(xs, ys, force=True)
+
+        bench(f"tell_many {n:,} points", run)
+
+
+def bench_ask():
+    print("\n=== ask 100 points ===")
+    for n_existing in [100, 1_000, 10_000]:
+        step = 2.0 / n_existing
+        xs = [-1.0 + step * i for i in range(n_existing)]
+        ys = [math.sin(x * 10) for x in xs]
+
+        def run(xs=xs, ys=ys):
+            l = Learner1D(bounds=(-1.0, 1.0))
+            l.tell_many(xs, ys, force=True)
+            l.ask(100, tell_pending=False)
+
+        bench(f"ask 100 (from {n_existing:,} pts)", run)
+
+
+def bench_full_loop():
+    print("\n=== full loop 10K points ===")
+    f = lambda x: math.sin(x * 10)
+
+    def run_serial():
+        l = Learner1D(bounds=(-1.0, 1.0))
+        l.run(f, n_points=10_000, batch_size=1)
+
+    def run_batched():
+        l = Learner1D(bounds=(-1.0, 1.0))
+        l.run(f, n_points=10_000, batch_size=100)
+
+    bench("serial (batch=1)", run_serial, n_iter=3)
+    bench("batched (batch=100)", run_batched, n_iter=3)
+
+
+if __name__ == "__main__":
+    print("Learner1D Benchmarks")
+    print("=" * 50)
+    bench_tell_single()
+    bench_tell_many_batch()
+    bench_ask()
+    bench_full_loop()

build.rs

@@ -0,0 +1,28 @@
+use std::process::Command;
+
+fn main() {
+    let python = std::env::var("PYO3_PYTHON")
+        .or_else(|_| std::env::var("PYTHON_SYS_EXECUTABLE"))
+        .unwrap_or_else(|_| "python3".to_owned());
+
+    let Ok(output) = Command::new(&python)
+        .args([
+            "-c",
+            "import sysconfig; print(sysconfig.get_config_var('LIBDIR') or '')",
+        ])
+        .output()
+    else {
+        return;
+    };
+    if !output.status.success() {
+        return;
+    }
+
+    let Ok(libdir) = String::from_utf8(output.stdout) else {
+        return;
+    };
+    let libdir = libdir.trim();
+    if !libdir.is_empty() {
+        println!("cargo:rustc-link-arg=-Wl,-rpath,{libdir}");
+    }
+}

pyproject.toml

@@ -60,6 +60,7 @@ select = ["ALL"]
 ignore = ["D", "COM812", "ISC001"]
 
 [tool.ruff.lint.per-file-ignores]
+"benches/*" = ["INP001", "T201", "ANN", "E741", "PLR2004", "E731"]
 "tests/*" = ["S101", "PLR2004", "ANN"]
 "examples/*" = ["INP001", "T201", "S101", "N813", "RUF001", "RUF002", "SIM105", "PERF203"]
 

python/adaptive_triangulation/__init__.py

@@ -23,6 +23,7 @@
 from __future__ import annotations
 
 from ._rust import (
+    Learner1D,
     SimplicesProxy,
     Triangulation,
     VertexToSimplicesProxy,
@@ -41,6 +42,7 @@
 )
 
 __all__: list[str] = [
+    "Learner1D",
     "SimplicesProxy",
     "Triangulation",
     "VertexToSimplicesProxy",

python/adaptive_triangulation/_rust.pyi

@@ -14,6 +14,47 @@ TransformLike: TypeAlias = Sequence[Sequence[float]] | npt.ArrayLike
 
 __version__: str
 
+class Learner1D:
+    def __init__(
+        self,
+        bounds: tuple[float, float],
+        loss_per_interval: object | None = None,
+    ) -> None: ...
+    def tell(self, x: float, y: float | Sequence[float]) -> None: ...
+    def tell_many(
+        self,
+        xs: list[float],
+        ys: list[float | Sequence[float]],
+        force: bool = False,
+    ) -> None: ...
+    def tell_pending(self, x: float) -> None: ...
+    def ask(self, n: int, tell_pending: bool = True) -> tuple[list[float], list[float]]: ...
+    def run(
+        self,
+        f: object,
+        *,
+        goal: float | None = None,
+        n_points: int | None = None,
+        batch_size: int = 1,
+    ) -> int: ...
+    @property
+    def loss(self) -> float: ...
+    @property
+    def npoints(self) -> int: ...
+    @property
+    def vdim(self) -> int | None: ...
+    def to_numpy(
+        self,
+    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: ...
+    def remove_unfinished(self) -> None: ...
+    @property
+    def pending_points(self) -> list[float]: ...
+    @property
+    def data(self) -> dict[float, float | list[float]]: ...
+    def intervals(self) -> list[tuple[float, float, float]]: ...
+    @property
+    def bounds(self) -> tuple[float, float]: ...
+
 class SimplicesProxy:
     def __contains__(self, simplex: Simplex) -> bool: ...
     def __iter__(self) -> Iterator[Simplex]: ...