diff --git a/.gitignore b/.gitignore index f22d2ec..6c3ab9f 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,8 @@ __pycache__ .doctrees/ jupyter_execute/ docs/build/ + +# Suite run outputs (regenerated each run; baseline.json is the committed reference) +benchmarks/results/suite_report.md +benchmarks/results/suite_results.json +benchmarks/results/tuned_frag/ diff --git a/DEV_NOTES.md b/DEV_NOTES.md index 840bcfa..088adbb 100644 --- a/DEV_NOTES.md +++ b/DEV_NOTES.md @@ -26,8 +26,14 @@ uv run --extra dev ./transpile_py310.py uv run --extra docs sphinx-build -b dirhtml docs_source/source docs # Run benchmark on example notebooks. -uv run --extra examples --extra docs python benchmark.py -uv run --extra examples --extra docs python benchmark.py --category robotics # Filter by category +uv run --extra dev --extra docs python benchmark.py +uv run --extra dev --extra docs python benchmark.py --category robotics # Filter by category + +# Benchmark + regression suite (BA / examples / pyroki / float32). See +# benchmarks/README.md. Run from the repo root. +uv run --extra dev --extra docs python -m benchmarks.suite # full run + report vs committed baseline +uv run --extra dev --extra docs python -m benchmarks.suite --quick # fast GPU-only inner loop +uv run --extra dev --extra docs python -m benchmarks.suite --gate # exit 1 on regression (CI) ``` ## Style Guidelines diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000..88adc80 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,100 @@ +# jaxls benchmarks + +Two things live here: + +1. **`suite/`** — the standing benchmark + regression suite. One command, + one report, a committed baseline to catch regressions. Use this for + day-to-day optimizer work and CI. +2. **Study scripts** — the harnesses behind the deep-dive in + [`results.md`](results.md) (Schur elimination, GPU optimization, the + LM/scaler investigation). Keep for reproducing those specific figures; + the suite is what you run normally. + +## The suite + +Run everything through `uv` (matching the Makefile / DEV_NOTES). The `dev` +extra brings `tyro`; `docs` brings `matplotlib` + `scikit-sparse` (cholmod). + +```bash +# Default run: the gate set (bundle_adjustment + float32), diff vs baseline. +uv run --extra dev --extra docs python -m benchmarks.suite + +# Fast inner loop while hill-climbing (GPU only, no slow CPU baselines, ~30s). +uv run --extra dev --extra docs python -m benchmarks.suite --quick --only bundle_adjustment + +# CI / regression gate: exit 1 if any metric regressed past tolerance. +uv run --extra dev --extra docs python -m benchmarks.suite --gate + +# Bless the current numbers as the new baseline (do this deliberately). +uv run --extra dev --extra docs python -m benchmarks.suite --update-baseline + +# Opt-in workloads (kept out of the deterministic gate; run on their own). +uv run --extra dev --extra docs python -m benchmarks.suite --only pyroki_ik # downstream IK regression +uv run --extra dev --extra docs python -m benchmarks.suite --only example_notebooks # ~minutes, 19 notebooks +``` + +GPU rows need a CUDA jaxlib in the environment (`uv pip install "jax[cuda12]"`); +without it the suite falls back to CPU. `pyroki_ik` additionally needs the +pyroki package installed. + +The default/gate set is `bundle_adjustment` + `float32_robustness` — fast and +in-process-reliable. `pyroki_ik` (needs pyroki installed) and +`example_notebooks` (19 jupyter subprocesses) are opt-in via `--only`: both +run fine standalone but are slow and/or unreliable as co-tenant subprocesses +of a CUDA-warmed parent, so they don't gate. + +Run from the repo root. Outputs: `results/suite_results.json` (raw metrics) +and `results/suite_report.md` (human-readable diff vs baseline). + +### How it works + +- **`metrics.py`** — the contract. Every measurement is a `Metric` with a + value, unit, `direction` (lower/higher better), and a regression + tolerance. The report and the gate both read this list, so a metric added + in a workload shows up in both automatically. +- **`workloads.py`** — each workload runs a piece of jaxls and emits + metrics. Thin wrappers over the shared kernels in `../bal.py`, + `../device_sweep.py`, `../trace_examples.py` — one implementation, not + two. Current workloads: `bundle_adjustment` (Schur vs full-system, + CPU+GPU, per-solve budgets), `example_notebooks` (every docs notebook's + headline solve), `pyroki_ik` (downstream IK A/B — catches the kind of + regression the scaler change caused), `float32_robustness` (Ladybug + float32 NaN/accuracy). +- **`baseline.py` / `baseline.json`** — the committed baseline and the diff + logic. Regression = a `lower_better` metric exceeding + `baseline * (1 + rel_tol) + abs_tol` (or the reverse for `higher_better`). + Time tolerances are loose (machine noise); correctness tolerances tight. + +### Adding a workload + +Write a `(SuiteConfig) -> WorkloadResult` in `workloads.py`, append metrics, +register it in `ALL`, then `--update-baseline`. Reuse `bal.run_k_iterations` +(warmup + min-of-repeats, returns accepted cost) for timing so the +methodology stays consistent. + +### Per-iteration timestamps + +`jaxls.record_iteration_times()` is a context manager that records a host +`perf_counter` timestamp per LM step (via a callback inside the solve), so one +solve yields a cost-vs-time trace (no matched-k re-solving). The timestamps +live in a host-side Python list — always float64, off the jitted path and out +of `SolveSummary` — so they resolve millisecond steps even without +`jax_enable_x64`. This is what the example traces and BA plots use; it is also +available to any jaxls user for profiling their own solves: + +```python +with jaxls.record_iteration_times() as times: + sol = problem.solve(init) +# times[i] - times[0] -> elapsed seconds to reach iteration i +``` + +## Study scripts (reproduce results.md) + +| script | what it produces | +|---|---| +| `device_sweep.py` | BA cost/time matrix + per-device and 3-up plots | +| `trace_examples.py` | example cost-vs-time traces (main vs PR) | +| `float32_check.py` | standalone float32 robustness check | + +BAL data auto-downloads to `/tmp`. The A/B scripts select the jaxls under +test via `PYTHONPATH` / a git worktree; see each script's docstring. diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/benchmarks/bal.py b/benchmarks/bal.py new file mode 100644 index 0000000..836083d --- /dev/null +++ b/benchmarks/bal.py @@ -0,0 +1,275 @@ +"""Bundle Adjustment in the Large (BAL) loader and reprojection cost. + +Dataset: https://grail.cs.washington.edu/projects/bal/ + +Camera model: angle-axis rotation + translation + focal length + two radial +distortion parameters; projection onto the negative-z image plane. +""" + +from __future__ import annotations + +import bz2 +import functools +import inspect +import time +import urllib.request +from pathlib import Path +from typing import Literal + +import jax +import jax.numpy as jnp +import numpy as onp + +import jaxls + + +def _analyze( + problem: jaxls.LeastSquaresProblem, + schur_elimination: Literal["auto", "off"], +) -> jaxls.AnalyzedLeastSquaresProblem: + """analyze() forwarding `schur_elimination`. jaxls@main predates + elimination (no such parameter); the shim lets the same benchmark scripts + run against both versions for before/after comparisons.""" + if "schur_elimination" in inspect.signature(problem.analyze).parameters: + return problem.analyze(schur_elimination=schur_elimination) + return problem.analyze() + + +def run_k_iterations( + problem: jaxls.AnalyzedLeastSquaresProblem, + initial_vals: jaxls.VarValues, + k: int, + *, + linear_solver: str, + repeats: int = 3, + lambda_initial: float | None = None, + warmup_budget_s: float | None = None, +) -> tuple[float, float]: + """Run exactly k LM iterations; return (accepted cost, min wall-clock). + + The shared measurement kernel for the matched-iteration benchmarks + (`device_sweep.py` and the benchmark suite). Each rule guards against a + real measurement bug: early termination is disabled so k is exact; one full + warmup solve with `jax.block_until_ready` absorbs compilation and + asynchronously dispatched device work; each timed point is the min of + `repeats`. The reported cost is recomputed from the returned (accepted) + solution — the solve summary's cost history also contains rejected + proposals, which would not be honest to report. + + Run on a specific device by passing `jax.device_put`-placed inputs. + """ + + # `lambda_initial`: bundle-adjustment curvature scales put workable + # damping around 1e1-1e3 (see results.md); the library default 5e-4 is + # tuned for ~unit-curvature problems. When set, it is applied uniformly + # to every method under comparison. + extra = ( + {} + if lambda_initial is None + else {"trust_region": jaxls.TrustRegionConfig(lambda_initial=lambda_initial)} + ) + + def solve() -> tuple[jaxls.VarValues, jaxls.SolveSummary]: + return problem.solve( + initial_vals, + linear_solver=linear_solver, # type: ignore[arg-type] + termination=jaxls.TerminationConfig( + max_iterations=k, early_termination=False + ), + verbose=False, + return_summary=True, + **extra, # type: ignore[arg-type] # only ever {"trust_region": ...} + ) + + # Warmup (absorbs compilation). If it blows past `warmup_budget_s`, this + # config is too slow to bother timing repeats — return the warmup's own + # (cost, time). Guards against the pathological CPU full-CG-on-large-BA + # cases whose single solve runs for minutes. + w_start = time.perf_counter() + warm_sol, _ = jax.block_until_ready(solve()) + warm_t = time.perf_counter() - w_start + if warmup_budget_s is not None and warm_t > warmup_budget_s: + cost = float(onp.asarray(problem._compute_cost_info(warm_sol).cost_total)) + return cost, warm_t + + elapsed = [] + sol = None + for _ in range(repeats): + start = time.perf_counter() + sol, _ = jax.block_until_ready(solve()) + elapsed.append(time.perf_counter() - start) + assert sol is not None + + final_cost = float(onp.asarray(problem._compute_cost_info(sol).cost_total)) + return final_cost, min(elapsed) + + +class CameraVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(9)): + """BAL camera: angle-axis (3), translation (3), focal, k1, k2.""" + + +class PointVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(3)): + """3D landmark.""" + + +def rodrigues(aa: jax.Array, x: jax.Array) -> jax.Array: + """Rotate `x` by the angle-axis vector `aa`.""" + theta = jnp.linalg.norm(aa) + 1e-12 + k = aa / theta + return ( + x * jnp.cos(theta) + + jnp.cross(k, x) * jnp.sin(theta) + + k * jnp.dot(k, x) * (1.0 - jnp.cos(theta)) + ) + + +def project(camera: jax.Array, point: jax.Array) -> jax.Array: + """BAL projection: world point -> 2D pixel coordinates.""" + p = rodrigues(camera[:3], point) + camera[3:6] + proj = -p[:2] / p[2] + f, k1, k2 = camera[6], camera[7], camera[8] + r2 = jnp.sum(proj**2) + return f * (1.0 + k1 * r2 + k2 * r2**2) * proj + + +def reprojection_cost( + vals: jaxls.VarValues, + camera: CameraVar, + point: PointVar, + observation: jax.Array, +) -> jax.Array: + return project(vals[camera], vals[point]) - observation + + +def download_bal(name: str, dest: Path) -> Path: + """Download and decompress a BAL problem file if not already present. + + `name` is the path under the BAL data root, e.g. + "ladybug/problem-49-7776-pre.txt.bz2". + """ + if dest.exists(): + return dest + url = f"https://grail.cs.washington.edu/projects/bal/data/{name}" + print(f"Downloading {url} ...") + compressed, _ = urllib.request.urlretrieve(url) + dest.write_bytes(bz2.decompress(Path(compressed).read_bytes())) + return dest + + +@functools.lru_cache(maxsize=4) +def _parse_bal( + path_str: str, +) -> tuple[onp.ndarray, onp.ndarray, onp.ndarray, onp.ndarray, onp.ndarray]: + """Parse a BAL problem file into (cam_idx, pt_idx, obs, cameras, points). + + Cached: the slow per-observation Python parse runs once even when the + same problem is built with several `analyze()` configurations.""" + with Path(path_str).open() as f: + tokens = f.read().split() + it = iter(tokens) + n_cams, n_pts, n_obs = int(next(it)), int(next(it)), int(next(it)) + + cam_idx = onp.empty(n_obs, dtype=onp.int64) + pt_idx = onp.empty(n_obs, dtype=onp.int64) + obs = onp.empty((n_obs, 2), dtype=onp.float64) + for i in range(n_obs): + cam_idx[i] = int(next(it)) + pt_idx[i] = int(next(it)) + obs[i, 0] = float(next(it)) + obs[i, 1] = float(next(it)) + + cameras = onp.array([float(next(it)) for _ in range(9 * n_cams)]).reshape(n_cams, 9) + points = onp.array([float(next(it)) for _ in range(3 * n_pts)]).reshape(n_pts, 3) + return cam_idx, pt_idx, obs, cameras, points + + +def load_bal( + path: Path, + schur_elimination: Literal["auto", "off"] = "auto", +) -> tuple[jaxls.AnalyzedLeastSquaresProblem, jaxls.VarValues]: + """Parse a BAL problem file and build the jaxls problem + initial values.""" + cam_idx, pt_idx, obs, cameras, points = _parse_bal(str(path)) + n_cams, n_pts = cameras.shape[0], points.shape[0] + + problem = _analyze( + jaxls.LeastSquaresProblem( + costs=[ + jaxls.Cost( + reprojection_cost, + ( + CameraVar(jnp.array(cam_idx)), + PointVar(jnp.array(pt_idx)), + jnp.array(obs), + ), + name="reprojection", + ) + ], + variables=[CameraVar(jnp.arange(n_cams)), PointVar(jnp.arange(n_pts))], + ), + schur_elimination, + ) + initial_vals = jaxls.VarValues.make( + [ + CameraVar(jnp.arange(n_cams)).with_value(jnp.array(cameras)), + PointVar(jnp.arange(n_pts)).with_value(jnp.array(points)), + ] + ) + return problem, initial_vals + + +def make_toy_ba( + seed: int = 0, + n_cams: int = 30, + n_pts: int = 700, + obs_per_point: int = 6, + schur_elimination: Literal["auto", "off"] = "auto", +) -> tuple[jaxls.AnalyzedLeastSquaresProblem, jaxls.VarValues]: + """Synthetic, well-conditioned BA problem. Small enough that a full dense + solve is feasible, for the exactness anchor and the well-conditioned + counter-case.""" + rng = onp.random.default_rng(seed) + cams_gt = onp.concatenate( + [ + rng.normal(0, 0.1, (n_cams, 3)), + rng.normal(0, 1.0, (n_cams, 2)), + rng.normal(10.0, 0.5, (n_cams, 1)), + onp.full((n_cams, 1), 500.0), + onp.zeros((n_cams, 2)), + ], + axis=1, + ) + pts_gt = rng.normal(0, 2.0, (n_pts, 3)) + + cam_idx = onp.concatenate( + [rng.choice(n_cams, size=obs_per_point, replace=False) for _ in range(n_pts)] + ) + pt_idx = onp.repeat(onp.arange(n_pts), obs_per_point) + + project_batch = jax.vmap(project) + obs_clean = project_batch(jnp.array(cams_gt[cam_idx]), jnp.array(pts_gt[pt_idx])) + obs = jnp.array(onp.asarray(obs_clean) + rng.normal(0, 1.0, (len(cam_idx), 2))) + + problem = _analyze( + jaxls.LeastSquaresProblem( + costs=[ + jaxls.Cost( + reprojection_cost, + (CameraVar(jnp.array(cam_idx)), PointVar(jnp.array(pt_idx)), obs), + name="reprojection", + ) + ], + variables=[CameraVar(jnp.arange(n_cams)), PointVar(jnp.arange(n_pts))], + ), + schur_elimination, + ) + initial_vals = jaxls.VarValues.make( + [ + CameraVar(jnp.arange(n_cams)).with_value( + jnp.array(cams_gt + rng.normal(0, 0.01, cams_gt.shape)) + ), + PointVar(jnp.arange(n_pts)).with_value( + jnp.array(pts_gt + rng.normal(0, 0.05, pts_gt.shape)) + ), + ] + ) + return problem, initial_vals diff --git a/benchmarks/device_sweep.py b/benchmarks/device_sweep.py new file mode 100644 index 0000000..7cdd495 --- /dev/null +++ b/benchmarks/device_sweep.py @@ -0,0 +1,664 @@ +"""Cross-device benchmark: linear-solver methods on CPU vs GPU. + +Runs a matched-iteration study on both CPU and GPU, across a sweep of problem +sizes, comparing every available linear solver: + + - full CG : conjugate gradient on the full system (no elimination) + - full dense : dense Cholesky on the full system (small problems only) + - cholmod : sparse Cholesky on the full system (CPU only) + - Schur + dense : variable elimination, dense reduced solve + - Schur + CG : variable elimination, matrix-free reduced CG + - Schur + cholmod : variable elimination, sparse-direct reduced solve (CPU only) + +Methodology: run exactly k Levenberg-Marquardt iterations with early +termination off, for a sweep of k, and record +(accepted cost, wall-clock) at each k. One warmup solve per configuration +with `jax.block_until_ready` absorbs compilation and async dispatch; each +timed point is the min of `repeats`. + +Device selection is by `jax.devices(platform)`; the array inputs are placed +on the target device with `jax.device_put` and the solve is run there. The +CPU device is always available; the GPU rows are skipped if no GPU backend +is present. + +Usage: + uv run --extra dev --extra docs python benchmarks/device_sweep.py # run everything + uv run --extra dev --extra docs python benchmarks/device_sweep.py --replot # plots from JSON only + uv run --extra dev --extra docs python benchmarks/device_sweep.py --problems toy ladybug49 + uv run --extra dev --extra docs python benchmarks/device_sweep.py --devices cpu # CPU only +""" + +from __future__ import annotations + +import argparse +import gc +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Literal + +import jax + +jax.config.update("jax_enable_x64", True) + + +import jaxls # noqa: E402 + +from bal import download_bal, load_bal, make_toy_ba, run_k_iterations # noqa: E402 + +RESULTS_DIR = Path(__file__).parent / "results" + + +# --------------------------------------------------------------------------- +# Problem registry +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ProblemSpec: + name: str + title: str + ks: tuple[int, ...] + full_dense_ok: bool + """Whether a dense full-system solve is feasible (small problems only).""" + load: Callable[ + [Literal["auto", "off"]], + tuple[jaxls.AnalyzedLeastSquaresProblem, jaxls.VarValues], + ] + """load(schur_elimination) -> (problem, initial_vals).""" + + +def _bal_loader( + name: str, dest: str +) -> Callable[ + [Literal["auto", "off"]], + tuple[jaxls.AnalyzedLeastSquaresProblem, jaxls.VarValues], +]: + def load( + schur_elimination: Literal["auto", "off"], + ) -> tuple[jaxls.AnalyzedLeastSquaresProblem, jaxls.VarValues]: + path = download_bal(name, Path(dest)) + return load_bal(path, schur_elimination=schur_elimination) + + return load + + +PROBLEMS: dict[str, ProblemSpec] = { + "toy": ProblemSpec( + name="toy", + title="Toy BA (30 cams, 700 points)", + ks=(1, 2, 3, 4, 6, 8, 12, 16), + full_dense_ok=True, + load=lambda se: make_toy_ba(schur_elimination=se), + ), + "ladybug49": ProblemSpec( + name="ladybug49", + title="Ladybug-49 (49 cams, 7,776 pts, 31,843 obs)", + ks=(1, 2, 4, 6, 9, 13, 18, 24, 30), + full_dense_ok=False, + load=_bal_loader("ladybug/problem-49-7776-pre.txt.bz2", "/tmp/ladybug-49.txt"), + ), + # Larger: 138 cameras -> 1,242-dim reduced system. Stresses the dense + # O(n_keep^3) factorization and the GPU's strength at big matmuls. + "trafalgar138": ProblemSpec( + name="trafalgar138", + title="Trafalgar-138 (138 cams, 44,033 pts)", + ks=(1, 2, 4, 6, 9, 13, 18, 24, 30), + full_dense_ok=False, + load=_bal_loader( + "trafalgar/problem-138-44033-pre.txt.bz2", "/tmp/trafalgar-138.txt" + ), + ), +} + + +# Which methods to run, and how. `elimination=False` selects the +# full-system problem; `devices` restricts a method to a subset. +@dataclass(frozen=True) +class Method: + name: str + linear_solver: str + elimination: bool + devices: tuple[str, ...] | None = None # None -> all devices + full_dense_only: bool = False # only on problems where full dense is feasible + + +METHODS: tuple[Method, ...] = ( + Method("full CG", "conjugate_gradient", elimination=False), + Method("cholmod", "cholmod", elimination=False, devices=("cpu",)), + Method("Schur + dense", "dense_cholesky", elimination=True), + Method("Schur + CG", "conjugate_gradient", elimination=True), + # Variable elimination with a sparse-direct CHOLMOD factorization of the + # reduced system (Ceres/g2o-style). CPU only: CHOLMOD runs as a host + # callback. + Method("Schur + cholmod", "cholmod", elimination=True, devices=("cpu",)), + Method( + "full dense", + "dense_cholesky", + elimination=False, + full_dense_only=True, + ), +) + + +# --------------------------------------------------------------------------- +# Timing — the measurement kernel itself is bal.run_k_iterations, shared with +# the benchmark suite so they cannot drift apart methodologically. +# --------------------------------------------------------------------------- + + +def device_for(platform: str) -> "jax.Device | None": # type: ignore[name-defined] + try: + return jax.devices(platform)[0] + except RuntimeError: + return None + + +def run_study(spec: ProblemSpec, devices: list[str], repeats: int) -> dict: + """Run the full method x device x k grid for one problem.""" + print(f"\n## {spec.title}") + problem_elim, init = spec.load("auto") + problem_full, _ = spec.load("off") + + results: dict = {"title": spec.title, "ks": list(spec.ks), "runs": {}} + for platform in devices: + dev = device_for(platform) + if dev is None: + print(f" [{platform}] no device, skipping") + continue + print(f" --- device: {platform} ({dev}) ---") + # Place each problem variant on the device once; every (method, k) + # measurement reuses the placed arrays. + elim_d = jax.device_put(problem_elim, dev) + full_d = jax.device_put(problem_full, dev) + init_d = jax.device_put(init, dev) + for method in METHODS: + if method.devices is not None and platform not in method.devices: + continue + if method.full_dense_only and not spec.full_dense_ok: + continue + target = elim_d if method.elimination else full_d + key = f"{platform}:{method.name}" + costs, times = [], [] + try: + for k in spec.ks: + cost, t = run_k_iterations( + target, + init_d, + k, + linear_solver=method.linear_solver, + repeats=repeats, + ) + costs.append(cost) + times.append(t) + print(f" {method.name:<14} k={k:>2}: {t:>8.3f}s cost={cost:g}") + except Exception as e: # noqa: BLE001 + print(f" {method.name:<14} FAILED: {type(e).__name__}: {e}") + continue + results["runs"][key] = { + "method": method.name, + "device": platform, + "ks": list(spec.ks), + "costs": costs, + "times": times, + } + gc.collect() + + out = RESULTS_DIR / f"device_{spec.name}.json" + out.write_text(json.dumps(results, indent=2)) + print(f" wrote {out}") + return results + + +# Damping used for the cholmod-vs-Schur+cholmod comparison. Bundle-adjustment +# curvature scales want lambda ~1e1-1e3; 1e2 matches the benchmark suite +# (suite/workloads.py ba_lambda_initial) and is applied identically to both +# methods so the comparison is apples-to-apples. +_CHOLMOD_BA_LAMBDA = 1e2 + + +def run_cholmod_study(spec: ProblemSpec, repeats: int) -> dict: + """CPU-only matched-iteration study of full-system CHOLMOD vs Schur + + CHOLMOD (sparse-direct on the reduced system), both with the same tuned + damping. Writes device__cholmod.json.""" + print(f"\n## {spec.title} — CHOLMOD comparison") + problem_elim, init = spec.load("auto") + problem_full, _ = spec.load("off") + dev = device_for("cpu") + assert dev is not None # CPU is always available. + elim_d = jax.device_put(problem_elim, dev) + full_d = jax.device_put(problem_full, dev) + init_d = jax.device_put(init, dev) + + results: dict = {"title": spec.title, "ks": list(spec.ks), "runs": {}} + for label, target, elim in ( + ("cholmod", full_d, False), + ("Schur + cholmod", elim_d, True), + ): + del elim # `target` already selects the right problem variant. + costs, times = [], [] + for k in spec.ks: + cost, t = run_k_iterations( + target, + init_d, + k, + linear_solver="cholmod", + repeats=repeats, + lambda_initial=_CHOLMOD_BA_LAMBDA, + ) + costs.append(cost) + times.append(t) + print(f" {label:<16} k={k:>2}: {t:>8.3f}s cost={cost:g}") + results["runs"][f"cpu:{label}"] = { + "method": label, + "device": "cpu", + "ks": list(spec.ks), + "costs": costs, + "times": times, + } + gc.collect() + + out = RESULTS_DIR / f"device_{spec.name}_cholmod.json" + out.write_text(json.dumps(results, indent=2)) + print(f" wrote {out}") + return results + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- + +# Method styling. Full-system methods get warm/neutral colors and a thin line; +# Schur (variable-elimination) methods get a cool family and a thicker line, so +# elimination-vs-full reads at a glance without spending the linestyle (which +# encodes tuned-PR "-" vs main-baseline "--"). +_COLOR = { + "full CG": "#d62728", # red + "full dense": "#7f7f7f", # gray + "cholmod": "#e377c2", # pink + "Schur + dense": "#1f77b4", # blue + "Schur + CG": "#2ca02c", # green + "Schur + cholmod": "#17becf", # teal +} +_SCHUR_LINEWIDTH = 2.8 +_FULL_LINEWIDTH = 1.4 + +# Draw order: full-system methods first, then Schur, so the legend groups them. +_METHOD_ORDER = ( + "full CG", + "full dense", + "cholmod", + "Schur + dense", + "Schur + CG", + "Schur + cholmod", +) + + +def _is_schur(method: str) -> bool: + return method.startswith("Schur") + + +def _load_runs(name: str) -> "tuple[dict, dict, str] | None": + """Gather plot data for one problem. Returns (final_runs, main_runs, title), + each runs dict keyed by "device:method". + + Sources, in priority order: the tuned PR JSON, then the plain sweep JSON + (so methods added after the tuned pair was generated still appear; a + device:method already present from the tuned data is not overwritten). The + tuned main-baseline JSON is returned separately for the dashed comparison.""" + final: dict = {} + main: dict = {} + title = None + tuned_final = RESULTS_DIR / f"device_{name}_tuned_final.json" + if tuned_final.exists(): + d = json.loads(tuned_final.read_text()) + title = d["title"] + final.update(d["runs"]) + tuned_main = RESULTS_DIR / f"device_{name}_tuned_main.json" + if tuned_main.exists(): + main.update(json.loads(tuned_main.read_text())["runs"]) + plain = RESULTS_DIR / f"device_{name}.json" + if plain.exists(): + d = json.loads(plain.read_text()) + title = title or d["title"] + for k, r in d["runs"].items(): + final.setdefault(k, r) # don't overwrite a tuned curve + if not final: + return None + return final, main, title or name + + +def _time_to(run: dict, target: float) -> float | None: + """Wall-clock at which a run first reaches `target` cost (or below).""" + for t, c in zip(run["times"], run["costs"]): + if c <= target: + return t + return None + + +def plot_study(name: str) -> None: + """One figure per device, two stacked panels telling the elimination story: + + - top: suboptimality gap (cost - best-observed optimum) vs wall-clock, + log-log. A method's curve plunging to the floor marks when it reaches the + solution; its horizontal position is the time-to-solution. Schur + (variable-elimination) methods are cool-colored and thick, full-system + methods warm/neutral and thin, so elimination-vs-full reads at a glance. + - bottom: speedup-to-solution of each method over the fastest full-system + method (bars). This is the headline number. + + The tuned main-system baseline (if present) is drawn dashed in the top + panel for context; it is not included in the speedup bars.""" + loaded = _load_runs(name) + if loaded is None: + return + final, main, title = loaded + + for device in ("cpu", "gpu"): + runs = {r["method"]: r for r in final.values() if r["device"] == device} + if not runs: + continue + main_d = {r["method"]: r for r in main.values() if r["device"] == device} + out = RESULTS_DIR / f"device_{name}_{device}.png" + _render_study_figure(runs, main_d, f"{title} — {device.upper()}", out) + + +def plot_cholmod_comparison(name: str) -> None: + """Focused CPU comparison: full-system CHOLMOD vs Schur + CHOLMOD, from + device__cholmod.json. Same two-panel layout as plot_study; the + speedup bar is Schur + CHOLMOD over full-system CHOLMOD.""" + path = RESULTS_DIR / f"device_{name}_cholmod.json" + if not path.exists(): + return + data = json.loads(path.read_text()) + runs = {r["method"]: r for r in data["runs"].values() if r["device"] == "cpu"} + if not runs: + return + out = RESULTS_DIR / f"device_{name}_cholmod.png" + _render_study_figure(runs, {}, f"{data['title']} — CHOLMOD (CPU)", out) + + +def _render_study_figure(runs: dict, main: dict, title: str, out: "Path") -> None: + """Render one two-panel study figure (top: suboptimality gap vs wall-clock; + bottom: speedup-to-solution bars over the fastest full-system method) for a + set of method runs (keyed by method name). `main` holds optional dashed + main-system baselines.""" + import matplotlib.pyplot as plt + import numpy as np + + methods = [m for m in _METHOD_ORDER if m in runs] + if not methods: + return + copt = min(min(r["costs"]) for r in runs.values()) + # "Everyone reaches this" threshold, for an apples-to-apples speedup. + target = max(r["costs"][-1] for r in runs.values()) * 1.03 + + fig = plt.figure(figsize=(7.0, 7.2), constrained_layout=True) + gs = fig.add_gridspec(2, 1, height_ratios=[1.5, 1.0]) + ax = fig.add_subplot(gs[0]) + ax2 = fig.add_subplot(gs[1]) + + # --- top: suboptimality gap vs wall-clock --- + for m in methods: + color, is_schur = _COLOR[m], _is_schur(m) + r = runs[m] + gap = [max(c - copt, 1.0) for c in r["costs"]] + ax.plot( + r["times"], + gap, + marker="o", + ms=3.5, + lw=_SCHUR_LINEWIDTH if is_schur else _FULL_LINEWIDTH, + color=color, + label=m, + zorder=3 if is_schur else 2, + solid_capstyle="round", + ) + # Main-system baselines, dashed and thin, for context. + for m, r in main.items(): + if m not in _COLOR: + continue + gap = [max(c - copt, 1.0) for c in r["costs"]] + ax.plot( + r["times"], + gap, + marker="o", + ms=2.5, + lw=1.3, + ls="--", + color=_COLOR[m], + label=f"main: {m}", + zorder=1, + ) + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_xlabel("wall-clock (s)") + ax.set_ylabel("cost above optimum") + ax.grid(True, which="major", alpha=0.25) + ax.grid(True, which="minor", alpha=0.07) + ax.legend(fontsize=8, framealpha=0.92, loc="lower left") + + # --- bottom: speedup-to-solution over fastest full-system method --- + tt = {m: _time_to(runs[m], target) for m in methods} + full_times = [t for m in methods if (t := tt[m]) is not None and not _is_schur(m)] + bar_methods = [m for m in methods if tt[m] is not None] + if full_times and bar_methods: + base = min(full_times) + speed = [base / t for m in bar_methods if (t := tt[m]) is not None] + xs = np.arange(len(bar_methods)) + ax2.bar( + xs, + speed, + width=0.72, + color=[_COLOR[m] for m in bar_methods], + edgecolor=["#222" if _is_schur(m) else "none" for m in bar_methods], + linewidth=1.3, + zorder=3, + ) + ax2.axhline(1.0, color="k", lw=0.9, ls="--", alpha=0.55, zorder=1) + ax2.set_yscale("log") + ax2.set_ylim(top=max(speed) * 2.3) + ax2.set_xticks(xs) + ax2.set_xticklabels( + [ + m.replace("Schur + ", "Schur\n").replace("full ", "full\n") + for m in bar_methods + ], + fontsize=8, + ) + ax2.set_ylabel("speedup to solution\n(vs fastest full-system)") + ax2.grid(True, axis="y", alpha=0.22, zorder=0) + for x, s in zip(xs, speed): + ax2.text( + float(x), + s * 1.08, + (f"{s:.0f}×" if s >= 10 else f"{s:.1f}×" if s >= 1 else f"{s:.2g}×"), + ha="center", + va="bottom", + fontsize=8.5, + fontweight="bold", + ) + + fig.suptitle(title, fontsize=12, fontweight="bold") + fig.savefig(out, dpi=150) + print(f"wrote {out}") + plt.close(fig) + + +def plot_ba_comparison() -> None: + """One figure, three subplots (toy / Ladybug-49 / Trafalgar-138), GPU. + + Each subplot is cost vs wall-clock with a marker per LM step, comparing + main's full-system CG baseline against the PR's Schur+dense and + Schur+CG. Reads the tuned PR/main JSON pair per problem.""" + import matplotlib.pyplot as plt + + problems = [ + ("toy", "Toy (30 cams)"), + ("ladybug49", "Ladybug-49"), + ("trafalgar138", "Trafalgar-138"), + ] + # (label, source-file suffix, method, color, marker) + series = [ + ("full CG (main)", "main", "full CG", "#d62728", "o"), + ("Schur+dense (PR)", "final", "Schur + dense", "#1f77b4", "^"), + ("Schur+CG (PR)", "final", "Schur + CG", "#2ca02c", "s"), + ] + + def _running_best(run: dict, c0: float) -> tuple[list[float], list[float]]: + # "Best solution found by time t." Matched-k runs independent solves + # per k, so an unconverged inexact-CG baseline can be non-monotone in + # k; cummin gives the honest, comparable curve. Step 0 (initial cost + # at t=0) is the common anchor every method descends from. + times = [0.0] + list(run["times"]) + best = [c0] + list(run["costs"]) + for i in range(1, len(best)): + best[i] = min(best[i], best[i - 1]) + return times, best + + def _time_to_cost(run: dict, c0: float, target: float) -> float | None: + """First wall-clock time the running-best cost reaches `target`.""" + for t, b in zip(*_running_best(run, c0)): + if b <= target: + return t + return None + + initial = json.loads((RESULTS_DIR / "ba_initial_cost.json").read_text()) + fig, axes = plt.subplots(1, 3, figsize=(16, 5.0)) + for ax, (prob, title) in zip(axes, problems): + c0 = initial.get(prob) + cache: dict[str, dict] = {} + + def _load(suffix: str) -> dict | None: + if suffix not in cache: + path = RESULTS_DIR / f"device_{prob}_tuned_{suffix}.json" + cache[suffix] = json.loads(path.read_text()) if path.exists() else {} + return cache[suffix] or None + + # Collect the runs present for this problem. + runs = [] # (label, run, color, marker) + for label, suffix, method, color, marker in series: + data = _load(suffix) + run = data["runs"].get(f"gpu:{method}") if data else None + if run is not None: + runs.append((label, run, color, marker)) + + # Threshold: the *worst* converged cost across the three methods — + # i.e. the best cost every method actually reaches. Labeling each + # curve's crossing time there gives an apples-to-apples "how long to + # the same solution?" that every method can answer (a stricter + # threshold would leave the slowest-converging method unlabeled). A + # small tolerance absorbs matched-k jitter so a method that lands + # ~0.1% above its floor still counts as having reached it. + threshold = max(min(run["costs"]) for _, run, _, _ in runs) if runs else None + target = threshold * 1.003 if threshold is not None else None + + crossings = [] # (label, time, color) to annotate after curves drawn + for label, run, color, marker in runs: + times, best = _running_best(run, c0) + ax.plot(times, best, marker=marker, ms=5, color=color, label=label) + t_hit = _time_to_cost(run, c0, target) if target is not None else None + if t_hit is not None and t_hit > 0: + crossings.append((label, t_hit, color)) + if run.get("budget_stopped_at_k") is not None: + ax.plot(times[-1], best[-1], marker="x", ms=11, mew=2, color=color) + + if threshold is not None: + ax.axhline(threshold, color="0.4", lw=1.4, ls="--", zorder=0) + ax.annotate( + "shared cost target", + (0, threshold), + textcoords="offset points", + xytext=(4, 4), + ha="left", + va="bottom", + fontsize=8, + color="0.4", + ) + # Mark where each method crosses the threshold and label the absolute + # time at that point — staggered vertically so the labels don't + # collide, with a leader dot on the threshold line. + crossings.sort(key=lambda c: c[1]) + for i, (label, t_hit, color) in enumerate(crossings): + ax.plot([t_hit], [threshold], marker="o", ms=7, color=color, zorder=5) + txt = f"{t_hit * 1e3:.0f} ms" if t_hit < 1.0 else f"{t_hit:.1f} s" + ax.annotate( + txt, + (t_hit, threshold), + textcoords="offset points", + xytext=(0, 14 + 16 * i), + ha="center", + va="bottom", + fontsize=11, + fontweight="bold", + color=color, + arrowprops=dict(arrowstyle="-", color=color, lw=0.8), + ) + ax.set_yscale("log") + # symlog: linear through t=0 (so step 0 shows), log beyond — the + # fast Schur methods (~0.1 s) and the slow baseline (~10-30 s) + # otherwise can't share a readable x-axis. + ax.set_xscale("symlog", linthresh=1e-2) + ax.set_xlim(left=0) + ax.set_xlabel("wall-clock (s, symlog)", fontsize=11) + ax.set_ylabel("accepted cost", fontsize=11) + ax.set_title(title, fontsize=13) + ax.tick_params(labelsize=10) + ax.grid(True, which="major", ls="-", lw=0.4, alpha=0.3) + ax.legend(fontsize=9, loc="upper right") + fig.suptitle( + "Bundle adjustment on GPU: cost vs wall-clock — Schur elimination " + "vs full-system CG\n(marker per LM step from step 0; dashed line = " + "shared cost target reached by all three methods, labels = time each " + "reaches it; × = budget cut-off)", + fontsize=13, + ) + fig.tight_layout() + out = RESULTS_DIR / "ba_comparison_gpu.png" + fig.savefig(out, dpi=150) + print(f"wrote {out}") + plt.close(fig) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--replot", action="store_true") + parser.add_argument( + "--problems", nargs="+", default=list(PROBLEMS), choices=list(PROBLEMS) + ) + parser.add_argument( + "--devices", nargs="+", default=["cpu", "gpu"], choices=["cpu", "gpu"] + ) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument( + "--cholmod", + action="store_true", + help="CPU CHOLMOD vs Schur+CHOLMOD comparison instead of the full sweep.", + ) + args = parser.parse_args() + RESULTS_DIR.mkdir(exist_ok=True) + + print("JAX devices:", jax.devices()) + for plat in ("cpu", "gpu"): + d = device_for(plat) + print(f" {plat}: {d}") + + for prob_name in args.problems: + spec = PROBLEMS[prob_name] + if args.cholmod: + if not args.replot: + run_cholmod_study(spec, args.repeats) + plot_cholmod_comparison(spec.name) + else: + if not args.replot: + run_study(spec, args.devices, args.repeats) + plot_study(spec.name) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/float32_check.py b/benchmarks/float32_check.py new file mode 100644 index 0000000..11b52c1 --- /dev/null +++ b/benchmarks/float32_check.py @@ -0,0 +1,55 @@ +"""float32 robustness check (Ladybug-49, direct Schur path). + +Forming S = H_cc - W V^{-1} W^T cancels catastrophically in float32; a naive +implementation NaNs in the Cholesky factorization. Requirement: a full +float32 solve with zero NaNs, converging to a cost comparable to float64. + +Usage: uv run --extra dev --extra docs python benchmarks/float32_check.py +""" + +from __future__ import annotations + +from pathlib import Path + +import jax +import numpy as onp + + +def run(x64: bool) -> onp.ndarray: + jax.config.update("jax_enable_x64", x64) + # Imports happen after the precision flag so arrays pick up the dtype. + import jaxls + + from bal import download_bal, load_bal + + path = download_bal( + "ladybug/problem-49-7776-pre.txt.bz2", Path("/tmp/ladybug-49.txt") + ) + problem, initial_vals = load_bal(path) + _, summary = problem.solve( + initial_vals, + linear_solver="dense_cholesky", + termination=jaxls.TerminationConfig(max_iterations=30, early_termination=False), + verbose=False, + return_summary=True, + ) + return onp.asarray(summary.cost_history) + + +def main() -> None: + history_f32 = run(x64=False) + nan_count = int(onp.isnan(history_f32).sum()) + best_f32 = float(history_f32[history_f32 > 0].min()) + print(f"float32: NaNs in cost history: {nan_count}") + print(f"float32: best cost: {best_f32:.1f}") + assert nan_count == 0, "float32 robustness FAILED: NaNs in cost history" + + history_f64 = run(x64=True) + best_f64 = float(history_f64[history_f64 > 0].min()) + print(f"float64: best cost: {best_f64:.1f}") + print(f"relative gap (f32 vs f64): {abs(best_f32 - best_f64) / best_f64:.2e}") + print("PASS") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/results.md b/benchmarks/results.md new file mode 100644 index 0000000..5b8c719 --- /dev/null +++ b/benchmarks/results.md @@ -0,0 +1,490 @@ +# Variable elimination (Schur complement) — results + +Branch `schur-elimination-v2`: a from-scratch reimplementation of variable +elimination for bundle adjustment, measured with the matched-iteration +methodology below. Elimination is automatic by default +(`analyze(schur_elimination="auto")`): `solve()` eliminates dominant +block-diagonal variable types (chosen from the problem structure, logged) and +all three linear solvers use the reduced system — dense Cholesky and CG solve +it densely / matrix-free, and cholmod factors it sparse-directly. +`schur_elimination="off"` opts out; the "full CG"/"full dense" baselines below +use it. Hardware: Apple Silicon CPU, float64 +unless noted. Absolute times are hardware-dependent; the exactness results +and the matched-iteration *orderings* are the portable claims. + +**Bottom line.** Variable elimination is correct (exact Newton steps), +robust (float32 with zero NaNs), and **beats the full-CG baseline at every +matched iteration count on real bundle adjustment, on both cost and +wall-clock**: at k=30 on Ladybug-49, Schur+dense is **13× faster to a lower +cost** and Schur+CG is **3–4× faster to the same cost**. On the +well-conditioned toy problem — where there is nothing to win — both Schur +paths still run slightly faster than full CG. Two Levenberg-Marquardt fixes +discovered during this work (see below) also substantially improved *all* +solvers, including the full-CG baseline itself. + +## Methodology: matched outer iterations + +"Time to convergence" is confounded by termination: a solver that stops +earlier looks faster for reasons unrelated to the linear algebra. Instead we +run **exactly k LM iterations** (early termination off), for a sweep of k, +and compare (accepted cost, wall-clock) at matched k. Timing rules: one full +warmup solve per configuration with `jax.block_until_ready` (absorbs +compilation and asynchronous dispatch), minimum of 3 repeats, and the +full-CG baseline doubles as a control line for the noise floor. Reported +costs are those of the **returned (accepted) solution** — the solve summary's +cost history also contains rejected proposals, which would not be honest to +report. + +Reproduce: `uv run --extra dev --extra docs python benchmarks/device_sweep.py` (raw data saved to JSON; +`--replot` regenerates plots without re-running). + +## Correctness — exact reduced steps (pass) + +The direct Schur path solves the same damped normal equations as a full +dense Cholesky, exactly: + +- Single damped steps match an explicit full dense solve to **~1e-10 + relative** across a sweep of damping values (`test_schur_single_step_exactness`), + including problems with multiple kept variable types. +- LM cost trajectories match full dense to ~5e-9 relative on a small BA + problem (`test_schur_dense_exactness`); visible in the toy table below, + where Schur+dense and full dense produce identical costs at every k. +- On Ladybug-49, the reduced step satisfies the full-system normal equations + to `|Hd-b|/|b| ~ 1e-11`, verified against a host-side sparse solve. + +## Robustness — float32 (pass) + +Forming `S = H_cc − W V⁻¹ Wᵀ` is a difference of nearly-equal SPD matrices +and cancels catastrophically in float32. The reduced solve applies Jacobi +(diagonal) scaling plus a precision-adaptive Tikhonov floor (~2e-3 relative +in float32, ~1e-15 in float64 so float64 stays exact). Ladybug-49 end-to-end +in float32 (`benchmarks/float32_check.py`): + +- **0 NaNs**; best cost 26,785 (float32) vs 26,781 (float64) — a 1.5e-4 + relative gap. + +## Performance — Ladybug-49 (real BAL: 49 cameras, 7,776 points, 31,843 observations) + +> **Note:** the table below was timed before the column-scaler splice +> (see "The column-scaling math"). Under the final scaler, matched-k +> costs improve slightly (k=30: full CG 26,768, Schur+dense 26,746, +> Schur+CG 26,767) and the orderings are unchanged; post-splice timing +> runs coincided with other load on the machine (the unchanged +> full-dense control line ran 4× slow) and were discarded rather than +> published. Re-time with `uv run --extra dev --extra docs python benchmarks/device_sweep.py` on a +> quiet machine. + +Accepted cost / wall-clock at matched k (float64, min of 3, CPU): + +| k | full CG (baseline) | Schur + dense | Schur + CG | +|---:|--------------------|---------------------|---------------------| +| 1 | 41,432 / 0.02 s | 42,254 / 0.03 s | 40,919 / 0.04 s | +| 2 | 29,466 / 0.10 s | 28,718 / 0.06 s | 28,824 / 0.09 s | +| 4 | 27,206 / 0.30 s | 27,175 / 0.10 s | 27,212 / 0.20 s | +| 6 | 26,964 / 0.64 s | 26,938 / 0.13 s | 26,961 / 0.33 s | +| 9 | 26,890 / 1.30 s | 26,861 / 0.19 s | 26,888 / 0.56 s | +| 13 | 26,862 / 2.38 s | 26,829 / 0.28 s | 26,860 / 0.91 s | +| 18 | 26,840 / 3.92 s | 26,807 / 0.37 s | 26,839 / 1.37 s | +| 24 | 26,822 / 5.95 s | 26,791 / 0.50 s | 26,821 / 1.93 s | +| 30 | 26,810 / 8.39 s | **26,779 / 0.62 s** | **26,809 / 2.57 s** | + +- **Schur+dense dominates at every k from k=2 on: lower cost *and* less + time.** At k=30 it is **13× faster** than full CG and reaches a lower cost + than full CG ever does; it already beats full CG's best-ever cost + (26,810 @ 8.39 s) by k=18 (26,807 @ 0.37 s) — **~22× less wall-clock to + full CG's best cost.** +- **Schur+CG tracks full CG's cost trajectory almost exactly at ~3× less + wall-clock** (2.57 s vs 8.39 s at k=30). Per-iteration, both CG variants + pay more as the Eisenstat-Walker forcing term tightens, but the reduced + 441-dim system needs far fewer inner iterations than the full 23,769-dim + one. +- Why: each outer iteration, full CG re-solves the ill-conditioned full + system from scratch with a runaway inner-iteration count near convergence. + The Schur path eliminates the points analytically (block-diagonal V), and + the per-iteration cost of the direct reduced solve is flat: ~2.5 ms of + damping-independent preparation plus ~17 ms per damping value, dominated by + the pair-block products for `W V⁻¹ Wᵀ`. + +## Performance — Toy (30 cameras, 700 points): the well-conditioned counter-case + +All solvers reach the optimum (5,913.72) by k≈4; there is nothing for +conditioning to win. At k=16: Schur+dense **0.04 s**, Schur+CG **0.09 s**, +full CG 0.13 s, full dense 6.4 s. So elimination costs nothing on easy +problems (it is in fact mildly faster here too), and full dense is ~150× +slower — the usual `O(n³)` cliff that makes it infeasible at real scale. + +## Two Levenberg-Marquardt fixes that this work surfaced + +> **Superseded in part — see "Final configuration" at the end.** Fix 1 (the +> spliced scaler) was later found to regress downstream workloads whose +> lambdas were tuned against the legacy convention (pyroki IK, M3500 pose +> graphs) and was **reverted**; the rejected-step termination fixes plus +> Nielsen lambda escalation recover the BAL behavior the splice was built +> for. Fix 2 stands. + +Both were found by asking why *exact* Newton steps were being rejected on +raw BAL data while inexact CG steps were accepted. They live in the shared +LM loop and improve every solver: + +1. **Jacobian column scaling: spliced Levenberg/Marquardt damping.** See + the derivation below. *(Later reverted; see above.)* +2. **Predicted-cost double-scaling bug.** The trust-region quality ratio + computed the predicted residual as `A_scaled @ (scaler * delta)`, which + applies the column scaling twice; the physical prediction is + `A_scaled @ delta`. This mildly distorted accept/reject decisions before, + and would have been catastrophic with the corrected scaler. + +### The column-scaling math (historical: the splice was reverted) + +Each LM step solves `min ‖A S δ + r‖² + λ‖δ‖²` with `S = diag(s)` and +applies `Δ = S δ`; equivalently `(AᵀA + λ S⁻²) Δ = −Aᵀr`. So choosing +`s_i` chooses the per-coordinate damping `d_i = λ / s_i²`: + +- **Levenberg** (`s = 1`, `d = λ`): one λ shared across curvatures that + differ by orders of magnitude — on raw BAL, exact Newton steps are + rejected for ~12 straight iterations at the default λ. +- **Marquardt/Moré** (`s = 1/n`, `d = λn²`; Ceres): fixes BA, but + amplifies weak columns without bound and re-denominates every + existing tuned λ — pyroki's collision benchmarks then converge to + worse local minima. +- **Legacy jaxls** (`s = (2+n)/(1+n) ∈ (1, 2]`): ≈ Levenberg, and the + convention downstream λs were tuned against. + +The splice evaluated here was whichever scales a column less: + +``` +s(n) = min( (2+n)/(1+n) , (3/2)/n ) # evaluated, then reverted +``` + +A min of two decreasing curves switches at their intersection, so fixing +the crossover at unit norm — the scale the Marquardt branch normalizes +*to* — pins the numerator: `c = s_legacy(1) = 3/2`. Below unit norm this +is exactly the legacy scaler; above it, exactly scale-invariant damping +`d = (4/9) λn²`. + +**Why it was reverted:** "below unit norm" turned out not to protect real +downstream problems — IK Jacobian columns (norms 6–11) and pose-graph +columns land *above* the crossover, so their tuned λs were silently +re-denominated by ~50×: pyroki IK-Beam fell from 99.8% to 94.6% success +(p99 error 0.13→16.7 mm) and the M3500 g2o example converged 5× worse. +Meanwhile the splice's own motivation dissolved: with the rejected-step +termination criteria gated (see "A real LM bug" below) and Nielsen lambda +escalation, the legacy scaler reaches *better* Ladybug-49 full-CG costs +than the splice ever did. See "Final configuration". + +A consequence worth stating plainly: performance numbers measured before +these fixes (including those of the previous Schur implementation on the +`schur-variable-elimination` branch, which reported full CG plateauing at +33,137) are not comparable with this table. The honest comparison is +within-run, against the same LM loop — which is what the table shows. + +## Caveats and the GPU question + +- The direct reduced solve is dense in the kept block: `O(n_keep²)` memory + and `O(n_keep³)` factorization. Fine to roughly 1–2k camera tangent + dimensions; beyond that the reduced system should go through a sparse + Cholesky (the co-visibility pattern is already computed by the plan). +- All numbers are CPU. On GPU, dense Cholesky is a weak spot, but + **Schur+CG is matrix-free end-to-end** — its assembly uses + `segment_sum` (contiguous reductions, no scatter atomics) and its + advantage comes from the reduced system's lower inner-iteration count, + which is sequential and cannot be parallelized away. So the GPU-favored + path is Schur+CG; validate there before relying on it. +- float64 is recommended for the direct path; float32 works (0 NaNs) with + accuracy bounded by the deliberate ~2e-3 Tikhonov floor. + +## Reproduce + +```bash +uv run --extra dev --extra docs python benchmarks/device_sweep.py # BA cost/time matrix + plots +uv run --extra dev --extra docs python benchmarks/float32_check.py # float32 robustness (Ladybug-49) +pytest tests/test_schur.py # correctness/validation suite +``` + +BAL data downloads to /tmp automatically. Raw arrays are saved to +`benchmarks/results/*.json`; `device_sweep.py --replot` regenerates plots. + +--- + +# CPU vs GPU study and GPU assembly optimization + +Added 2026-06-12. The numbers above are Apple-Silicon CPU. This section +re-runs the matched-iteration study on a second machine across **both CPU +and GPU**, and documents a GPU-driven optimization of the per-iteration +assembly. Hardware: AMD Threadripper PRO 5955WX (16C/32T) + NVIDIA RTX 4090, +float64. + +> **Config note:** the measurements in this part were taken with the +> (since-reverted) spliced scaler plus the termination fixes. The +> qualitative conclusions — assembly bottleneck, broadcast-vs-GEMM, the +> termination bug, Schur-vs-full-system gaps — are config-independent; +> the shipped-configuration BA tables live in "Final configuration" below. + +Harness: `benchmarks/device_sweep.py` (same matched-k methodology: exactly k +LM iterations, early termination off, warmup + min-of-3, per device via +`jax.device_put`). It compares full CG, cholmod (CPU only), Schur+dense, +Schur+CG, and full dense (small problems only). + +## Where the GPU time went: batched einsum vs broadcast + +Profiling one Schur+dense outer iteration on Ladybug-49/GPU (per-phase +timing) showed the cost was *not* the Cholesky factorization but the +**assembly of the reduced matrix**: + +| phase | before | after | +|-------------------|--------:|--------:| +| `prepare_schur` | 2.37 ms | 0.22 ms | +| `assemble_dense_S`| 7.75 ms | 0.55 ms | +| `solve_spd_scaled`| 1.56 ms | 1.56 ms | +| full inner step | 8.66 ms | 2.10 ms | + +The assembly is built from many small batched products +(`H_cc` Gram blocks `J^T J`, the cross blocks `W`, and the pair products of +`W V^{-1} W^T`). These were written as `jnp.einsum` / `lax.dot_general`. +Each has a **tiny contraction axis** — the residual dimension (2 for a +reprojection cost) or the eliminated-variable dimension (3 for a landmark). +On the 4090, XLA lowers such a batched GEMM to a kernel that is +**15–30× slower** than the mathematically identical +broadcast-multiply-sum. Measured on the 91,243-pair product of Ladybug-49: + +| form | time | +|-----------------------------------|--------:| +| `einsum('ptf,psf->pts', ...)` | 5.13 ms | +| `lax.dot_general` (batched) | 5.13 ms | +| `sum(a[...,None,:]*b[...,:,None,...], axis)` | **0.17 ms** | + +The fix (`_batched_gram`, `_batched_outer_last` in `_schur.py`) replaces the +five matrix-producing einsums with the broadcast form. It is bit-identical +to the einsum on random inputs (max diff 0.0) and changes no results: the +single damped step still matches a from-scratch full dense solve to +**2.6e-14 relative**, and all 53 tests pass. The dense inner step is now +bottlenecked by the irreducible Cholesky solve, as it should be. + +A dead end worth recording: replacing the pair-list assembly with one big +dense `camera×point` matmul is **slower** (8.4 ms), because the camera×point +occupancy is only ~8% on Ladybug — the dense matmul wastes >90% of its FLOPs +on structural zeros. The sparse pair structure must be kept. + +## Ladybug-49 — before/after, both devices + +Wall-clock at k=30 (float64, min of 3), Schur paths only (the methods the +optimization touches): + +| method | CPU before | CPU after | GPU before | GPU after | +|---------------|-----------:|----------:|-----------:|----------:| +| Schur + dense | 1.45 s | **1.04 s** | 0.339 s | **0.079 s** | +| Schur + CG | 15.5 s | (≈ same)¹ | 0.463 s | **0.357 s** | + +¹ The CPU Schur+CG re-run shared all 32 threads with another job and timed +slower (artifact, not a regression). The CG inner loop's matrix-*vector* +einsums were measured separately and are NOT GEMM-pathological — XLA already +lowers them well, so no broadcast rewrite applies there. +**GPU Schur+dense improves 4.3×; CPU Schur+dense 1.4×.** +Returned costs are identical to the pre-change run (≤1e-10). + +## Cross-device picture (optimized) + +Wall-clock at k=30 on Ladybug-49 (float64), all methods: + +| method | CPU | GPU | +|---------------|---------:|---------:| +| Schur + dense | **1.04 s** | **0.079 s** | +| Schur + CG | ~14.8 s | 0.357 s | +| cholmod | 4.16 s | n/a | +| full CG | 33.0 s | 0.88 s | + +- **Schur+dense is the fastest path on *both* devices** — including GPU, + correcting the earlier prediction above that "dense Cholesky is a GPU weak + spot, so Schur+CG is the GPU-favored path." With the assembly cost removed, + the 441×441 Cholesky is cheap on the 4090 and the dense path wins outright. +- The GPU rehabilitates full CG (33 s → 0.88 s), but it is still ~11× slower + than GPU Schur+dense. +- cholmod remains the best *full-system* CPU path, but Schur+dense beats it 4×. + +## Trafalgar-138 — the gap at scale (165,899 observations) + +Trafalgar-138 (138 cameras, 44,033 points, 165,899 observations) has a +**133,341-dim full system**; elimination reduces it to **1,242** dims. The +two slowest baselines were capped once their cost had visibly plateaued +(CPU full CG at k≤9, CPU Schur+CG at k≤13) — each further point cost +minutes of wall-clock to confirm an already-flat line. + +### A real LM bug, found via the first Trafalgar run + +The first run of this study showed every *exact-step* solver (Schur+dense +177,142; cholmod 178,707; full CG 179,039) plateauing well above where +inexact Schur+CG kept descending (173,653). Tracing one stalled solve +showed the mechanism: a **rejected** proposal whose (untaken) step was tiny +tripped the parameter-tolerance convergence criterion, which exits the +lambda-escalation loop before lambda can rise; on rejection lambda is kept, +so every subsequent outer iteration re-proposed the bit-identical step +forever. The cost criterion had already been gated to accepted proposals +(see the LM fixes above); the parameter criterion needed the same gate. +With the one-line fix (`converged_parameters ... & accepted`, +`_solvers.py`), **every solver converges to ~173.6k** — and the exact-step +methods now reach the *lowest* costs, as they should. + +Wall-clock / accepted cost after the fix (float64; the two slowest +baselines capped once their behavior was established): + +| method | device | k | time | cost | +|---------------|--------|---:|---------:|------------:| +| full CG | CPU | 9 | 97.4 s | 174,420 | +| full CG | GPU | 30 | 9.55 s | 173,660 | +| cholmod | CPU | 30 | 17.9 s | 173,659 | +| Schur + dense | CPU | 30 | 4.93 s | **173,632** | +| Schur + dense | GPU | 30 | **0.229 s** | **173,632** | +| Schur + CG | CPU | 30 | 354 s | 173,653 | +| Schur + CG | GPU | 30 | 5.33 s | 173,653 | + +- **Schur vs naive CG at scale:** at matched k=9, CPU full CG needs 97.4 s + (174,420) while GPU Schur+dense is at a comparable cost (174,989) in + 0.070 s — a **~1,400× wall-clock gap**, up from ~30× (CPU) / ~11× (GPU) + on Ladybug-49. Even on the same GPU, Schur+dense reaches a lower cost + than full CG ever does at 1/40th the wall-clock (0.23 s vs 9.6 s). +- **Schur+dense is now both the fastest and the most accurate** path on + both devices (173,632). CPU Schur+CG produces the same trajectory as GPU + Schur+CG ~50× slower — each CG matvec is a pass over 165,899 + observations, which the GPU parallelizes and the CPU serializes. +- cholmod again the best full-system path, and again beaten by Schur+dense + on both time (3.6× at k=30) and cost. + +--- + +# Final configuration (the shipped state) + +Added 2026-06-13. After the cross-device study above, two downstream +regressions traced to the spliced column scaler forced a configuration +change; this section records the final state and its benchmark numbers. +Per-problem A/B tables are reproducible via the benchmark suite +(`python -m benchmarks.suite --gate` against a committed baseline). + +**What shipped:** + +- **Legacy column scaler** (`1/(1+n) + 1`). The splice silently + re-denominated tuned lambdas for any problem with above-unit column + norms: pyroki IK-Beam fell from 99.8% to 94.6% success (p99 error + 0.13 → 16.7 mm) and the M3500 g2o example converged 5× worse + (679.8 vs 137.9). Both match `main` exactly under the final config. +- **Rejected-step termination gates + AL/lambda_max fixes + stable gain + ratio** (see "A real LM bug" above; all solvers converge ~2% lower on + Trafalgar with the gates in place). +- **Nielsen-style accelerating lambda escalation**: on each consecutive + rejection within an outer step, the escalation factor itself doubles, so + a grossly under-damped lambda recovers in O(log log) tries. This is the + rejection-side half of H.B. Nielsen (1999), "Damping Parameter in + Marquardt's Method", IMM-REP-1999-05, DTU; the decrease-on-acceptance + rule remains jaxls's plain halving. A no-op when the first proposal is + accepted, so tuned solves are unaffected. +- **Broadcast-form batched products** (`utils._batched_*`) in Schur + assembly and the block-Jacobi preconditioner (5–30× over einsum on GPU). + +**Lambda guidance for bundle adjustment:** BA curvature scales put +workable uniform damping around `lambda_initial ≈ 1e1–1e3` (a three-decade +plateau); the library default 5e-4 cold-starts the exact-step paths +through a long rejection sweep. The BA tables below pass +`lambda_initial=1e2` to *every* method, so the comparison isolates the +linear solver. + +## Bundle adjustment, main vs this PR (lambda_initial=1e2, float64) + +GPU, accepted cost / outer LM steps / wall-clock at the largest k within +a 25 s/solve budget; "(times out)" marks rows stopped by the budget. The +baseline is jaxls@main's full-system CG — what `solve()` runs on these +problems without this PR. (CPU numbers, cholmod, and the rest come from +`benchmarks/device_sweep.py`; main's full dense crashes on an int64/int32 +index bug this PR fixes.) + +| problem | full CG (main) | **Schur+dense (PR)** | Schur+CG (PR) | +|---|---|---|---| +| Toy | 5,913.89 / 16 it / 0.075 s | 5,913.72 / 16 it / **0.025 s** | 5,913.72 / 16 it / 0.032 s | +| Ladybug-49 | 27,394 / 30 it / 9.50 s | **26,779 / 30 it / 0.080 s** | 27,711 / 30 it / 0.50 s | +| Trafalgar-138 | 202,122 / 6 it / 31.3 s (times out) | **173,627 / 30 it / 0.234 s** | 173,627 / 30 it / 5.22 s | + +See `results/ba_comparison_gpu.png`: cost vs wall-clock per problem, one +marker per LM step starting from the shared step-0 initial cost +(running-best cost, since the unconverged full-CG baseline is non-monotone +in matched k; symlog x so t=0 and the multi-second baseline tail are both +legible). + +- **Ladybug-49:** Schur+dense reaches a lower cost than full CG ever does, + **119× faster** at matched 30 steps. +- **Trafalgar-138:** full CG times out at 6 steps far from convergence + (CPU exceeds 650 s *per solve* past the first step); **Schur+dense + reaches 173,627 in 0.234 s** — >130× faster at a 14% lower cost. Each + Schur+dense step is ~40× cheaper than a full-CG step here, so it both + takes useful steps and takes them faster. +- Toy: everything agrees at the optimum in the same step count; Schur + paths are fastest there too. + +## Reproduce (cross-device) + +```bash +uv run --extra dev --extra docs python benchmarks/device_sweep.py # CPU+GPU, all problems +uv run --extra dev --extra docs python benchmarks/device_sweep.py --devices gpu # GPU only +``` + +Requires a CUDA jaxlib for the GPU rows; falls back to CPU-only otherwise. +Raw arrays in `benchmarks/results/device_*.json`. + +--- + +# Performance ceiling: why single-solve time is what it is + +Added 2026-06-14. After the GPU-assembly and LM-fix wins above, a focused +optimization pass (target: 2x on single-solve wall-clock) profiled where the +time actually goes. The conclusion is structural and worth recording so +future work starts from it rather than re-deriving it. + +## The solve is host-dispatch-bound, not compute-bound + +On Trafalgar-138 (GPU, float64, lambda_initial=1e2), a 30-iteration dense +solve takes ~256 ms = ~8.5 ms/iter. But: + +- **Dispatch-only time (no `block_until_ready`) equals full time** (256 ms == + 256 ms): the host enqueue loop, not the device, is the critical path. The + GPU is mostly idle. +- The whole solve is **0.64 GFLOP** against **1.76 GB** of memory traffic — + it runs at ~0.01% of the 4090's float64 peak. Nothing is compute-limited. +- Each iteration compiles to **~175 fusions + ~9 custom-calls ≈ 184 kernel + launches**, run serially (~46 µs each). Per-iteration kernel count is + nearly identical for toy (138) and Trafalgar (175) — the overhead is + structural, which is why on small problems the linear solve is ~1% of + wall-clock and the other ~99% is this fixed per-iteration chain. + +So the bottleneck is **the number of small serial kernels per outer step**, +each doing little work. This reframes every solver-internal optimization: + +| lever | measured effect | why | +|---|---|---| +| GPU broadcast-vs-einsum assembly | ~4x on assembly (landed) | real, but assembly is ~0.5 ms of 8.5 ms/iter | +| Hoist column-norm scaler out of loop | ~2-5%/iter (landed) | one fewer scatter/iter | +| float32 Cholesky + iterative refinement | **rejected** | real Schur S is ill-conditioned (cond ~1e6); f32 refinement stalls at ~4% error | +| gather vs vmap(dynamic_slice) in multiply | ~0% | XLA lowers both to the same kernel here | +| einsum->broadcast in reduced_rhs/back_sub | ~1% | matvec einsums are not GEMM-pathological | +| analytic vs autodiff reprojection Jacobian | ~0% (cost identical) | the Jacobian is one big vmapped fusion either way; op *count* is high but it is not the dispatch bottleneck | +| remove per-iter timestamp callback | ~0.5% | JAX pipelines the callback; it is not the bottleneck | + +The linear solver — where the obvious optimizations live — is **1-3% of +per-iteration wall-clock** on these problems. A faster Cholesky or CG matvec +cannot move the total. + +## What would actually move it + +1. **Batch/vmap across problems.** Because the GPU is idle between + dispatches, solving N problems at once amortizes the per-iteration kernel + chain over N. This is how downstream users already get throughput + (pyroki's batch-2000 vmapped IK); per-problem cost collapses. The single + highest-leverage path, and it needs no solver changes. +2. **Fewer, larger kernels per outer step.** A real single-solve 2x needs the + per-iteration graph to compile to a handful of big fused kernels instead + of ~184 small ones — a structural rewrite of the cost/Jacobian/assemble/ + solve chain, research-grade and risky, with uncertain payoff. +3. **Analytic Jacobians help op *count* but not this bottleneck** on GPU; + they matter more on CPU and for compile time. Worth offering for + op-heavy costs, but not the single-solve GPU lever it first appears to be. + +The honest summary: a 2x single-solve speedup on toy/Ladybug/Trafalgar is +**not reachable from solver-internal changes** — the bottleneck is the fixed +per-iteration kernel-dispatch chain. Throughput (batching) is the productive +direction. diff --git a/benchmarks/results/ba_comparison_gpu.png b/benchmarks/results/ba_comparison_gpu.png new file mode 100644 index 0000000..15baf08 Binary files /dev/null and b/benchmarks/results/ba_comparison_gpu.png differ diff --git a/benchmarks/results/ba_initial_cost.json b/benchmarks/results/ba_initial_cost.json new file mode 100644 index 0000000..dc3341e --- /dev/null +++ b/benchmarks/results/ba_initial_cost.json @@ -0,0 +1,5 @@ +{ + "toy": 107025.0897669929, + "ladybug49": 1701824.9213558466, + "trafalgar138": 40327138.03197177 +} \ No newline at end of file diff --git a/benchmarks/results/device_ladybug49_tuned_final.json b/benchmarks/results/device_ladybug49_tuned_final.json new file mode 100644 index 0000000..0e81284 --- /dev/null +++ b/benchmarks/results/device_ladybug49_tuned_final.json @@ -0,0 +1,254 @@ +{ + "title": "Ladybug-49 (49 cams, 7,776 pts, 31,843 obs)", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "runs": { + "gpu:full CG": { + "method": "full CG", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 44398.523841905386, + 34263.90846859208, + 31559.920150212267, + 31269.906259454496, + 30270.60950791607, + 28232.591992503345, + 27995.29747203082, + 27576.625470306986, + 27064.72368409084 + ], + "times": [ + 0.028773650992661715, + 0.7660990839940496, + 3.1202277870033868, + 3.304774535004981, + 3.905058675969485, + 3.939488784992136, + 6.039854708011262, + 7.300608842982911, + 8.477442481962498 + ], + "budget_stopped_at_k": 30 + }, + "gpu:Schur + dense": { + "method": "Schur + dense", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 42163.8945511792, + 33773.20484760012, + 31485.357497420984, + 30166.83731704591, + 28270.16262435279, + 27186.794129008587, + 26885.876038218405, + 26845.317993228593, + 26779.12384370059 + ], + "times": [ + 0.003518136974889785, + 0.006013084028381854, + 0.011377584014553577, + 0.016695685975719243, + 0.024663987976964563, + 0.03528096800437197, + 0.04850225895643234, + 0.06450995797058567, + 0.08039847196778283 + ] + }, + "gpu:Schur + CG": { + "method": "Schur + CG", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 41534.049793307335, + 33791.414774562974, + 31962.852301198123, + 30351.47063657897, + 31279.72555445456, + 30093.30798425196, + 29426.494285873276, + 29690.08417888387, + 27711.259421350715 + ], + "times": [ + 0.0036068640183657408, + 0.015862094005569816, + 0.04587929096305743, + 0.08762639702763408, + 0.10959422902669758, + 0.17219199001556262, + 0.25412007700651884, + 0.2604773289640434, + 0.49532757402630523 + ] + }, + "cpu:full CG": { + "method": "full CG", + "device": "cpu", + "ks": [ + 1, + 2 + ], + "costs": [ + 44384.01544574153, + 34266.891352345156 + ], + "times": [ + 0.9358694970142096, + 26.498676098010037 + ], + "budget_stopped_at_k": 2 + }, + "cpu:cholmod": { + "method": "cholmod", + "device": "cpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 42163.89422732871, + 33773.20509436272, + 31485.358101100435, + 30166.831465946932, + 28270.185359737818, + 27186.73260019729, + 26885.908331466988, + 26845.02636817331, + 26779.463886570076 + ], + "times": [ + 0.19488017697585747, + 0.39569709298666567, + 0.7258570889825933, + 0.8911163749871776, + 1.4267730970168486, + 1.9956432770122774, + 2.6679983059875667, + 3.5063375490135513, + 4.412380089983344 + ] + }, + "cpu:Schur + dense": { + "method": "Schur + dense", + "device": "cpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 42163.89455125185, + 33773.20484759631, + 31485.357497056815, + 30166.827837645324, + 28270.176637799275, + 27186.795871402443, + 26885.876455855425, + 26845.302264487957, + 26779.117871800758 + ], + "times": [ + 0.05062148900469765, + 0.07989321101922542, + 0.1334534739726223, + 0.19153855397598818, + 0.26279049302684143, + 0.3881818869849667, + 0.5241721370257437, + 0.6881244990509003, + 0.8598640309646726 + ] + }, + "cpu:Schur + CG": { + "method": "Schur + CG", + "device": "cpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24 + ], + "costs": [ + 41534.049793297425, + 33791.417487668456, + 31879.86189440849, + 30309.39491733489, + 30309.39491733489, + 30093.24202880105, + 29429.740870901485, + 27498.42651762844 + ], + "times": [ + 0.10916121600894257, + 0.7404052279889584, + 2.155251582036726, + 4.438389771035872, + 6.653298849996645, + 8.59673287998885, + 12.974949212977663, + 19.304635751002934 + ], + "budget_stopped_at_k": 24 + } + } +} \ No newline at end of file diff --git a/benchmarks/results/device_ladybug49_tuned_main.json b/benchmarks/results/device_ladybug49_tuned_main.json new file mode 100644 index 0000000..bac9f1c --- /dev/null +++ b/benchmarks/results/device_ladybug49_tuned_main.json @@ -0,0 +1,107 @@ +{ + "title": "Ladybug-49 (49 cams, 7,776 pts, 31,843 obs)", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "runs": { + "gpu:full CG": { + "method": "full CG", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 44426.98569366009, + 34262.87463568737, + 31561.622786884298, + 31174.030065448173, + 29761.89299330177, + 85377.80319279799, + 28014.55935304365, + 26926.801560326552, + 27394.112501388066 + ], + "times": [ + 0.030902878032065928, + 0.7624542599660344, + 3.204482298984658, + 3.3086227129679173, + 3.385593814018648, + 4.262224388017785, + 5.192163272004109, + 5.0953683410189115, + 9.5042430239846 + ] + }, + "cpu:full CG": { + "method": "full CG", + "device": "cpu", + "ks": [ + 1, + 2 + ], + "costs": [ + 44384.01544574153, + 34266.891352345156 + ], + "times": [ + 1.111208993010223, + 30.737980152014643 + ], + "budget_stopped_at_k": 2 + }, + "cpu:cholmod": { + "method": "cholmod", + "device": "cpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 42163.89422732871, + 33773.20509436272, + 31485.358101100435, + 30166.831465946932, + 28270.185359737818, + 27452.166027966265, + 27257.022959206173, + 27163.23447947068, + 27105.95375585732 + ], + "times": [ + 0.1859285940299742, + 0.2793644049670547, + 0.720436344970949, + 0.9424537149607204, + 1.4259259660029784, + 1.7930776689900085, + 2.411875197023619, + 3.5361848100437783, + 4.354561993968673 + ] + } + } +} \ No newline at end of file diff --git a/benchmarks/results/device_toy_tuned_final.json b/benchmarks/results/device_toy_tuned_final.json new file mode 100644 index 0000000..de43607 --- /dev/null +++ b/benchmarks/results/device_toy_tuned_final.json @@ -0,0 +1,321 @@ +{ + "title": "Toy BA (30 cams, 700 points)", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "runs": { + "gpu:full CG": { + "method": "full CG", + "device": "gpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6053.6714787005385, + 5949.332078946396, + 5943.341410810073, + 5940.122536640705, + 5928.503624263827, + 5916.455003213524, + 5913.795128889351, + 5913.723888461816 + ], + "times": [ + 0.00274576002266258, + 0.009499750973191112, + 0.01673922804184258, + 0.023683217994403094, + 0.03559224400669336, + 0.04438814701279625, + 0.05869549699127674, + 0.07327541202539578 + ] + }, + "gpu:Schur + dense": { + "method": "Schur + dense", + "device": "gpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6042.873293249699, + 5948.93243915592, + 5943.085723824704, + 5939.901807994504, + 5928.377259055442, + 5916.42972567745, + 5913.79360398411, + 5913.723884249477 + ], + "times": [ + 0.0023326000082306564, + 0.0038166700396686792, + 0.005389527010265738, + 0.006984325998928398, + 0.010033948987256736, + 0.013113197986967862, + 0.01932212000247091, + 0.02533750794827938 + ] + }, + "gpu:Schur + CG": { + "method": "Schur + CG", + "device": "gpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6046.5152995471735, + 5948.910686844465, + 5943.086504616785, + 5939.908603844247, + 5928.3794181895155, + 5916.429465397216, + 5913.793649813369, + 5913.723885442613 + ], + "times": [ + 0.0015120220487006009, + 0.003221767023205757, + 0.0049773299833759665, + 0.0068494210136123, + 0.01104865298839286, + 0.0153563889907673, + 0.0235561880399473, + 0.03162046102806926 + ] + }, + "gpu:full dense": { + "method": "full dense", + "device": "gpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6042.873293249704, + 5948.932439155922, + 5943.085723824706, + 5939.901807994507, + 5928.377259055444, + 5916.429725677453, + 5913.793603984112, + 5913.723884249483 + ], + "times": [ + 0.09415952005656436, + 0.1771324870060198, + 0.2653433160157874, + 0.3536949289846234, + 0.5299657309660688, + 0.7063257619738579, + 1.0603889460326172, + 1.4172241009655409 + ] + }, + "cpu:full CG": { + "method": "full CG", + "device": "cpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6053.671478590877, + 5949.332234427411, + 5943.3414681106615, + 5940.122556006943, + 5928.5037083917605, + 5916.455008542737, + 5913.795128643678, + 5913.723888460615 + ], + "times": [ + 0.008974610012955964, + 0.025270113022997975, + 0.05573083704803139, + 0.12637252401327714, + 0.22535767004592344, + 0.20870086597278714, + 0.3065639559645206, + 0.46634163899580017 + ] + }, + "cpu:cholmod": { + "method": "cholmod", + "device": "cpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6042.873303465806, + 5948.9324409982255, + 5943.085724852076, + 5939.901810914198, + 5928.377280416571, + 5916.429763473134, + 5913.793609929149, + 5913.723885359342 + ], + "times": [ + 0.011362346005626023, + 0.019273590005468577, + 0.02797503600595519, + 0.037693549995310605, + 0.05130907503189519, + 0.07437769399257377, + 0.10483085503801703, + 0.14864225598284975 + ] + }, + "cpu:Schur + dense": { + "method": "Schur + dense", + "device": "cpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6042.873293249703, + 5948.932439155927, + 5943.085723824703, + 5939.901807994512, + 5928.377259055448, + 5916.429725677453, + 5913.7936039841115, + 5913.72388424948 + ], + "times": [ + 0.0074668360175564885, + 0.013688042003195733, + 0.012365487054921687, + 0.018556818016804755, + 0.026108724996447563, + 0.03553987597115338, + 0.05397293099667877, + 0.06639832200016826 + ] + }, + "cpu:Schur + CG": { + "method": "Schur + CG", + "device": "cpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6046.515299547178, + 5948.910686844468, + 5943.086504616787, + 5939.908603844245, + 5928.379418189511, + 5916.4294653972165, + 5913.793649813366, + 5913.723885442619 + ], + "times": [ + 0.006070072005968541, + 0.013979682989884168, + 0.030009463021997362, + 0.06063217303017154, + 0.11540134600363672, + 0.1693657399737276, + 0.22181008901679888, + 0.3048389629693702 + ] + }, + "cpu:full dense": { + "method": "full dense", + "device": "cpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6042.873293249706, + 5948.932439155925, + 5943.085723824703, + 5939.901807994506, + 5928.377259055445, + 5916.429725677452, + 5913.793603984118, + 5913.723884249483 + ], + "times": [ + 0.24085840297630057, + 0.5119776339852251, + 0.9035006380290724, + 0.9187898819800466, + 1.5281315150205046, + 1.9343808419653215, + 2.998356012045406, + 4.268886885023676 + ] + } + } +} \ No newline at end of file diff --git a/benchmarks/results/device_toy_tuned_main.json b/benchmarks/results/device_toy_tuned_main.json new file mode 100644 index 0000000..cd4bd9f --- /dev/null +++ b/benchmarks/results/device_toy_tuned_main.json @@ -0,0 +1,117 @@ +{ + "title": "Toy BA (30 cams, 700 points)", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "runs": { + "gpu:full CG": { + "method": "full CG", + "device": "gpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6053.67147877712, + 5949.3321719944715, + 5943.341442157899, + 5940.122492859633, + 5928.503715905392, + 5916.455003385644, + 5913.8910021555475, + 5913.891003611683 + ], + "times": [ + 0.002824869006872177, + 0.009612152993213385, + 0.017356152005959302, + 0.024621959018986672, + 0.036923223990015686, + 0.04604607500368729, + 0.06085291696945205, + 0.0747278219787404 + ] + }, + "cpu:full CG": { + "method": "full CG", + "device": "cpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6053.671478590877, + 5949.332234427411, + 5943.3414681106615, + 5940.122556006943, + 5928.5037083917605, + 5916.455008542737, + 5913.89100359044, + 5913.89100359044 + ], + "times": [ + 0.010096842015627772, + 0.02423832204658538, + 0.06142262800130993, + 0.09281042398652062, + 0.11602018598932773, + 0.2558243770035915, + 0.3149716969928704, + 0.3923199909622781 + ] + }, + "cpu:cholmod": { + "method": "cholmod", + "device": "cpu", + "ks": [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16 + ], + "costs": [ + 6042.873303465806, + 5948.9324409982255, + 5943.085724852076, + 5939.901810914198, + 5928.377280416571, + 5916.429763473134, + 5913.888169365026, + 5913.888169365026 + ], + "times": [ + 0.011468568991404027, + 0.026372456981334835, + 0.03235538996523246, + 0.04263989697210491, + 0.059231457009445876, + 0.07388252799864858, + 0.1054272599867545, + 0.1467865909799002 + ] + } + } +} \ No newline at end of file diff --git a/benchmarks/results/device_trafalgar138_tuned_final.json b/benchmarks/results/device_trafalgar138_tuned_final.json new file mode 100644 index 0000000..dea08cc --- /dev/null +++ b/benchmarks/results/device_trafalgar138_tuned_final.json @@ -0,0 +1,205 @@ +{ + "title": "Trafalgar-138 (138 cams, 44,033 pts)", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "runs": { + "gpu:full CG": { + "method": "full CG", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6 + ], + "costs": [ + 335312.0371299817, + 217982.95256344584, + 210088.90658246638, + 202087.13497821003 + ], + "times": [ + 0.053141399985179305, + 9.254745967045892, + 22.7864464269951, + 32.050926144002005 + ], + "budget_stopped_at_k": 6 + }, + "gpu:Schur + dense": { + "method": "Schur + dense", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 219842.52620743986, + 215171.67481489124, + 209794.85628479754, + 202080.73254066124, + 185178.1353425056, + 176725.4948564962, + 175905.28549251636, + 174926.77429803723, + 173626.54778743518 + ], + "times": [ + 0.009056083974428475, + 0.017085411003790796, + 0.03354103502351791, + 0.05008119798731059, + 0.0742526879766956, + 0.10611930402228609, + 0.14117078896379098, + 0.18932925001718104, + 0.2342980530229397 + ] + }, + "gpu:Schur + CG": { + "method": "Schur + CG", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 235373.83568897998, + 215444.25520515544, + 209779.3409786772, + 202075.48918591783, + 185173.30238355056, + 176730.77446211388, + 175904.88971609776, + 175094.4622275972, + 173626.6124626888 + ], + "times": [ + 0.0066279839957132936, + 0.16602578799938783, + 0.5144766619778238, + 1.0812328499741852, + 1.5759166239877231, + 2.1413880109903403, + 2.969115466985386, + 4.255199619976338, + 5.214814551000018 + ] + }, + "cpu:cholmod": { + "method": "cholmod", + "device": "cpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24 + ], + "costs": [ + 219842.52639703467, + 215171.6752305055, + 209794.85854522337, + 202080.74626005784, + 185178.2492307129, + 176725.59248664894, + 175905.37833155054, + 174926.95455926386 + ], + "times": [ + 0.7661450759624131, + 1.385755588999018, + 2.842199510021601, + 4.266641339985654, + 5.704402634990402, + 8.575095239037182, + 11.813907900010236, + 15.040278232016135 + ], + "budget_stopped_at_k": 24 + }, + "cpu:Schur + dense": { + "method": "Schur + dense", + "device": "cpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "costs": [ + 219842.5262074525, + 215171.67481488996, + 209794.8562847995, + 202080.73253981167, + 185178.13533317327, + 176725.49329096868, + 175905.28540574195, + 174926.78904950432, + 173626.54781461268 + ], + "times": [ + 0.24849610100500286, + 0.4125385520164855, + 0.6918708649463952, + 0.9679526060353965, + 1.3943158480105922, + 1.9753825340303592, + 2.6404840190079994, + 3.6093236900051124, + 4.368649238022044 + ] + }, + "cpu:Schur + CG": { + "method": "Schur + CG", + "device": "cpu", + "ks": [ + 1, + 2, + 4 + ], + "costs": [ + 235373.83568962163, + 215444.11739599748, + 209779.3440807712 + ], + "times": [ + 0.37656114296987653, + 13.602246335998643, + 41.136300699959975 + ], + "budget_stopped_at_k": 4 + } + } +} \ No newline at end of file diff --git a/benchmarks/results/device_trafalgar138_tuned_main.json b/benchmarks/results/device_trafalgar138_tuned_main.json new file mode 100644 index 0000000..d6d707b --- /dev/null +++ b/benchmarks/results/device_trafalgar138_tuned_main.json @@ -0,0 +1,74 @@ +{ + "title": "Trafalgar-138 (138 cams, 44,033 pts)", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24, + 30 + ], + "runs": { + "gpu:full CG": { + "method": "full CG", + "device": "gpu", + "ks": [ + 1, + 2, + 4, + 6 + ], + "costs": [ + 334636.37499285047, + 218175.07825511973, + 210165.0481230358, + 202121.978937964 + ], + "times": [ + 0.06612372200470418, + 9.014368803997058, + 22.027035192993935, + 31.26639480201993 + ], + "budget_stopped_at_k": 6 + }, + "cpu:cholmod": { + "method": "cholmod", + "device": "cpu", + "ks": [ + 1, + 2, + 4, + 6, + 9, + 13, + 18, + 24 + ], + "costs": [ + 219842.52639703467, + 215171.6752305055, + 209794.85854522337, + 202080.74626005784, + 185178.2492307129, + 185178.2492307129, + 185178.2492307129, + 185178.2492307129 + ], + "times": [ + 0.7047256489749998, + 1.3251609980361536, + 2.790073114039842, + 4.0458978620008565, + 6.1389207990141585, + 8.350827991962433, + 11.543715320003685, + 15.13785265799379 + ], + "budget_stopped_at_k": 24 + } + } +} \ No newline at end of file diff --git a/benchmarks/results/example_traces.png b/benchmarks/results/example_traces.png new file mode 100644 index 0000000..438ed70 Binary files /dev/null and b/benchmarks/results/example_traces.png differ diff --git a/benchmarks/suite/__init__.py b/benchmarks/suite/__init__.py new file mode 100644 index 0000000..71caa17 --- /dev/null +++ b/benchmarks/suite/__init__.py @@ -0,0 +1,14 @@ +"""jaxls benchmark + regression suite. + +A single entry point (`python -m benchmarks.suite`) that runs a set of +workloads, reduces each to named scalar metrics, and either: + + - writes a results JSON and a markdown report (reporting mode), or + - diffs the run against a committed baseline JSON and exits nonzero if any + metric regressed beyond its tolerance (gate mode, for CI / hill climbing). + +Workloads (see `workloads.py`): bundle-adjustment matrix (Schur vs +full-system, CPU+GPU), example-notebook convergence, pyroki IK downstream, +and float32 robustness. The metric layer (`metrics.py`) is the contract: +add a metric there and both the report and the regression gate pick it up. +""" diff --git a/benchmarks/suite/__main__.py b/benchmarks/suite/__main__.py new file mode 100644 index 0000000..82dcd49 --- /dev/null +++ b/benchmarks/suite/__main__.py @@ -0,0 +1,107 @@ +"""jaxls benchmark + regression suite — one entry point. + +Run via uv from the repo root (the `dev` extra brings tyro; `docs` brings +matplotlib + scikit-sparse): + + uv run --extra dev --extra docs python -m benchmarks.suite # full run, report to stdout + uv run --extra dev --extra docs python -m benchmarks.suite --quick # fast GPU-only inner loop + uv run --extra dev --extra docs python -m benchmarks.suite --gate # diff vs baseline, exit 1 on regression + uv run --extra dev --extra docs python -m benchmarks.suite --update-baseline + uv run --extra dev --extra docs python -m benchmarks.suite --only bundle_adjustment float32_robustness +""" + +from __future__ import annotations + +import dataclasses +import json +import sys +from pathlib import Path + +import jax + +jax.config.update("jax_enable_x64", True) + +import tyro # noqa: E402 + +from . import baseline as bl # noqa: E402 +from . import workloads as wl # noqa: E402 +from .metrics import Metric, WorkloadResult # noqa: E402 + + +@dataclasses.dataclass +class Args: + quick: bool = False + """Fast inner-loop tier: GPU only, fewer problems/k, no slow CPU baselines.""" + gate: bool = False + """Regression-gate mode: diff against the committed baseline and exit 1 if + any metric regressed beyond tolerance.""" + update_baseline: bool = False + """Overwrite the committed baseline with this run's metrics.""" + only: tuple[str, ...] = () + """Run only these workloads (default: all); see benchmarks/README.md for + the choices, or pass an unknown name to have them listed.""" + report_path: Path = Path("benchmarks/results/suite_report.md") + results_path: Path = Path("benchmarks/results/suite_results.json") + repeats: int = 3 + + +def main(args: Args) -> None: + # --quick truncates the k-ladder, so its costs differ from the full tier + # by design; gating/updating the shared baseline from a quick run would + # compare apples to oranges. The baseline is always a full-tier artifact. + if args.quick and (args.gate or args.update_baseline): + sys.exit( + "--quick is for eyeballing during development; it is not comparable " + "to the (full-tier) baseline. Use a full run for --gate / " + "--update-baseline." + ) + cfg = wl.SuiteConfig(quick=args.quick, repeats=args.repeats) + names = args.only or wl.DEFAULT + bad = [n for n in names if n not in wl.ALL] + if bad: + sys.exit(f"unknown workload(s): {bad}; choices: {list(wl.ALL)}") + + results = [] + for name in names: + print(f"\n########## {name} ##########", flush=True) + try: + results.append(wl.ALL[name](cfg)) + except Exception as e: # noqa: BLE001 + import traceback + + traceback.print_exc() + results.append(WorkloadResult(name=name, skipped=f"crashed: {e}")) + + meta = { + "jaxls": __import__("jaxls").__file__, + "devices": wl._device_list(cfg), + "quick": args.quick, + "baseline": str(bl.BASELINE_PATH.name) if bl.BASELINE_PATH.exists() else "none", + } + run_json = bl.results_to_json(results, meta) + args.results_path.parent.mkdir(parents=True, exist_ok=True) + args.results_path.write_text(json.dumps(run_json, indent=2)) + print(f"\nwrote {args.results_path}") + + if args.update_baseline: + bl.BASELINE_PATH.write_text(json.dumps(run_json, indent=2)) + print(f"updated baseline: {bl.BASELINE_PATH}") + return + + current = {k: Metric.from_json(k, v) for k, v in run_json["metrics"].items()} + baseline = bl.load_metrics(bl.BASELINE_PATH) if bl.BASELINE_PATH.exists() else {} + diffs = bl.diff_against_baseline(current, baseline) + report = bl.render_report(diffs, run_json["skipped"], meta) + args.report_path.write_text(report) + print(f"wrote {args.report_path}\n") + print(report) + + if args.gate: + n = sum(d.regressed for d in diffs) + if not baseline: + sys.exit("gate mode but no baseline committed; run --update-baseline first") + sys.exit(1 if n else 0) + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/benchmarks/suite/baseline.json b/benchmarks/suite/baseline.json new file mode 100644 index 0000000..8fc0499 --- /dev/null +++ b/benchmarks/suite/baseline.json @@ -0,0 +1,318 @@ +{ + "meta": { + "jaxls": "/home/brent/jaxls/src/jaxls/__init__.py", + "devices": [ + "gpu", + "cpu" + ], + "quick": false, + "baseline": "baseline.json" + }, + "metrics": { + "toy.gpu.full_cg.cost": { + "value": 5913.723888460381, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "toy.gpu.full_cg.time": { + "value": 0.07491348299663514, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "toy.gpu.schur_dense.cost": { + "value": 5913.723884249479, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "toy.gpu.schur_dense.time": { + "value": 0.02682858897605911, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "toy.gpu.schur_cg.cost": { + "value": 5913.723885442614, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "toy.gpu.schur_cg.time": { + "value": 0.032702618977054954, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "toy.cpu.full_cg.cost": { + "value": 5913.723888460916, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "toy.cpu.full_cg.time": { + "value": 0.2743255959940143, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "toy.cpu.schur_dense.cost": { + "value": 5913.723884249483, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "toy.cpu.schur_dense.time": { + "value": 0.0633065280271694, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "toy.cpu.schur_cg.cost": { + "value": 5913.72388544262, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "toy.cpu.schur_cg.time": { + "value": 0.3201562390313484, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.gpu.full_cg.cost": { + "value": 26997.88885164329, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.gpu.full_cg.time": { + "value": 8.418146082025487, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.gpu.schur_dense.cost": { + "value": 26779.09611144552, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.gpu.schur_dense.time": { + "value": 0.07681672601029277, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.gpu.schur_cg.cost": { + "value": 26874.3014152529, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.gpu.schur_cg.time": { + "value": 0.5777458370430395, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.cpu.full_cg.cost": { + "value": 34263.44280161518, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "budget-stopped at k=2" + }, + "ladybug49.cpu.full_cg.time": { + "value": 29.228925258968957, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "budget-stopped at k=2" + }, + "ladybug49.cpu.schur_dense.cost": { + "value": 26779.116647798262, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.cpu.schur_dense.time": { + "value": 0.8472943399683572, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "ladybug49.cpu.schur_cg.cost": { + "value": 26887.171530858235, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "budget-stopped at k=30" + }, + "ladybug49.cpu.schur_cg.time": { + "value": 26.030747681041248, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "budget-stopped at k=30" + }, + "trafalgar138.gpu.full_cg.cost": { + "value": 210085.48831159686, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "budget-stopped at k=4" + }, + "trafalgar138.gpu.full_cg.time": { + "value": 26.193290175986476, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "budget-stopped at k=4" + }, + "trafalgar138.gpu.schur_dense.cost": { + "value": 173626.54810044504, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "trafalgar138.gpu.schur_dense.time": { + "value": 0.22934570698998868, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "trafalgar138.gpu.schur_cg.cost": { + "value": 173626.61255608813, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "trafalgar138.gpu.schur_cg.time": { + "value": 5.172860280028544, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "trafalgar138.cpu.full_cg.cost": { + "value": 218215.04808663815, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "budget-stopped at k=2" + }, + "trafalgar138.cpu.full_cg.time": { + "value": 610.3361833519884, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "budget-stopped at k=2" + }, + "trafalgar138.cpu.schur_dense.cost": { + "value": 173626.54786369647, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "" + }, + "trafalgar138.cpu.schur_dense.time": { + "value": 4.959632423997391, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "" + }, + "trafalgar138.cpu.schur_cg.cost": { + "value": 209779.3431301261, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.02, + "abs_tol": 0.0, + "notes": "budget-stopped at k=4" + }, + "trafalgar138.cpu.schur_cg.time": { + "value": 30.598396679968573, + "unit": "s", + "direction": "lower_better", + "rel_tol": 0.3, + "abs_tol": 0.0, + "notes": "budget-stopped at k=4" + }, + "float32.ladybug49.nan_count": { + "value": 0.0, + "unit": "", + "direction": "lower_better", + "rel_tol": 0.0, + "abs_tol": 0.0, + "notes": "" + }, + "float32.ladybug49.best_cost": { + "value": 28277.06640625, + "unit": "cost", + "direction": "lower_better", + "rel_tol": 0.05, + "abs_tol": 0.0, + "notes": "float32 direct Schur path" + } + }, + "skipped": {} +} \ No newline at end of file diff --git a/benchmarks/suite/baseline.py b/benchmarks/suite/baseline.py new file mode 100644 index 0000000..f4e078b --- /dev/null +++ b/benchmarks/suite/baseline.py @@ -0,0 +1,99 @@ +"""Load/save run JSON, diff against a committed baseline, render the report.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +from .metrics import Metric, WorkloadResult + +SUITE_DIR = Path(__file__).resolve().parent +BASELINE_PATH = SUITE_DIR / "baseline.json" + + +def results_to_json(results: list[WorkloadResult], meta: dict) -> dict: + metrics: dict[str, dict] = {} + skipped: dict[str, str] = {} + for r in results: + if r.skipped is not None: + skipped[r.name] = r.skipped + for m in r.metrics: + metrics[m.key] = m.to_json() + return {"meta": meta, "metrics": metrics, "skipped": skipped} + + +def load_metrics(path: Path) -> dict[str, Metric]: + data = json.loads(path.read_text()) + return {k: Metric.from_json(k, v) for k, v in data.get("metrics", {}).items()} + + +@dataclass +class Diff: + key: str + current: Metric + baseline: Metric | None + regressed: bool + + @property + def rel_change(self) -> float | None: + if self.baseline is None or self.baseline.value == 0: + return None + return (self.current.value - self.baseline.value) / abs(self.baseline.value) + + +def diff_against_baseline( + current: dict[str, Metric], baseline: dict[str, Metric] +) -> list[Diff]: + diffs = [] + for key, cur in current.items(): + base = baseline.get(key) + diffs.append(Diff(key, cur, base, cur.regressed(base) if base else False)) + return diffs + + +def render_report(diffs: list[Diff], skipped: dict[str, str], meta: dict) -> str: + lines = [ + "# jaxls benchmark suite report", + "", + f"- jaxls: `{meta.get('jaxls', '?')}`", + f"- devices: {meta.get('devices', '?')} | quick: {meta.get('quick', False)}", + f"- baseline: `{meta.get('baseline', 'none')}`", + "", + ] + regressions = [d for d in diffs if d.regressed] + if regressions: + lines += [f"## ⚠️ {len(regressions)} regression(s)", ""] + lines += ["| metric | baseline | current | change |", "|---|---:|---:|---:|"] + for d in regressions: + # regressed is only set when a baseline exists (see + # diff_against_baseline), so d.baseline is non-None here. + assert d.baseline is not None + rc = f"{d.rel_change:+.1%}" if d.rel_change is not None else "—" + lines.append( + f"| {d.key} | {d.baseline.value:.6g} | {d.current.value:.6g} | {rc} |" + ) + lines.append("") + else: + lines += ["## ✅ no regressions", ""] + + lines += [ + "## all metrics", + "", + "| metric | current | baseline | change | unit |", + "|---|---:|---:|---:|---|", + ] + for d in sorted(diffs, key=lambda x: x.key): + b = f"{d.baseline.value:.6g}" if d.baseline else "—" + rc = f"{d.rel_change:+.1%}" if d.rel_change is not None else "—" + flag = " ⚠️" if d.regressed else "" + note = f" _{d.current.notes}_" if d.current.notes else "" + lines.append( + f"| {d.key}{flag} | {d.current.value:.6g} | {b} | {rc} | " + f"{d.current.unit}{note} |" + ) + if skipped: + lines += ["", "## skipped workloads", ""] + for name, why in skipped.items(): + lines.append(f"- **{name}**: {why}") + return "\n".join(lines) + "\n" diff --git a/benchmarks/suite/metrics.py b/benchmarks/suite/metrics.py new file mode 100644 index 0000000..500e1d7 --- /dev/null +++ b/benchmarks/suite/metrics.py @@ -0,0 +1,73 @@ +"""The metric contract: a flat namespace of named scalar measurements. + +Every workload emits `Metric`s. A metric has a value, a unit, and a +`direction` (is lower or higher better?) plus a default regression +tolerance. The regression gate (`baseline.py`) and the report +(`report.py`) both consume this same list, so adding a measurement in one +place makes it show up in the report and the gate automatically. + +Metric `key`s are stable identifiers ("ladybug49.gpu.schur_dense.cost") — +they are the JSON keys and the baseline keys, so keep them deterministic. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + + +@dataclass(frozen=True) +class Metric: + key: str + value: float + unit: str = "" + direction: Literal["lower_better", "higher_better", "neutral"] = "lower_better" + rel_tol: float = 0.10 + """Fractional regression tolerance: a `lower_better` metric is flagged + when it grows past value*(1+rel_tol). Time metrics are noisy, so they + default looser than correctness metrics (set per-metric at creation).""" + abs_tol: float = 0.0 + """Absolute slack added to the tolerance band (for near-zero values).""" + notes: str = "" + + def regressed(self, baseline: "Metric") -> bool: + """True if `self` is worse than `baseline` beyond tolerance.""" + if self.direction == "neutral": + return False + band = abs(baseline.value) * self.rel_tol + self.abs_tol + if self.direction == "lower_better": + return self.value > baseline.value + band + return self.value < baseline.value - band + + def to_json(self) -> dict: + return { + "value": self.value, + "unit": self.unit, + "direction": self.direction, + "rel_tol": self.rel_tol, + "abs_tol": self.abs_tol, + "notes": self.notes, + } + + @staticmethod + def from_json(key: str, d: dict) -> "Metric": + return Metric( + key=key, + value=d["value"], + unit=d.get("unit", ""), + direction=d.get("direction", "lower_better"), + rel_tol=d.get("rel_tol", 0.10), + abs_tol=d.get("abs_tol", 0.0), + notes=d.get("notes", ""), + ) + + +@dataclass +class WorkloadResult: + """Output of one workload: its metrics plus optional artifact paths + (plots, raw JSON) for the report.""" + + name: str + metrics: list[Metric] = field(default_factory=list) + artifacts: list[str] = field(default_factory=list) + skipped: str | None = None # reason, if the workload could not run diff --git a/benchmarks/suite/pyroki_ik.py b/benchmarks/suite/pyroki_ik.py new file mode 100644 index 0000000..ea47ce6 --- /dev/null +++ b/benchmarks/suite/pyroki_ik.py @@ -0,0 +1,208 @@ +"""A/B regression benchmark: pyroki IK-Beam on two jaxls versions. + +The pyroki side of /tmp/pyroki/benchmark/ik_benchmark.py (cuRobo comparison +and torch removed; targets generated by FK on uniformly sampled joint +configurations, same recipe the original uses for its q_sample). Select the +jaxls under test via PYTHONPATH: + + python /tmp/pyroki_ab.py --label branch + PYTHONPATH=/tmp/jaxls-main/src python /tmp/pyroki_ab.py --label main +""" + +import argparse +import json +import time +import xml.etree.ElementTree as ET +from io import StringIO + +import jax +import jaxlie +import numpy as np +import yourdfpy +from jax import lax +from jax import numpy as jnp +from robot_descriptions.loaders.yourdfpy import load_robot_description + +import jaxls +import pyroki as pk # type: ignore[import-not-found] + + +def newton_raphson(f, x, iters): + def update(x, _): + y = x - f(x) / jax.grad(f)(x) + return y, None + + x, _ = lax.scan(update, 1.0, length=iters) + return x + + +def roberts_sequence(num_points, dim, root): + basis = 1 - (1 / root ** (1 + jnp.arange(dim))) + n = jnp.arange(num_points) + x = n[:, None] * basis[None, :] + x, _ = jnp.modf(x) + return x + + +class PyrokiIkBeamHelper: + """Verbatim from pyroki/benchmark/ik_benchmark.py.""" + + def __init__(self): + urdf = load_robot_description("panda_description") + xml_tree = urdf.write_xml() + for joint in xml_tree.findall('.//joint[@type="prismatic"]'): + joint.set("type", "fixed") + for tag in ("axis", "limit", "dynamics"): + child = joint.find(tag) + if child is not None: + joint.remove(child) + xml_str = ET.tostring(xml_tree.getroot(), encoding="unicode") + urdf = yourdfpy.URDF.load(StringIO(xml_str)) + assert urdf.validate() + + robot = pk.Robot.from_urdf(urdf) + ee_link_name = "panda_hand_tcp" + self.target_link_index = jnp.array(robot.links.names.index(ee_link_name)) + self.robot = robot + exp = robot.joints.num_actuated_joints + self.root = newton_raphson(lambda x: x ** (exp + 1) - x - 1, 1.0, 10_000) + + def solve_ik(self, target_wxyz: jax.Array, target_position: jax.Array) -> jax.Array: + num_seeds_init: int = 64 + num_seeds_final: int = 4 + total_steps: int = 16 + init_steps: int = 6 + + def solve_one(initial_q, lambda_initial, max_iters): + joint_var = robot.joint_var_cls(0) + factors = [ + pk.costs.pose_cost_analytic_jac( + robot, + joint_var, + jaxlie.SE3.from_rotation_and_translation( + jaxlie.SO3(target_wxyz), target_position + ), + self.target_link_index, + pos_weight=10.0, + ori_weight=5.0, + ), + pk.costs.limit_cost(robot, joint_var, weight=50.0), + ] + sol, summary = ( + jaxls.LeastSquaresProblem(factors, [joint_var]) + .analyze() + .solve( + initial_vals=jaxls.VarValues.make( + [joint_var.with_value(initial_q)] + ), + verbose=False, + linear_solver="dense_cholesky", + termination=jaxls.TerminationConfig( + max_iterations=max_iters, + early_termination=False, + ), + trust_region=jaxls.TrustRegionConfig(lambda_initial=lambda_initial), + return_summary=True, + ) + ) + return sol[joint_var], summary + + vmapped_solve = jax.vmap(solve_one, in_axes=(0, 0, None)) + robot = self.robot + initial_qs = robot.joints.lower_limits + roberts_sequence( + num_seeds_init, robot.joints.num_actuated_joints, self.root + ) * (robot.joints.upper_limits - robot.joints.lower_limits) + + initial_sols, summary = vmapped_solve( + initial_qs, jnp.full(initial_qs.shape[:1], 10.0), init_steps + ) + best_initial_sols = jnp.argsort( + summary.cost_history[jnp.arange(num_seeds_init), -1] + )[:num_seeds_final] + best_sols, summary = vmapped_solve( + initial_sols[best_initial_sols], + summary.lambda_history[jnp.arange(num_seeds_init), -1][best_initial_sols], + total_steps - init_steps, + ) + return best_sols[ + jnp.argmin( + summary.cost_history[jnp.arange(num_seeds_final), summary.iterations] + ) + ] + + def forward_kinematics(self, q): + return self.robot.forward_kinematics(jnp.asarray(q))[self.target_link_index] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--label", default="suite") + ap.add_argument( + "--batch-sizes", type=int, nargs="+", default=[1, 10, 100, 1000, 2000] + ) + ap.add_argument("--out", default=None) + args = ap.parse_args() + + print(f"jaxls under test: {jaxls.__file__}") + + helper = PyrokiIkBeamHelper() + batched_ik = jax.jit(jax.vmap(helper.solve_ik)) + batched_fk = jax.jit(jax.vmap(helper.forward_kinematics)) + + rng = np.random.default_rng(2) # Same role as the benchmark's torch seed 2. + lo = np.asarray(helper.robot.joints.lower_limits) + hi = np.asarray(helper.robot.joints.upper_limits) + + results = {"label": args.label, "jaxls": jaxls.__file__, "rows": []} + for b in args.batch_sizes: + q_sample = rng.uniform(lo, hi, size=(b, lo.shape[0])) + target = np.asarray(batched_fk(q_sample)) + wxyz, pos = jnp.asarray(target[..., 0:4]), jnp.asarray(target[..., 4:7]) + + # Warmup (compile), then min of 3. + sols = jax.block_until_ready(batched_ik(wxyz, pos)) + times = [] + for _ in range(3): + t0 = time.perf_counter() + sols = jax.block_until_ready(batched_ik(wxyz, pos)) + times.append(time.perf_counter() - t0) + + # Errors of the solved poses vs targets. + reached = np.asarray(batched_fk(np.asarray(sols))) + pos_err = np.linalg.norm(reached[..., 4:7] - target[..., 4:7], axis=-1) + ori_err = np.asarray( + jnp.abs( + ( + jaxlie.SO3(jnp.asarray(reached[..., 0:4])).inverse() + @ jaxlie.SO3(jnp.asarray(target[..., 0:4])) + ).log() + ).max(axis=-1) + ) + succ = float(np.mean((pos_err < 1e-3) & (ori_err < 1e-3)) * 100) + row = { + "batch": b, + "time_ms": min(times) * 1e3, + "succ_pct": succ, + "pos_err_mm_mean": float(pos_err.mean() * 1e3), + "pos_err_mm_p99": float(np.quantile(pos_err, 0.99) * 1e3), + "ori_err_mean": float(ori_err.mean()), + } + results["rows"].append(row) + print( + f"batch={b:>5} {row['time_ms']:>9.2f} ms succ={succ:6.2f}% " + f"pos_err mean={row['pos_err_mm_mean']:.5f} mm p99={row['pos_err_mm_p99']:.5f} mm " + f"ori_err mean={row['ori_err_mean']:.2e}", + flush=True, + ) + + import json as _json + + print("__PYROKI__" + _json.dumps(results["rows"]) + "__END__", flush=True) + out = args.out or f"/tmp/pyroki_ab_{args.label}.json" + with open(out, "w") as f: + json.dump(results, f, indent=2) + print("wrote", out) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/suite/workloads.py b/benchmarks/suite/workloads.py new file mode 100644 index 0000000..99bf21a --- /dev/null +++ b/benchmarks/suite/workloads.py @@ -0,0 +1,350 @@ +"""Workloads: each runs a piece of jaxls and reduces it to `Metric`s. + +These are deliberately thin wrappers over the existing harness modules +(`bal`, `device_sweep`, `float32_check`) so there is one implementation of +the measurement kernels, not two. Every workload is a function +`(cfg) -> WorkloadResult` registered in `ALL`. + +Timing note: metrics use `bal.run_k_iterations` (warmup + min-of-repeats), +so they are execution times, not compile times. Per-solve budgets keep the +slow full-system baselines from dominating wall-clock; a budget-stopped row +still yields its last cost/time, flagged in the metric notes. +""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from pathlib import Path + +import jax + +# benchmarks/ on path for the shared harness modules. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import bal # noqa: E402 +import device_sweep as ds # noqa: E402 + +from .metrics import Metric, WorkloadResult # noqa: E402 + + +@dataclass(frozen=True) +class SuiteConfig: + quick: bool = False + """Quick tier: GPU only, drop the slow CPU full-system baselines, fewer + k. For the hill-climbing inner loop.""" + devices: tuple[str, ...] = ("gpu", "cpu") + repeats: int = 3 + ba_lambda_initial: float = 1e2 + """Uniform damping for the BA matrix (BA curvature scale; see + results.md). Applied to every method so the comparison is solver-only.""" + solve_budget_s: float = 25.0 + """Skip the rest of a method's k-ladder once one solve exceeds this.""" + + +def _device_list(cfg: SuiteConfig) -> list[str]: + want = ["gpu"] if cfg.quick else list(cfg.devices) + return [p for p in want if ds.device_for(p) is not None] + + +def _subprocess_env() -> dict: + """Env for subprocess workloads. The parent suite process holds GPU + memory (and leaves it fragmented after the BA warmups), so a child + targeting the same GPU OOMs — even with preallocation off — when the + child is memory-hungry (pyroki's batch-2000 vmap IK). Strategy: + + 1. If another GPU is free, pin the child to it (CUDA_VISIBLE_DEVICES). + 2. Otherwise share the parent's GPU but disable preallocation and CUDA + command buffers (graph capture needs a large contiguous block). + + Falls back gracefully to CPU-only env when no GPU / nvidia-smi. + """ + import os + import subprocess + + env = { + **os.environ, + "XLA_PYTHON_CLIENT_PREALLOCATE": "false", + "XLA_FLAGS": ( + os.environ.get("XLA_FLAGS", "") + " --xla_gpu_enable_command_buffer=" + ).strip(), + } + # Pick a GPU the parent isn't using, if visible to nvidia-smi. + try: + out = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,memory.used", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=10, + ).stdout + free = [ + line.split(",")[0].strip() + for line in out.strip().splitlines() + if int(line.split(",")[1]) < 500 # MiB used -> effectively idle + ] + # Don't reuse a GPU the parent already pinned itself to. + parent = os.environ.get("CUDA_VISIBLE_DEVICES", "") + free = [g for g in free if g != parent] + if free: + env["CUDA_VISIBLE_DEVICES"] = free[0] + except Exception: # noqa: BLE001 + pass + return env + + +def _run_marker_subprocess(argv: list[str], marker: str, timeout: float): + """Run argv as a subprocess (with _subprocess_env so it doesn't OOM as a + co-tenant of the CUDA-warmed parent) and parse JSON printed between + `marker` and `__END__`. Returns (parsed, None) on success or + (None, skip_reason) — surfacing the subprocess's real stderr instead of a + cryptic missing-marker error.""" + import json + import subprocess + + try: + p = subprocess.run( + argv, capture_output=True, text=True, timeout=timeout, env=_subprocess_env() + ) + out = p.stdout + if marker not in out: + tail = (p.stderr or "").strip().splitlines()[-3:] + return None, f"subprocess produced no result; stderr: {' | '.join(tail)}" + seg = out[out.index(marker) + len(marker) : out.index("__END__")] + return json.loads(seg), None + except Exception as e: # noqa: BLE001 + return None, f"{type(e).__name__}: {e}" + + +# -------------------------------------------------------------------------- +# Bundle adjustment matrix +# -------------------------------------------------------------------------- + +_BA_METHODS = ( + # (label, linear_solver, elimination) + ("full_cg", "conjugate_gradient", False), + ("schur_dense", "dense_cholesky", True), + ("schur_cg", "conjugate_gradient", True), +) + + +def bundle_adjustment(cfg: SuiteConfig) -> WorkloadResult: + res = WorkloadResult(name="bundle_adjustment") + problems = ( + ["toy", "ladybug49"] if cfg.quick else ["toy", "ladybug49", "trafalgar138"] + ) + for prob in problems: + spec = ds.PROBLEMS[prob] + ks = spec.ks[: 5 if cfg.quick else len(spec.ks)] + elim, init = spec.load("auto") + full, _ = spec.load("off") + for platform in _device_list(cfg): + dev = ds.device_for(platform) + elim_d = jax.device_put(elim, dev) + full_d = jax.device_put(full, dev) + init_d = jax.device_put(init, dev) + for label, solver, elimination in _BA_METHODS: + target = elim_d if elimination else full_d + last_cost = last_t = float("nan") + stopped = None + for k in ks: + cost, t = bal.run_k_iterations( + target, + init_d, + k, + linear_solver=solver, + repeats=1 if cfg.quick else cfg.repeats, + lambda_initial=cfg.ba_lambda_initial, + warmup_budget_s=cfg.solve_budget_s, + ) + last_cost, last_t = cost, t + if t > cfg.solve_budget_s: + stopped = k + break + base = f"{prob}.{platform}.{label}" + note = f"budget-stopped at k={stopped}" if stopped else "" + res.metrics.append( + Metric( + f"{base}.cost", + last_cost, + "cost", + "lower_better", + rel_tol=0.02, + notes=note, + ) + ) + res.metrics.append( + Metric( + f"{base}.time", + last_t, + "s", + "lower_better", + rel_tol=0.30, + notes=note, + ) + ) + return res + + +# -------------------------------------------------------------------------- +# float32 robustness (Ladybug, direct Schur path) +# -------------------------------------------------------------------------- + + +def float32_robustness(cfg: SuiteConfig) -> WorkloadResult: + res = WorkloadResult(name="float32_robustness") + # float32_check flips the x64 flag at import time, so run it in a + # subprocess to avoid contaminating the rest of the suite. + # Run float32_check.run(x64=False) in a subprocess: it flips the global + # x64 flag at import, so it must not share this process. Reuses the + # importable module rather than re-spelling the solve recipe. + root = str(Path(__file__).resolve().parent.parent) + code = ( + f"import sys; sys.path.insert(0, {root!r})\n" + "import numpy as onp, json, float32_check\n" + "ch = float32_check.run(x64=False)\n" + "print('__F32__' + json.dumps({'nans': int(onp.isnan(ch).sum()),\n" + " 'best': float(onp.nanmin(ch))}) + '__END__')\n" + ) + d, skip = _run_marker_subprocess( + [sys.executable, "-c", code], "__F32__", timeout=300 + ) + if skip is not None: + res.skipped = skip + return res + assert d is not None # skip is None => parsed result present + res.metrics.append( + Metric( + "float32.ladybug49.nan_count", + float(d["nans"]), + "", + "lower_better", + rel_tol=0.0, + abs_tol=0.0, + ) + ) + res.metrics.append( + Metric( + "float32.ladybug49.best_cost", + float(d["best"]), + "cost", + "lower_better", + rel_tol=0.05, + notes="float32 direct Schur path", + ) + ) + return res + + +# -------------------------------------------------------------------------- +# pyroki IK-Beam downstream (optional; needs pyroki + the A/B script) +# -------------------------------------------------------------------------- + + +def pyroki_ik(cfg: SuiteConfig) -> WorkloadResult: + res = WorkloadResult(name="pyroki_ik") + try: + import pyroki # type: ignore[import-not-found] # noqa: F401 + except ImportError: + res.skipped = "pyroki not installed (pip install the pyroki repo)" + return res + script = Path(__file__).resolve().parent / "pyroki_ik.py" + if not script.exists(): + res.skipped = "benchmarks/suite/pyroki_ik.py missing" + return res + batches = ["100", "1000"] if cfg.quick else ["100", "1000", "2000"] + rows, skip = _run_marker_subprocess( + [sys.executable, str(script), "--batch-sizes", *batches], + "__PYROKI__", + timeout=600, + ) + if skip is not None: + res.skipped = skip + return res + assert rows is not None # skip is None => parsed result present + for r in rows: + b = r["batch"] + res.metrics.append( + Metric( + f"pyroki.b{b}.success_pct", + r["succ_pct"], + "%", + "higher_better", + rel_tol=0.02, + ) + ) + res.metrics.append( + Metric( + f"pyroki.b{b}.p99_pos_err_mm", + r["pos_err_mm_p99"], + "mm", + "lower_better", + rel_tol=0.25, + ) + ) + res.metrics.append( + Metric( + f"pyroki.b{b}.time_ms", r["time_ms"], "ms", "lower_better", rel_tol=0.30 + ) + ) + return res + + +# -------------------------------------------------------------------------- +# Example notebooks (convergence to a relative cost floor) +# -------------------------------------------------------------------------- + + +def example_notebooks(cfg: SuiteConfig) -> WorkloadResult: + """Reduce each notebook's headline solve to (final cost, steps, time) + via the trace harness. Skipped in --quick (notebook execution is slow).""" + res = WorkloadResult(name="example_notebooks") + if cfg.quick: + res.skipped = "skipped in --quick (notebook execution is slow)" + return res + import trace_examples as tx + + traces = tx.trace_single(src="src", timeout=300.0) + for name, trace in traces.items(): + if trace.get("iterations", 0) <= 0: + continue + res.metrics.append( + Metric( + f"example.{name}.final_cost", + float(min(trace["cost_history"])), + "cost", + "lower_better", + rel_tol=0.05, + ) + ) + res.metrics.append( + Metric( + f"example.{name}.steps", + float(trace["iterations"]), + "steps", + "lower_better", + rel_tol=0.20, + ) + ) + return res + + +ALL = { + "bundle_adjustment": bundle_adjustment, + "float32_robustness": float32_robustness, + "pyroki_ik": pyroki_ik, + "example_notebooks": example_notebooks, +} + +# Default set for a no-argument run / the regression gate: the fast, +# in-process-reliable workloads. Two workloads are opt-in (run via +# `--only`), kept out of the gate so it stays deterministic: +# - example_notebooks: 19 jupyter subprocesses, ~minutes. +# - pyroki_ik: needs pyroki installed, and its memory-hungry batched IK is +# unreliable as a *co-tenant* subprocess of a CUDA-warmed parent (cuSolver +# init / command-buffer OOM); it runs fine standalone (`--only pyroki_ik`) +# and is the canonical downstream-regression check. +DEFAULT = ("bundle_adjustment", "float32_robustness") diff --git a/benchmarks/trace_examples.py b/benchmarks/trace_examples.py new file mode 100644 index 0000000..ade0581 --- /dev/null +++ b/benchmarks/trace_examples.py @@ -0,0 +1,287 @@ +"""Per-iteration cost-vs-time traces for the example notebooks. + +Uses the in-solve timestamps recorded via `jaxls.record_iteration_times()` +(host callback per outer LM step), so a single solve yields a full +convergence trace — no matched-k re-solving. Every top-level solve in a +notebook is instrumented and the one with the largest relative cost drop is +kept (the headline optimization, not a warmup/forward/demo solve); the +captured (cost_history, elapsed_s, iterations) is dumped as JSON, then +plotted as a grid of cost-vs-time subplots (one per notebook), running-best +cost with a marker per LM step from step 0. + + uv run --extra dev --extra docs python benchmarks/trace_examples.py # run + plot + uv run --extra dev --extra docs python benchmarks/trace_examples.py --plot # plot from saved JSON +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import tempfile +from pathlib import Path + +EXAMPLES_DIR = Path(__file__).parent.parent / "docs/source/examples" +RESULTS = Path(__file__).parent / "results" +TRACE_JSON = RESULTS / "example_traces.json" + +# Instrumentation: capture EVERY top-level solve, then (at finalize) emit +# the one with the largest cost reduction — the notebook's headline +# optimization, not an incidental warmup/forward/demo solve. +# +# Per-step timestamps come from `jaxls.record_iteration_times()` (host +# float64 list, off the jitted path). The `main` config predates that API +# and instead exposes `SolveSummary.time_history`; we feature-detect and use +# whichever is present so both sides yield real per-step times. A warmup +# solve absorbs compilation before the timed one. +INSTRUMENT = """ +import jax, jaxls, jaxls._problem, json as _json +import numpy as _onp +_traces = [] +_orig = jaxls._problem.AnalyzedLeastSquaresProblem.solve +_HAS_RECORDER = hasattr(jaxls, "record_iteration_times") + +def _traced(self, *a, **k): + import jax.core + def _tr(x, s=None): + s = set() if s is None else s + if id(x) in s: return False + s.add(id(x)) + if isinstance(x, jax.core.Tracer): return True + return any(_tr(v, s) for v in getattr(x, "__dict__", {}).values()) + if _tr(self) or _tr(a) or _tr(k): + return _orig(self, *a, **k) + k = dict(k); want = k.get("return_summary", False); k["return_summary"] = True + if _HAS_RECORDER: + with jaxls.record_iteration_times() as _times: + jax.block_until_ready(_orig(self, *a, **k)) # warmup / compile + _times.clear() + sol, summ = _orig(self, *a, **k) + jax.block_until_ready((sol, summ)) + elapsed = [t - _times[0] for t in _times] + else: + # Legacy main: timestamps always recorded in SolveSummary.time_history. + jax.block_until_ready(_orig(self, *a, **k)) # warmup / compile + sol, summ = _orig(self, *a, **k) + jax.block_until_ready((sol, summ)) + th = _onp.asarray(summ.time_history) + elapsed = (th - th[0]).tolist() + n = int(summ.iterations) + # cost_history holds init at [0] + one slot per step but is only + # max_iterations long, so it caps at init + (max_iterations-1) steps; the + # timestamp recorder has no such cap. Align both to the shorter length. + ch = _onp.asarray(summ.cost_history)[: n + 1].tolist() + m = min(len(ch), len(elapsed)) + _traces.append({ + "iterations": n, + "cost_history": ch[:m], + "elapsed_s": elapsed[:m], + }) + return (sol, summ) if want else sol + +jaxls._problem.AnalyzedLeastSquaresProblem.solve = _traced +""" + +FINALIZE = """ +import json as _json, math as _math +def _drop(r): + c = r["cost_history"] + if len(c) < 2 or c[0] <= 0 or not _math.isfinite(c[0]): + return -1.0 + return (c[0] - min(c)) / c[0] # relative reduction +_best = max(_traces, key=_drop, default=None) if _traces else None +if _best is not None and _drop(_best) > 0: + print("__TRACE__" + _json.dumps(_best) + "__END__", flush=True) +""" + + +# jaxls under test per label; PYTHONPATH for the notebook subprocess. +CONFIGS = {"main": "/tmp/jaxls-main/src", "PR": "src"} + + +def _list_notebooks() -> list[Path]: + """Source notebooks under the examples tree. Excludes checkpoints and the + `tmp*.ipynb` instrumented copies this script writes (a crash/kill can skip + their cleanup, and we must not re-trace those stragglers next run).""" + return sorted( + p + for p in EXAMPLES_DIR.rglob("*.ipynb") + if ".ipynb_checkpoints" not in str(p) and not p.name.startswith("tmp") + ) + + +def run_one(nb_path: Path, src: str, timeout: float) -> dict | None: + import nbformat + + nb = nbformat.read(open(nb_path), as_version=4) + nb.cells.insert(0, nbformat.v4.new_code_cell(source=INSTRUMENT)) + nb.cells.append(nbformat.v4.new_code_cell(source=FINALIZE)) + with tempfile.NamedTemporaryFile( + mode="w", suffix=".ipynb", delete=False, dir=nb_path.parent + ) as f: + nbformat.write(nb, f) + tmp = f.name + try: + p = subprocess.run( + [sys.executable, "-m", "jupyter", "execute", tmp, "--inplace"], + capture_output=True, + text=True, + timeout=timeout, + cwd=str(EXAMPLES_DIR.parent.parent.parent), + # PREALLOCATE=false: when run under the benchmark suite, the parent + # process already holds most of the GPU; without this the notebook + # subprocess OOMs. + env={ + **__import__("os").environ, + "PYTHONPATH": src, + "XLA_PYTHON_CLIENT_PREALLOCATE": "false", + }, + ) + ex = nbformat.read(open(tmp), as_version=4) + for cell in ex.cells: + for out in cell.get("outputs", []): + txt = out.get("text", "") + if "__TRACE__" in txt: + s = txt.index("__TRACE__") + len("__TRACE__") + e = txt.index("__END__", s) + return json.loads(txt[s:e]) + if p.returncode != 0: + print(f" {nb_path.stem}: exit {p.returncode}") + return None + except subprocess.TimeoutExpired: + print(f" {nb_path.stem}: timeout") + return None + finally: + Path(tmp).unlink(missing_ok=True) + + +def trace_single(src: str = "src", timeout: float = 300.0) -> dict: + """Trace every notebook under one jaxls (PYTHONPATH=`src`). Returns + {name: trace_rec}. Used by the benchmark suite for single-config metrics + (the A/B `run_all` traces two configs).""" + out: dict = {} + nbs = _list_notebooks() + for i, nb in enumerate(nbs): + print(f"[{i + 1}/{len(nbs)}] {nb.stem}", flush=True) + rec = run_one(nb, src, timeout) + if rec and rec["iterations"] > 0: + out[nb.stem] = rec + return out + + +def run_all(timeout: float) -> dict: + traces: dict = {} + nbs = _list_notebooks() + for label, src in CONFIGS.items(): + print(f"\n=== config: {label} ({src}) ===", flush=True) + for i, nb in enumerate(nbs): + name = nb.stem + print(f"[{label} {i + 1}/{len(nbs)}] {name}", flush=True) + rec = run_one(nb, src, timeout) + if rec and rec["iterations"] > 0: + traces.setdefault(name, {})[label] = rec + print( + f" {rec['iterations']} steps, " + f"cost {rec['cost_history'][0]:.4g} -> {rec['cost_history'][-1]:.4g}, " + f"{rec['elapsed_s'][-1]:.3f}s" + ) + TRACE_JSON.write_text(json.dumps(traces, indent=2)) + print(f"wrote {TRACE_JSON}") + return traces + + +def _running_best(costs: list[float]) -> list[float]: + # cost_history records each step's proposed cost, so a rejected step can + # bump it up; cummin shows the best solution found so far. + best = list(costs) + for i in range(1, len(best)): + best[i] = min(best[i], best[i - 1]) + return best + + +def _speedup_label(traces_for_name: dict) -> str: + """PR-vs-main speedup: ratio of wall-clock to reach main's final cost. + Empty string unless both configs are present and the ratio is meaningful.""" + main, pr = traces_for_name.get("main"), traces_for_name.get("PR") + if main is None or pr is None: + return "" + target = min(main["cost_history"]) * 1.003 + t_main = main["elapsed_s"][-1] + + def time_to(rec: dict) -> float | None: + for t, b in zip(rec["elapsed_s"], _running_best(rec["cost_history"])): + if b <= target: + return t + return None + + t_pr = time_to(pr) + if t_pr is None or t_pr <= 0 or t_main <= 0: + return "" + factor = t_main / t_pr + return f" ({factor:.1f}× faster)" if factor >= 1.2 else "" + + +def plot(traces: dict) -> None: + import math + + import matplotlib.pyplot as plt + + style = {"main": ("#d62728", "o", "--"), "PR": ("#1f77b4", "^", "-")} + names = sorted(traces) + ncol = 4 + nrow = math.ceil(len(names) / ncol) + fig, axes = plt.subplots( + nrow, ncol, figsize=(4.5 * ncol, 3.4 * nrow), squeeze=False + ) + for ax in axes.flat: + ax.set_visible(False) + for ax, name in zip(axes.flat, names): + ax.set_visible(True) + for label, (color, marker, ls) in style.items(): + rec = traces[name].get(label) + if rec is None: + continue + best = _running_best(rec["cost_history"]) + ax.plot( + rec["elapsed_s"], + best, + marker=marker, + ms=5, + ls=ls, + color=color, + label=f"{label} ({rec['iterations']} steps)", + ) + ax.set_yscale("log") + ax.set_xscale("symlog", linthresh=1e-2) + ax.set_xlim(left=0) + ax.set_title(name + _speedup_label(traces[name]), fontsize=11) + ax.set_xlabel("wall-clock (s)", fontsize=10) + ax.set_ylabel("cost", fontsize=10) + ax.tick_params(labelsize=9) + ax.grid(True, which="major", ls="-", lw=0.4, alpha=0.3) + ax.legend(fontsize=9) + fig.suptitle( + "jaxls examples: cost vs wall-clock per LM step — main (dashed) vs PR " + "(solid)\nreal in-solve timestamps on both; speedup = time to reach " + "main's final cost", + fontsize=14, + ) + fig.tight_layout() + out = RESULTS / "example_traces.png" + fig.savefig(out, dpi=150) + print(f"wrote {out}") + plt.close(fig) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--plot", action="store_true") + ap.add_argument("--timeout", type=float, default=300.0) + args = ap.parse_args() + traces = json.loads(TRACE_JSON.read_text()) if args.plot else run_all(args.timeout) + plot(traces) + + +if __name__ == "__main__": + main() diff --git a/src/jaxls/__init__.py b/src/jaxls/__init__.py index 70625c3..29e1936 100644 --- a/src/jaxls/__init__.py +++ b/src/jaxls/__init__.py @@ -18,6 +18,7 @@ from ._problem import AnalyzedLeastSquaresProblem as AnalyzedLeastSquaresProblem from ._problem import LeastSquaresProblem as LeastSquaresProblem from ._solvers import ConjugateGradientConfig as ConjugateGradientConfig + from ._solvers import record_iteration_times as record_iteration_times from ._solvers import SolveSummary as SolveSummary from ._solvers import TerminationConfig as TerminationConfig from ._solvers import TrustRegionConfig as TrustRegionConfig @@ -42,6 +43,7 @@ ) from ._py310._problem import LeastSquaresProblem as LeastSquaresProblem from ._py310._solvers import ConjugateGradientConfig as ConjugateGradientConfig + from ._py310._solvers import record_iteration_times as record_iteration_times from ._py310._solvers import SolveSummary as SolveSummary from ._py310._solvers import TerminationConfig as TerminationConfig from ._py310._solvers import TrustRegionConfig as TrustRegionConfig diff --git a/src/jaxls/_preconditioning.py b/src/jaxls/_preconditioning.py index 05ef064..40b4aa3 100644 --- a/src/jaxls/_preconditioning.py +++ b/src/jaxls/_preconditioning.py @@ -5,6 +5,8 @@ import jax from jax import numpy as jnp +from .utils import _batched_gram + if TYPE_CHECKING: from ._problem import AnalyzedLeastSquaresProblem from ._sparse_matrices import BlockRowSparseMatrix @@ -80,8 +82,11 @@ def make_block_jacobi_precoditioner( ) ) - # f: factor, r: residual, v: variable, t/a: tangent - gram_blocks = jnp.einsum("frvt,frva->fvta", A_blocks, A_blocks) + # f: factor, r: residual, v: variable, t/a: tangent. The + # broadcast helper avoids einsum's slow batched-GEMM lowering + # for the tiny residual contraction; see `utils._batched_gram`. + A_fvrt = jnp.swapaxes(A_blocks, 1, 2) + gram_blocks = _batched_gram(A_fvrt, A_fvrt) assert gram_blocks.shape == ( num_costs, num_vars, diff --git a/src/jaxls/_problem.py b/src/jaxls/_problem.py index 7f12e08..c250293 100644 --- a/src/jaxls/_problem.py +++ b/src/jaxls/_problem.py @@ -21,6 +21,7 @@ from ._analyzed_cost import _AnalyzedCost, _augment_constraint_cost from ._cost import Cost, CostKind, CustomJacobianCache +from ._schur import EliminationPlan from ._solvers import ( ConjugateGradientConfig, NonlinearSolver, @@ -143,7 +144,11 @@ def show( max_variables=max_variables, ) - def analyze(self, use_onp: bool = False) -> AnalyzedLeastSquaresProblem: + def analyze( + self, + use_onp: bool = False, + schur_elimination: Literal["auto", "off"] | tuple[type[Var], ...] = "auto", + ) -> AnalyzedLeastSquaresProblem: """Analyze sparsity pattern of least squares problem. Needed before solving. Processes all costs and variables to compute the sparse Jacobian structure, @@ -152,6 +157,31 @@ def analyze(self, use_onp: bool = False) -> AnalyzedLeastSquaresProblem: Args: use_onp: If True, use numpy instead of jax.numpy for index computations. Can be faster for problem setup on CPU. + schur_elimination: Controls Schur-complement variable elimination. + When a dominant block-diagonal variable type is eliminated (for + example, landmarks in bundle adjustment), solves run on the much + smaller, better-conditioned reduced system and then + back-substitute the eliminated variables. All three linear + solvers use the plan: "dense_cholesky" and "conjugate_gradient" + solve the reduced system densely / matrix-free, and "cholmod" + factors it sparse-directly. + + - ``"auto"`` (default): automatically eliminate a dominant + block-diagonal variable type if one exists, otherwise solve + the full system. + - ``"off"``: skip elimination and solve the full system — + useful for debugging or benchmarking against the + non-eliminated solve. + - a tuple of variable types, e.g. ``(LandmarkVar,)``: eliminate + exactly these types. Each must be block-diagonal (no single + cost may touch more than one variable of the type), or + building the plan raises a ``ValueError``. + + Only a single level of elimination is currently supported: the + eliminated types are removed in one Schur step and the remaining + types form the reduced system. (Nested / multi-level elimination + — eliminating further types from the already-reduced system — is + not implemented.) Returns: An AnalyzedLeastSquaresProblem ready for solving. @@ -372,7 +402,7 @@ def _sort_key(x: Any) -> str: shape=(residual_dim_sum, tangent_dim_sum), ) - return AnalyzedLeastSquaresProblem( + analyzed = AnalyzedLeastSquaresProblem( _stacked_costs=tuple(stacked_costs), _cost_counts=tuple(cost_counts), _sorted_ids_from_var_type=sorted_ids_from_var_type, @@ -384,6 +414,51 @@ def _sort_key(x: Any) -> str: _residual_dim=residual_dim_sum, ) + # Precompute the Schur-complement elimination plan when a dominant + # block-diagonal variable type exists (for example, landmarks in + # bundle adjustment). This is host-side index structure that depends + # only on the problem's sparsity, so it belongs with the rest of the + # analysis; `solve()` decides whether to use it. + from ._schur import ( + _TracedVariableIdsError, + build_elimination_plan, + infer_eliminate, + ) + + if schur_elimination == "auto": + eliminate = infer_eliminate(analyzed) + elif schur_elimination == "off": + eliminate = () + elif isinstance(schur_elimination, tuple): + eliminate = schur_elimination + else: + raise ValueError( + "schur_elimination must be 'auto', 'off', or a tuple of " + f"variable types; got {schur_elimination!r}." + ) + if len(eliminate) > 0: + try: + elimination = build_elimination_plan(analyzed, eliminate) + except _TracedVariableIdsError: + # Variable IDs are tracers (analyze() itself is being + # traced); solves will run without elimination. + logger.info( + "Variable elimination: variable IDs are traced; solves " + "will not use elimination" + ) + else: + logger.info( + "Variable elimination: eliminating {} ({} of {} tangent " + "dims); reduced system is {}-dimensional", + ", ".join(var_type.__name__ for var_type in eliminate), + analyzed._tangent_dim - elimination.reduced_dim, + analyzed._tangent_dim, + elimination.reduced_dim, + ) + with jdc.copy_and_mutate(analyzed, validate=False) as analyzed: + analyzed._elimination = elimination + return analyzed + @jdc.pytree_dataclass class AnalyzedLeastSquaresProblem: @@ -398,6 +473,11 @@ class AnalyzedLeastSquaresProblem: _tangent_start_from_var_type: jdc.Static[dict[type[Var[Any]], int]] _tangent_dim: jdc.Static[int] _residual_dim: jdc.Static[int] + _elimination: EliminationPlan | None = None + """Schur-complement elimination plan, precomputed by `analyze()` when a + dominant block-diagonal variable type is eliminated. Used by `solve()` for + all reduced-solve paths (dense Cholesky, CG, and CHOLMOD); None when + `analyze(schur_elimination="off")` or no type was eliminated.""" @overload def solve( @@ -450,7 +530,13 @@ def solve( Args: initial_vals: Initial values for the variables. If None, default values will be used. - linear_solver: The linear solver to use. + linear_solver: The linear solver to use. When a dominant + block-diagonal variable type exists (for example, landmarks + in bundle adjustment), "conjugate_gradient" and + "dense_cholesky" automatically eliminate it via the Schur + complement and solve the much smaller, better-conditioned + reduced system instead; the decision is logged. "cholmod" + always factors the full system. trust_region: Configuration for Levenberg-Marquardt trust region. termination: Configuration for termination criteria. sparse_mode: The representation to use for sparse matrix @@ -488,6 +574,15 @@ def solve( conjugate_gradient_config = linear_solver linear_solver = "conjugate_gradient" + # Schur-complement variable elimination: the plan was precomputed by + # `analyze()` when a dominant block-diagonal variable type exists + # (for example, landmarks in bundle adjustment); the linear solves + # then run on the much smaller, better-conditioned reduced system and + # back-substitute the eliminated variables. All three reduced-solve + # paths are supported: dense Cholesky, matrix-free CG, and CHOLMOD + # (sparse-direct on the reduced system, the Ceres/g2o combination). + elimination = self._elimination + # Create unified solver (handles both constrained and unconstrained). solver = NonlinearSolver( linear_solver, @@ -497,6 +592,7 @@ def solve( sparse_mode, verbose, augmented_lagrangian if has_constraints else None, + elimination, ) return solver.solve( problem=self, initial_vals=initial_vals, return_summary=return_summary diff --git a/src/jaxls/_py310/_analyzed_cost.py b/src/jaxls/_py310/_analyzed_cost.py index f87be7c..11d3c74 100644 --- a/src/jaxls/_py310/_analyzed_cost.py +++ b/src/jaxls/_py310/_analyzed_cost.py @@ -55,6 +55,7 @@ def compute_residual_flat(self, vals: Any, *args: Any) -> Any: def _make( cost: Any, ) -> Any: + if cost.kind != "l2_squared": return _augment_constraint_cost(cost) diff --git a/src/jaxls/_py310/_augmented_lagrangian.py b/src/jaxls/_py310/_augmented_lagrangian.py index 5df2c67..8b29b9f 100644 --- a/src/jaxls/_py310/_augmented_lagrangian.py +++ b/src/jaxls/_py310/_augmented_lagrangian.py @@ -82,6 +82,7 @@ def initialize_al_state( config: Any, verbose: Any = False, ) -> Any: + is_inequality = tuple( cost.kind in ("constraint_leq_zero", "constraint_geq_zero") for cost in problem._stacked_costs @@ -217,6 +218,7 @@ def update_problem_al_params( problem: Any, al_state: Any, ) -> Any: + from ._analyzed_cost import AugmentedLagrangianParams updated_costs = [] @@ -253,6 +255,7 @@ def check_al_convergence( al_state: Any, config: Any, ) -> Any: + converged_absolute = ( jnp.maximum(al_state.snorm, al_state.constraint_violation) < config.tolerance_absolute diff --git a/src/jaxls/_py310/_cost.py b/src/jaxls/_py310/_cost.py index 4429d25..9f8a97a 100644 --- a/src/jaxls/_py310/_cost.py +++ b/src/jaxls/_py310/_cost.py @@ -97,6 +97,7 @@ def factory( jac_custom_with_cache_fn: Any = None, name: Any = None, ) -> Any: + def decorator( compute_residual: Any, ) -> Any: diff --git a/src/jaxls/_py310/_preconditioning.py b/src/jaxls/_py310/_preconditioning.py index a181180..87a622b 100644 --- a/src/jaxls/_py310/_preconditioning.py +++ b/src/jaxls/_py310/_preconditioning.py @@ -4,6 +4,8 @@ from jax import numpy as jnp +from .utils import _batched_gram + def make_point_jacobi_precoditioner( A_blocksparse: Any, @@ -32,6 +34,7 @@ def make_point_jacobi_precoditioner( def make_block_jacobi_precoditioner(graph: Any, A_blocksparse: Any) -> Any: + gram_diagonal_blocks = list() for var_type, ids in graph._tangent_ordering.ordered_dict_items( graph._sorted_ids_from_var_type @@ -66,7 +69,8 @@ def make_block_jacobi_precoditioner(graph: Any, A_blocksparse: Any) -> Any: ) ) - gram_blocks = jnp.einsum("frvt,frva->fvta", A_blocks, A_blocks) + A_fvrt = jnp.swapaxes(A_blocks, 1, 2) + gram_blocks = _batched_gram(A_fvrt, A_fvrt) assert gram_blocks.shape == ( num_costs, num_vars, diff --git a/src/jaxls/_py310/_problem.py b/src/jaxls/_py310/_problem.py index 19619dd..466cbf0 100644 --- a/src/jaxls/_py310/_problem.py +++ b/src/jaxls/_py310/_problem.py @@ -80,7 +80,12 @@ def show( max_variables=max_variables, ) - def analyze(self, use_onp: Any = False) -> Any: + def analyze( + self, + use_onp: Any = False, + schur_elimination: Any = "auto", + ) -> Any: + if use_onp: jnp = onp else: @@ -271,7 +276,7 @@ def _sort_key(x: Any) -> Any: shape=(residual_dim_sum, tangent_dim_sum), ) - return AnalyzedLeastSquaresProblem( + analyzed = AnalyzedLeastSquaresProblem( _stacked_costs=tuple(stacked_costs), _cost_counts=tuple(cost_counts), _sorted_ids_from_var_type=sorted_ids_from_var_type, @@ -283,6 +288,44 @@ def _sort_key(x: Any) -> Any: _residual_dim=residual_dim_sum, ) + from ._schur import ( + _TracedVariableIdsError, + build_elimination_plan, + infer_eliminate, + ) + + if schur_elimination == "auto": + eliminate = infer_eliminate(analyzed) + elif schur_elimination == "off": + eliminate = () + elif isinstance(schur_elimination, tuple): + eliminate = schur_elimination + else: + raise ValueError( + "schur_elimination must be 'auto', 'off', or a tuple of " + f"variable types; got {schur_elimination!r}." + ) + if len(eliminate) > 0: + try: + elimination = build_elimination_plan(analyzed, eliminate) + except _TracedVariableIdsError: + logger.info( + "Variable elimination: variable IDs are traced; solves " + "will not use elimination" + ) + else: + logger.info( + "Variable elimination: eliminating {} ({} of {} tangent " + "dims); reduced system is {}-dimensional", + ", ".join(var_type.__name__ for var_type in eliminate), + analyzed._tangent_dim - elimination.reduced_dim, + analyzed._tangent_dim, + elimination.reduced_dim, + ) + with jdc.copy_and_mutate(analyzed, validate=False) as analyzed: + analyzed._elimination = elimination + return analyzed + @jdc.pytree_dataclass class AnalyzedLeastSquaresProblem: @@ -295,6 +338,7 @@ class AnalyzedLeastSquaresProblem: _tangent_start_from_var_type: jdc.Static[Any] _tangent_dim: jdc.Static[Any] _residual_dim: jdc.Static[Any] + _elimination: Any = None def solve( self, @@ -326,6 +370,8 @@ def solve( conjugate_gradient_config = linear_solver linear_solver = "conjugate_gradient" + elimination = self._elimination + solver = NonlinearSolver( linear_solver, trust_region, @@ -334,6 +380,7 @@ def solve( sparse_mode, verbose, augmented_lagrangian if has_constraints else None, + elimination, ) return solver.solve( problem=self, initial_vals=initial_vals, return_summary=return_summary diff --git a/src/jaxls/_py310/_schur.py b/src/jaxls/_py310/_schur.py new file mode 100644 index 0000000..6b9d77e --- /dev/null +++ b/src/jaxls/_py310/_schur.py @@ -0,0 +1,903 @@ +from __future__ import annotations + +import dataclasses +from typing import Any +from typing_extensions import assert_never + +import jax +import jax_dataclasses as jdc +import numpy as onp +from jax import numpy as jnp + +from .utils import _batched_gram, _batched_matmul, _batched_outer_last + + +class _TracedVariableIdsError(ValueError): + pass + + +@dataclasses.dataclass(frozen=True) +class _TypeInfo: + count: Any + dim: Any + orig_start: Any + start: Any + + +@jdc.pytree_dataclass +class _Slot: + type_index: jdc.Static[Any] + is_kept: jdc.Static[Any] + col_start: jdc.Static[Any] + index: Any + + +@jdc.pytree_dataclass +class _Combo: + kept_type_index: jdc.Static[Any] + elim_type_index: jdc.Static[Any] + entries: jdc.Static[Any] + kept_index: Any + elim_index: Any + + +@jdc.pytree_dataclass +class _CrossPairs: + combo_a: jdc.Static[Any] + combo_b: jdc.Static[Any] + obs_a: Any + obs_b: Any + block_id: Any + block_rows: Any + block_cols: Any + num_blocks: jdc.Static[Any] + + +@jdc.pytree_dataclass +class _SparseSPattern: + rows: Any + cols: Any + + +@jdc.pytree_dataclass +class EliminationPlan: + kept_types: jdc.Static[Any] + elim_types: jdc.Static[Any] + group_slots: Any + combos: Any + pairs: Any + reduced_dim: jdc.Static[Any] + tangent_dim: jdc.Static[Any] + sparse_s_pattern: Any = None + + +def infer_eliminate( + problem: Any, +) -> Any: + + slots_per_group = [ + { + var_type: ids.shape[-1] + for var_type, ids in cost.sorted_ids_from_var_type.items() + } + for cost in problem._stacked_costs + ] + + def total_elim_slots_ok(candidate: Any) -> Any: + return all( + sum(slots.get(var_type, 0) for var_type in candidate) <= 1 + for slots in slots_per_group + ) + + sizes = { + var_type: (int(ids.shape[0]) * var_type.tangent_dim, int(ids.shape[0])) + for var_type, ids in problem._sorted_ids_from_var_type.items() + } + + candidates = sorted(sizes, key=lambda t: sizes[t], reverse=True) + + chosen = list() + for var_type in candidates: + if len(chosen) + 1 == len(candidates): + break + if total_elim_slots_ok(chosen + [var_type]): + chosen.append(var_type) + + eliminated_dim = sum(sizes[t][0] for t in chosen) + if eliminated_dim * 2 < problem._tangent_dim: + return () + return tuple(chosen) + + +def build_elimination_plan( + problem: Any, + eliminate: Any, +) -> Any: + + elim_set = list() + for var_type in eliminate: + if var_type not in elim_set: + elim_set.append(var_type) + + for var_type in elim_set: + if var_type not in problem._sorted_ids_from_var_type: + raise ValueError( + f"eliminate={var_type.__name__} was requested, but no variables " + "of this type exist in the problem." + ) + + def _concrete(array: Any) -> Any: + try: + return onp.asarray(array) + except Exception as e: + raise _TracedVariableIdsError( + "Variable elimination requires concrete variable IDs; the " + "problem appears to be traced inside jax.jit." + ) from e + + kept_types = list() + elim_types = list() + kept_pos = dict() + elim_pos = dict() + reduced_dim = 0 + for var_type, ids in problem._tangent_ordering.ordered_dict_items( + problem._sorted_ids_from_var_type + ): + (count,) = ids.shape + if var_type in elim_set: + elim_pos[var_type] = len(elim_types) + elim_types.append( + _TypeInfo( + count=count, + dim=var_type.tangent_dim, + orig_start=problem._tangent_start_from_var_type[var_type], + start=0, + ) + ) + else: + kept_pos[var_type] = len(kept_types) + kept_types.append( + _TypeInfo( + count=count, + dim=var_type.tangent_dim, + orig_start=problem._tangent_start_from_var_type[var_type], + start=reduced_dim, + ) + ) + reduced_dim += count * var_type.tangent_dim + + if len(kept_types) == 0: + raise ValueError( + "All variable types were marked for elimination; at least one " + "type must be kept to form the reduced system." + ) + + group_slots = list() + slot_indices_onp = list() + for cost in problem._stacked_costs: + slots = list() + indices_onp = list() + col = 0 + num_elim_slots = 0 + for var_type, ids in problem._tangent_ordering.ordered_dict_items( + cost.sorted_ids_from_var_type + ): + ids_onp = _concrete(ids) + sorted_ids = _concrete(problem._sorted_ids_from_var_type[var_type]) + assert ids_onp.ndim == 2 + for k in range(ids_onp.shape[-1]): + global_index = onp.searchsorted(sorted_ids, ids_onp[:, k]).astype( + onp.int32 + ) + is_kept = var_type not in elim_set + if not is_kept: + num_elim_slots += 1 + slots.append( + _Slot( + type_index=kept_pos[var_type] + if is_kept + else elim_pos[var_type], + is_kept=is_kept, + col_start=col, + index=jnp.asarray(global_index), + ) + ) + indices_onp.append(global_index) + col += var_type.tangent_dim + if num_elim_slots >= 2: + elim_names = ", ".join( + t.__name__ for t in cost.sorted_ids_from_var_type if t in elim_set + ) + raise ValueError( + f"Cost '{cost._get_name()}' couples multiple eliminated " + f"variables ({elim_names}); the eliminated block would not be " + "block-diagonal. Eliminate fewer variable types." + ) + group_slots.append(tuple(slots)) + slot_indices_onp.append(indices_onp) + + combo_entries = dict() + for g, slots in enumerate(group_slots): + elim_positions = [p for p, s in enumerate(slots) if not s.is_kept] + if len(elim_positions) == 0: + continue + (e_pos,) = elim_positions + for p, s in enumerate(slots): + if s.is_kept: + key = (s.type_index, slots[e_pos].type_index) + combo_entries.setdefault(key, []).append((g, p, e_pos)) + + combos = list() + combo_kept_onp = list() + combo_elim_onp = list() + for (kt_i, et_i), entries in sorted(combo_entries.items()): + kept_idx = onp.concatenate( + [slot_indices_onp[g][p] for g, p, _ in entries], axis=0 + ) + elim_idx = onp.concatenate( + [slot_indices_onp[g][e] for g, _, e in entries], axis=0 + ) + combos.append( + _Combo( + kept_type_index=kt_i, + elim_type_index=et_i, + entries=tuple(entries), + kept_index=jnp.asarray(kept_idx), + elim_index=jnp.asarray(elim_idx), + ) + ) + combo_kept_onp.append(kept_idx) + combo_elim_onp.append(elim_idx) + + pairs = list() + for ci_a in range(len(combos)): + for ci_b in range(ci_a, len(combos)): + if combos[ci_a].elim_type_index != combos[ci_b].elim_type_index: + continue + num_elim = elim_types[combos[ci_a].elim_type_index].count + pair_a, pair_b = _matching_pairs( + combo_elim_onp[ci_a], combo_elim_onp[ci_b], num_elim + ) + if ci_a == ci_b: + keep = pair_a < pair_b + pair_a, pair_b = pair_a[keep], pair_b[keep] + if pair_a.shape[0] == 0: + continue + block_rows_all = combo_kept_onp[ci_a][pair_a] + block_cols_all = combo_kept_onp[ci_b][pair_b] + num_cols_type = kept_types[combos[ci_b].kept_type_index].count + keys = block_rows_all.astype(onp.int64) * num_cols_type + block_cols_all + unique_keys, block_id = onp.unique(keys, return_inverse=True) + + order = onp.argsort(block_id, kind="stable") + pair_a, pair_b, block_id = pair_a[order], pair_b[order], block_id[order] + pairs.append( + _CrossPairs( + combo_a=ci_a, + combo_b=ci_b, + obs_a=jnp.asarray(pair_a.astype(onp.int32)), + obs_b=jnp.asarray(pair_b.astype(onp.int32)), + block_id=jnp.asarray(block_id.astype(onp.int32)), + block_rows=jnp.asarray( + (unique_keys // num_cols_type).astype(onp.int32) + ), + block_cols=jnp.asarray( + (unique_keys % num_cols_type).astype(onp.int32) + ), + num_blocks=int(unique_keys.shape[0]), + ) + ) + + plan = EliminationPlan( + kept_types=tuple(kept_types), + elim_types=tuple(elim_types), + group_slots=tuple(group_slots), + combos=tuple(combos), + pairs=tuple(pairs), + reduced_dim=reduced_dim, + tangent_dim=problem._tangent_dim, + ) + + with jdc.copy_and_mutate(plan, validate=False) as plan: + plan.sparse_s_pattern = build_sparse_s_pattern(plan) + return plan + + +def _matching_pairs(elim_a: Any, elim_b: Any, num_elim: Any) -> Any: + order_a = onp.argsort(elim_a, kind="stable") + order_b = onp.argsort(elim_b, kind="stable") + counts_a = onp.bincount(elim_a, minlength=num_elim) + counts_b = onp.bincount(elim_b, minlength=num_elim) + starts_a = onp.cumsum(counts_a) - counts_a + starts_b = onp.cumsum(counts_b) - counts_b + pairs_per_value = counts_a * counts_b + total = int(pairs_per_value.sum()) + value_of_pair = onp.repeat(onp.arange(num_elim), pairs_per_value) + offset = onp.arange(total) - onp.repeat( + onp.cumsum(pairs_per_value) - pairs_per_value, pairs_per_value + ) + a_local = offset // counts_b[value_of_pair] + b_local = offset % counts_b[value_of_pair] + return ( + order_a[starts_a[value_of_pair] + a_local], + order_b[starts_b[value_of_pair] + b_local], + ) + + +@jdc.pytree_dataclass +class SchurFactors: + plan: Any + v_blocks: Any + cross_blocks: Any + keep_diag: Any + hcc_dense: Any + hcc_offdiag: Any + keep_jacs: Any + b_keep: Any + b_elim: Any + + +def prepare_schur( + plan: Any, + A_blocksparse: Any, + ATb: Any, + linear_solver: Any, +) -> Any: + blocks_by_group = tuple( + block_row.blocks_concat for block_row in A_blocksparse.block_rows + ) + assert len(blocks_by_group) == len(plan.group_slots) + dtype = ATb.dtype + need_dense = linear_solver == "dense_cholesky" + need_sparse = linear_solver == "cholmod" + + def slot_jac(g: Any, slot: Any) -> Any: + info = (plan.kept_types if slot.is_kept else plan.elim_types)[slot.type_index] + return blocks_by_group[g][:, :, slot.col_start : slot.col_start + info.dim] + + v_blocks = [ + jnp.zeros((info.count, info.dim, info.dim), dtype=dtype) + for info in plan.elim_types + ] + keep_diag = [ + jnp.zeros((info.count, info.dim, info.dim), dtype=dtype) + for info in plan.kept_types + ] + for g, slots in enumerate(plan.group_slots): + for slot in slots: + J = slot_jac(g, slot) + gram = _batched_gram(J, J) + target = keep_diag if slot.is_kept else v_blocks + target[slot.type_index] = target[slot.type_index] + jax.ops.segment_sum( + gram, slot.index, num_segments=target[slot.type_index].shape[0] + ) + + cross_blocks = list() + for combo in plan.combos: + parts = [ + _batched_gram( + slot_jac(g, plan.group_slots[g][kp]), + slot_jac(g, plan.group_slots[g][ep]), + ) + for g, kp, ep in combo.entries + ] + cross_blocks.append( + parts[0] if len(parts) == 1 else jnp.concatenate(parts, axis=0) + ) + + b_keep = jnp.concatenate( + [ + ATb[info.orig_start : info.orig_start + info.count * info.dim] + for info in plan.kept_types + ] + ) + b_elim = tuple( + ATb[info.orig_start : info.orig_start + info.count * info.dim].reshape( + (info.count, info.dim) + ) + for info in plan.elim_types + ) + + hcc_offdiag = list() + if need_dense or need_sparse: + for g, slots in enumerate(plan.group_slots): + kept_slots = [s for s in slots if s.is_kept] + for a in range(len(kept_slots)): + for b in range(a + 1, len(kept_slots)): + slot_a, slot_b = kept_slots[a], kept_slots[b] + hcc_offdiag.append( + _batched_gram(slot_jac(g, slot_a), slot_jac(g, slot_b)) + ) + + hcc_dense = None + if need_dense: + hcc_dense = jnp.zeros((plan.reduced_dim, plan.reduced_dim), dtype=dtype) + for kt_i, info in enumerate(plan.kept_types): + rows = ( + info.start + + jnp.arange(info.count)[:, None] * info.dim + + jnp.arange(info.dim)[None, :] + ) + hcc_dense = hcc_dense.at[rows[:, :, None], rows[:, None, :]].add( + keep_diag[kt_i] + ) + i = 0 + for g, slots in enumerate(plan.group_slots): + kept_slots = [s for s in slots if s.is_kept] + for a in range(len(kept_slots)): + for b in range(a + 1, len(kept_slots)): + slot_a, slot_b = kept_slots[a], kept_slots[b] + blk = hcc_offdiag[i] + i += 1 + rows = _block_rows(plan.kept_types[slot_a.type_index], slot_a.index) + cols = _block_rows(plan.kept_types[slot_b.type_index], slot_b.index) + hcc_dense = hcc_dense.at[rows[:, :, None], cols[:, None, :]].add( + blk + ) + hcc_dense = hcc_dense.at[cols[:, :, None], rows[:, None, :]].add( + jnp.swapaxes(blk, 1, 2) + ) + + hcc_offdiag = [] + + keep_jacs = list() + if linear_solver == "conjugate_gradient": + for g, slots in enumerate(plan.group_slots): + kept = [slot_jac(g, s) for s in slots if s.is_kept] + if len(kept) == 0: + keep_jacs.append(None) + else: + keep_jacs.append( + kept[0] if len(kept) == 1 else jnp.concatenate(kept, axis=2) + ) + else: + keep_jacs = [None] * len(plan.group_slots) + + return SchurFactors( + plan=plan, + v_blocks=tuple(v_blocks), + cross_blocks=tuple(cross_blocks), + keep_diag=tuple(keep_diag), + hcc_dense=hcc_dense, + hcc_offdiag=tuple(hcc_offdiag), + keep_jacs=tuple(keep_jacs), + b_keep=b_keep, + b_elim=b_elim, + ) + + +def _block_rows(info: Any, index: Any) -> Any: + return info.start + index[:, None] * info.dim + jnp.arange(info.dim)[None, :] + + +def _type_slice(x: Any, info: Any) -> Any: + return x[info.start : info.start + info.count * info.dim].reshape( + (info.count, info.dim) + ) + + +def _invert_psd_blocks(blocks: Any) -> Any: + dim = blocks.shape[-1] + if dim == 1: + return 1.0 / blocks + if dim == 2: + a = blocks[..., 0, 0] + b = blocks[..., 0, 1] + c = blocks[..., 1, 0] + d = blocks[..., 1, 1] + det = a * d - b * c + inv = jnp.stack( + [ + jnp.stack([d, -b], axis=-1), + jnp.stack([-c, a], axis=-1), + ], + axis=-2, + ) + return inv / det[..., None, None] + if dim == 3: + m = blocks + c00 = m[..., 1, 1] * m[..., 2, 2] - m[..., 1, 2] * m[..., 2, 1] + c01 = m[..., 1, 2] * m[..., 2, 0] - m[..., 1, 0] * m[..., 2, 2] + c02 = m[..., 1, 0] * m[..., 2, 1] - m[..., 1, 1] * m[..., 2, 0] + c10 = m[..., 0, 2] * m[..., 2, 1] - m[..., 0, 1] * m[..., 2, 2] + c11 = m[..., 0, 0] * m[..., 2, 2] - m[..., 0, 2] * m[..., 2, 0] + c12 = m[..., 0, 1] * m[..., 2, 0] - m[..., 0, 0] * m[..., 2, 1] + c20 = m[..., 0, 1] * m[..., 1, 2] - m[..., 0, 2] * m[..., 1, 1] + c21 = m[..., 0, 2] * m[..., 1, 0] - m[..., 0, 0] * m[..., 1, 2] + c22 = m[..., 0, 0] * m[..., 1, 1] - m[..., 0, 1] * m[..., 1, 0] + det = m[..., 0, 0] * c00 + m[..., 0, 1] * c01 + m[..., 0, 2] * c02 + cof = jnp.stack( + [ + jnp.stack([c00, c10, c20], axis=-1), + jnp.stack([c01, c11, c21], axis=-1), + jnp.stack([c02, c12, c22], axis=-1), + ], + axis=-2, + ) + return cof / det[..., None, None] + return jnp.linalg.inv(blocks) + + +def _damped_vinv(factors: Any, lambd: Any) -> Any: + out = list() + damp = jnp.maximum(lambd, 1e-10) + for v in factors.v_blocks: + dim = v.shape[-1] + out.append(_invert_psd_blocks(v + damp * jnp.eye(dim, dtype=v.dtype))) + return out + + +def _reduced_rhs(factors: Any, vinv: Any) -> Any: + plan = factors.plan + b_red = factors.b_keep + u = [ + jnp.einsum("nef,nf->ne", vinv_e, b_e) + for vinv_e, b_e in zip(vinv, factors.b_elim) + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + w = jnp.einsum("mte,me->mt", cross, u[combo.elim_type_index][combo.elim_index]) + contrib = jax.ops.segment_sum(w, combo.kept_index, num_segments=info.count) + b_red = b_red.at[info.start : info.start + info.count * info.dim].add( + -contrib.flatten() + ) + return b_red + + +def _wvwt_terms(factors: Any, vinv: Any) -> Any: + plan = factors.plan + y_blocks = list() + diag = [ + jnp.zeros((info.count, info.dim, info.dim), dtype=factors.b_keep.dtype) + for info in plan.kept_types + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + y = _batched_matmul(cross, vinv[combo.elim_type_index][combo.elim_index]) + y_blocks.append(y) + kt_i = combo.kept_type_index + diag[kt_i] = diag[kt_i] + jax.ops.segment_sum( + _batched_outer_last(y, cross), + combo.kept_index, + num_segments=plan.kept_types[kt_i].count, + ) + return y_blocks, diag + + +def _assemble_dense_S(factors: Any, lambd: Any, vinv: Any) -> Any: + plan = factors.plan + assert factors.hcc_dense is not None + diag_idx = jnp.arange(plan.reduced_dim) + S = factors.hcc_dense.at[diag_idx, diag_idx].add(lambd) + + y_blocks, wvwt_diag = _wvwt_terms(factors, vinv) + for pair in plan.pairs: + combo_a = plan.combos[pair.combo_a] + combo_b = plan.combos[pair.combo_b] + info_a = plan.kept_types[combo_a.kept_type_index] + info_b = plan.kept_types[combo_b.kept_type_index] + blk = _batched_outer_last( + y_blocks[pair.combo_a][pair.obs_a], + factors.cross_blocks[pair.combo_b][pair.obs_b], + ) + summed = jax.ops.segment_sum( + blk, + pair.block_id, + num_segments=pair.num_blocks, + indices_are_sorted=True, + ) + rows = _block_rows(info_a, pair.block_rows) + cols = _block_rows(info_b, pair.block_cols) + + S = S.at[rows[:, :, None], cols[:, None, :]].add(-summed) + S = S.at[cols[:, :, None], rows[:, None, :]].add(-jnp.swapaxes(summed, 1, 2)) + for kt_i, info in enumerate(plan.kept_types): + rows = ( + info.start + + jnp.arange(info.count)[:, None] * info.dim + + jnp.arange(info.dim)[None, :] + ) + S = S.at[rows[:, :, None], rows[:, None, :]].add(-wvwt_diag[kt_i]) + return S + + +def _iter_S_block_sources(plan: Any): + + for kt_i, info in enumerate(plan.kept_types): + idx = onp.arange(info.count, dtype=onp.int64) + yield "diag", kt_i, (lambda info=info, idx=idx: (info, idx, info, idx)) + + j = 0 + for slots in plan.group_slots: + kept_slots = [s for s in slots if s.is_kept] + for a in range(len(kept_slots)): + for b in range(a + 1, len(kept_slots)): + slot_a, slot_b = kept_slots[a], kept_slots[b] + yield ( + "offdiag", + j, + lambda slot_a=slot_a, slot_b=slot_b: ( + plan.kept_types[slot_a.type_index], + onp.asarray(slot_a.index), + plan.kept_types[slot_b.type_index], + onp.asarray(slot_b.index), + ), + ) + j += 1 + + for j, pair in enumerate(plan.pairs): + yield ( + "pair", + j, + lambda pair=pair: ( + plan.kept_types[plan.combos[pair.combo_a].kept_type_index], + onp.asarray(pair.block_rows), + plan.kept_types[plan.combos[pair.combo_b].kept_type_index], + onp.asarray(pair.block_cols), + ), + ) + + +def _block_rows_onp(info: Any, index: Any) -> Any: + return ( + info.start + + index[:, None] * info.dim + + onp.arange(info.dim, dtype=onp.int64)[None, :] + ) + + +def build_sparse_s_pattern(plan: Any) -> Any: + rows_parts = list() + cols_parts = list() + + def emit(rows: Any, cols: Any) -> Any: + + rows_parts.append(onp.repeat(rows, cols.shape[1], axis=1).reshape(-1)) + cols_parts.append(onp.tile(cols, (1, rows.shape[1])).reshape(-1)) + + for kind, _, geometry in _iter_S_block_sources(plan): + info_r, idx_r, info_c, idx_c = geometry() + rows = _block_rows_onp(info_r, idx_r) + cols = _block_rows_onp(info_c, idx_c) + emit(rows, cols) + if kind != "diag": + emit(cols, rows) + + rows_all = ( + onp.concatenate(rows_parts) if rows_parts else onp.zeros((0,), dtype=onp.int64) + ) + cols_all = ( + onp.concatenate(cols_parts) if cols_parts else onp.zeros((0,), dtype=onp.int64) + ) + return _SparseSPattern( + rows=jnp.asarray(rows_all.astype(onp.int32)), + cols=jnp.asarray(cols_all.astype(onp.int32)), + ) + + +def _assemble_S_values(factors: Any, lambd: Any, vinv: Any) -> Any: + plan = factors.plan + y_blocks, wvwt_diag = _wvwt_terms(factors, vinv) + parts = list() + + for kind, idx, _ in _iter_S_block_sources(plan): + if kind == "diag": + info = plan.kept_types[idx] + blk = ( + factors.keep_diag[idx] + - wvwt_diag[idx] + + lambd * jnp.eye(info.dim, dtype=factors.b_keep.dtype) + ) + elif kind == "offdiag": + blk = factors.hcc_offdiag[idx] + else: + pair = plan.pairs[idx] + blk = -jax.ops.segment_sum( + _batched_outer_last( + y_blocks[pair.combo_a][pair.obs_a], + factors.cross_blocks[pair.combo_b][pair.obs_b], + ), + pair.block_id, + num_segments=pair.num_blocks, + indices_are_sorted=True, + ) + parts.append(blk.reshape(-1)) + if kind != "diag": + parts.append(jnp.swapaxes(blk, 1, 2).reshape(-1)) + + return ( + jnp.concatenate(parts) if parts else jnp.zeros((0,), dtype=factors.b_keep.dtype) + ) + + +def _solve_spd_scaled(S: Any, b: Any) -> Any: + diag = jnp.diagonal(S) + + scale = jnp.sqrt(jnp.maximum(jnp.abs(diag), jnp.finfo(S.dtype).tiny)) + S_scaled = S / (scale[:, None] * scale[None, :]) + eps = jnp.finfo(S.dtype).eps + floor = eps * (2e4 if S.dtype == jnp.float32 else 4.0) + diag_idx = jnp.arange(S.shape[0]) + S_scaled = S_scaled.at[diag_idx, diag_idx].add(floor) + factor = jax.scipy.linalg.cho_factor(S_scaled) + return jax.scipy.linalg.cho_solve(factor, b / scale) / scale + + +def _back_substitute(factors: Any, vinv: Any, dc: Any) -> Any: + plan = factors.plan + wt_dc = [ + jnp.zeros((info.count, info.dim), dtype=dc.dtype) for info in plan.elim_types + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + dc_obs = _type_slice(dc, info)[combo.kept_index] + et_i = combo.elim_type_index + wt_dc[et_i] = wt_dc[et_i] + jax.ops.segment_sum( + jnp.einsum("mte,mt->me", cross, dc_obs), + combo.elim_index, + num_segments=plan.elim_types[et_i].count, + ) + + delta = jnp.zeros(plan.tangent_dim, dtype=dc.dtype) + for info in plan.kept_types: + delta = delta.at[info.orig_start : info.orig_start + info.count * info.dim].set( + dc[info.start : info.start + info.count * info.dim] + ) + for et_i, info in enumerate(plan.elim_types): + dl = jnp.einsum("nef,nf->ne", vinv[et_i], factors.b_elim[et_i] - wt_dc[et_i]) + delta = delta.at[info.orig_start : info.orig_start + info.count * info.dim].set( + dl.flatten() + ) + return delta + + +def solve_schur_dense(factors: Any, lambd: Any) -> Any: + vinv = _damped_vinv(factors, lambd) + b_red = _reduced_rhs(factors, vinv) + S = _assemble_dense_S(factors, lambd, vinv) + dc = _solve_spd_scaled(S, b_red) + return _back_substitute(factors, vinv, dc) + + +def solve_schur_cholmod(factors: Any, lambd: Any) -> Any: + from ._solvers import _cholmod_solve_symmetric + + plan = factors.plan + pattern = plan.sparse_s_pattern + assert pattern is not None, ( + "solve_schur_cholmod requires plan.sparse_s_pattern; build the plan " + "with the sparse path enabled." + ) + + lam_eff = jnp.asarray(lambd) + 1e-5 + vinv = _damped_vinv(factors, lam_eff) + b_red = _reduced_rhs(factors, vinv) + s_values = _assemble_S_values(factors, lam_eff, vinv) + dc = _cholmod_solve_symmetric( + s_values, pattern.rows, pattern.cols, plan.reduced_dim, b_red + ) + return _back_substitute(factors, vinv, dc) + + +def solve_schur_cg( + factors: Any, + lambd: Any, + cg_config: Any, + prev_state: Any, +) -> Any: + from ._solvers import _ConjugateGradientState + + plan = factors.plan + + lam_eff = lambd + 1e-5 + vinv = _damped_vinv(factors, lam_eff) + b_red = _reduced_rhs(factors, vinv) + + if cg_config.preconditioner is None: + precondition = lambda x: x + else: + _, wvwt_diag = _wvwt_terms(factors, vinv) + diag_blocks = [ + factors.keep_diag[kt_i] + + (lam_eff + 1e-6) * jnp.eye(info.dim, dtype=b_red.dtype) + - wvwt_diag[kt_i] + for kt_i, info in enumerate(plan.kept_types) + ] + if cg_config.preconditioner == "block_jacobi": + precond_inv = [jnp.linalg.inv(block) for block in diag_blocks] + + def precondition(x: Any) -> Any: + parts = list() + for kt_i, info in enumerate(plan.kept_types): + xs = _type_slice(x, info) + parts.append( + jnp.einsum("nij,nj->ni", precond_inv[kt_i], xs).flatten() + ) + return jnp.concatenate(parts) + elif cg_config.preconditioner == "point_jacobi": + inv_diag = [ + 1.0 / jnp.diagonal(block, axis1=1, axis2=2) for block in diag_blocks + ] + + def precondition(x: Any) -> Any: + parts = list() + for kt_i, info in enumerate(plan.kept_types): + parts.append((_type_slice(x, info) * inv_diag[kt_i]).flatten()) + return jnp.concatenate(parts) + else: + assert_never(cg_config.preconditioner) + + def matvec(x: Any) -> Any: + out = lam_eff * x + + for g, slots in enumerate(plan.group_slots): + keep_jac = factors.keep_jacs[g] + if keep_jac is None: + continue + kept_slots = [s for s in slots if s.is_kept] + gathered = list() + for slot in kept_slots: + info = plan.kept_types[slot.type_index] + gathered.append(_type_slice(x, info)[slot.index]) + x_cat = ( + gathered[0] if len(gathered) == 1 else jnp.concatenate(gathered, axis=1) + ) + residual_space = jnp.einsum("crk,ck->cr", keep_jac, x_cat) + back = jnp.einsum("crk,cr->ck", keep_jac, residual_space) + col = 0 + for slot in kept_slots: + info = plan.kept_types[slot.type_index] + contrib = jax.ops.segment_sum( + back[:, col : col + info.dim], + slot.index, + num_segments=info.count, + ) + out = out.at[info.start : info.start + info.count * info.dim].add( + contrib.flatten() + ) + col += info.dim + + wt_x = [ + jnp.zeros((info.count, info.dim), dtype=x.dtype) for info in plan.elim_types + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + xs = _type_slice(x, info)[combo.kept_index] + et_i = combo.elim_type_index + wt_x[et_i] = wt_x[et_i] + jax.ops.segment_sum( + jnp.einsum("mte,mt->me", cross, xs), + combo.elim_index, + num_segments=plan.elim_types[et_i].count, + ) + z = [jnp.einsum("nef,nf->ne", vinv_e, wt_e) for vinv_e, wt_e in zip(vinv, wt_x)] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + back = jnp.einsum( + "mte,me->mt", cross, z[combo.elim_type_index][combo.elim_index] + ) + contrib = jax.ops.segment_sum( + back, combo.kept_index, num_segments=info.count + ) + out = out.at[info.start : info.start + info.count * info.dim].add( + -contrib.flatten() + ) + return out + + b_norm = jnp.linalg.norm(b_red) + eta = jnp.minimum( + cg_config.eisenstat_walker_gamma + * (b_norm / (prev_state.ATb_norm_prev + 1e-7)) + ** cg_config.eisenstat_walker_alpha, + cg_config.tolerance_max, + ) + eta = jnp.maximum(cg_config.tolerance_min, jnp.minimum(eta, prev_state.eta)) + + dc, _ = jax.scipy.sparse.linalg.cg( + A=matvec, + b=b_red, + x0=jnp.zeros_like(b_red), + maxiter=plan.reduced_dim, + tol=eta, + M=precondition, + ) + delta = _back_substitute(factors, vinv, dc) + return delta, _ConjugateGradientState(ATb_norm_prev=b_norm, eta=eta) diff --git a/src/jaxls/_py310/_solvers.py b/src/jaxls/_py310/_solvers.py index 6d22b4f..27c5aaf 100644 --- a/src/jaxls/_py310/_solvers.py +++ b/src/jaxls/_py310/_solvers.py @@ -1,11 +1,16 @@ from __future__ import annotations from typing import Any +import contextlib +import dataclasses +import functools +import time from typing_extensions import assert_never import jax import jax.flatten_util import jax_dataclasses as jdc +import numpy as onp import scipy import scipy.sparse from jax import numpy as jnp @@ -21,6 +26,12 @@ update_al_state, update_problem_al_params, ) +from ._schur import ( + prepare_schur, + solve_schur_cg, + solve_schur_cholmod, + solve_schur_dense, +) from ._sparse_matrices import SparseCooMatrix, SparseCsrMatrix from .utils import jax_log @@ -71,6 +82,98 @@ def _cholmod_solve_on_host( return cost.solve_A(ATb) +_cholmod_symmetric_analyze_cache: Any = {} + + +def _cholmod_solve_symmetric( + s_values: Any, + rows: Any, + cols: Any, + reduced_dim: Any, + b: Any, +) -> Any: + return jax.pure_callback( + functools.partial(_cholmod_solve_symmetric_on_host, reduced_dim=reduced_dim), + b, + s_values, + rows, + cols, + b, + vmap_method="sequential", + ) + + +def _cholmod_solve_symmetric_on_host( + s_values: Any, + rows: Any, + cols: Any, + b: Any, + *, + reduced_dim: Any, +) -> Any: + import sksparse.cholmod + + rows_onp = onp.asarray(rows) + cols_onp = onp.asarray(cols) + S = scipy.sparse.coo_matrix( + (onp.asarray(s_values), (rows_onp, cols_onp)), + shape=(reduced_dim, reduced_dim), + ).tocsc() + + cache_key = (rows_onp.tobytes(), cols_onp.tobytes(), reduced_dim) + factor = _cholmod_symmetric_analyze_cache.get(cache_key, None) + if factor is None: + factor = sksparse.cholmod.analyze(S) + _cholmod_symmetric_analyze_cache[cache_key] = factor + + max_cache_size = 512 + if len(_cholmod_symmetric_analyze_cache) > max_cache_size: + _cholmod_symmetric_analyze_cache.pop( + next(iter(_cholmod_symmetric_analyze_cache)) + ) + + factor = factor.cholesky(S) + return factor.solve_A(onp.asarray(b)) + + +_active_iteration_time_recorder: Any = None + + +@contextlib.contextmanager +def record_iteration_times() -> Any: + global _active_iteration_time_recorder + prev = _active_iteration_time_recorder + times: Any = [] + _active_iteration_time_recorder = times + + _clear_solve_cache() + try: + yield times + finally: + _active_iteration_time_recorder = prev + _clear_solve_cache() + + +def _clear_solve_cache() -> Any: + + NonlinearSolver.solve.clear_cache() + + +def _record_iteration_time(anchor: Any) -> Any: + if _active_iteration_time_recorder is None: + return + + def _stamp(_anchor: Any) -> Any: + if _active_iteration_time_recorder is not None: + _active_iteration_time_recorder.append(time.perf_counter()) + + jax.debug.callback(_stamp, anchor) + + +def _compute_jacobian_scaler(column_norms: Any) -> Any: + return 1.0 / (1.0 + column_norms) + 1.0 + + @jdc.pytree_dataclass class _ConjugateGradientState: ATb_norm_prev: Any @@ -154,6 +257,7 @@ class _LmInnerState: sol_proposed: Any local_delta: Any summary: Any + lambda_growth: Any = dataclasses.field(default_factory=lambda: jnp.array(2.0)) @jdc.pytree_dataclass @@ -163,6 +267,7 @@ class _LmOuterState: lambd: Any jacobian_scaler: Any al_state: Any + last_step_accepted: Any = dataclasses.field(default_factory=lambda: jnp.array(True)) @jdc.pytree_dataclass @@ -174,6 +279,14 @@ class NonlinearSolver: sparse_mode: jdc.Static[Any] verbose: jdc.Static[Any] augmented_lagrangian: Any = None + elimination: Any = None + + def _resolve_cg_config(self) -> Any: + if isinstance(self.linear_solver, ConjugateGradientConfig): + return self.linear_solver + if self.conjugate_gradient_config is not None: + return self.conjugate_gradient_config + return ConjugateGradientConfig() @jdc.jit def solve( @@ -191,6 +304,8 @@ def solve( if self.trust_region is not None: lambda_history = lambda_history.at[0].set(self.trust_region.lambda_initial) + _record_iteration_time(cost_info.cost_total) + al_state: Any = None if self.augmented_lagrangian is not None: al_state = initialize_al_state( @@ -205,6 +320,12 @@ def solve( cost_info = problem._compute_cost_info(vals) cost_history = cost_history.at[0].set(cost_info.cost_nonconstraint) + jacobian_scaler = _compute_jacobian_scaler( + problem._compute_jac_values( + vals, cost_info.jac_cache + ).compute_column_norms() + ) + state = _LmOuterState( solution=_SolutionState( vals=vals, @@ -213,11 +334,7 @@ def solve( if self.linear_solver != "conjugate_gradient" else _ConjugateGradientState( ATb_norm_prev=0.0, - eta=( - ConjugateGradientConfig() - if self.conjugate_gradient_config is None - else self.conjugate_gradient_config - ).tolerance_max, + eta=self._resolve_cg_config().tolerance_max, ), ), summary=SolveSummary( @@ -230,17 +347,24 @@ def solve( lambd=self.trust_region.lambda_initial if self.trust_region is not None else 0.0, - jacobian_scaler=jnp.ones(problem._tangent_dim), + jacobian_scaler=jacobian_scaler, al_state=al_state, ) if self.termination.early_termination: def should_continue(state: Any) -> Any: + basic_checks = ~jnp.isnan(state.solution.cost_info.cost_total) & ( state.summary.iterations < self.termination.max_iterations ) + if self.trust_region is not None and self.augmented_lagrangian is None: + basic_checks = basic_checks & ~( + ~state.last_step_accepted + & (state.lambd >= self.trust_region.lambda_max) + ) + if self.augmented_lagrangian is None: return basic_checks & ~jnp.any(state.summary.termination_criteria) @@ -293,10 +417,12 @@ def lm_inner_step( A_multiply: Any, AT_multiply: Any, ATb: Any, + schur_factors: Any = None, ) -> Any: + if self.trust_region is not None: lambd = jnp.minimum( - inner_state.lambd * self.trust_region.lambda_factor, + inner_state.lambd * inner_state.lambda_growth, self.trust_region.lambda_max, ) else: @@ -306,15 +432,25 @@ def lm_inner_step( self._log_state(problem, sol_prev, inner_state.summary.iterations, lambd) cg_state: Any = None - if ( + if schur_factors is not None: + if self.linear_solver == "conjugate_gradient": + assert isinstance(sol_prev.cg_state, _ConjugateGradientState) + local_delta, cg_state = solve_schur_cg( + schur_factors, lambd, self._resolve_cg_config(), sol_prev.cg_state + ) + elif self.linear_solver == "dense_cholesky": + local_delta = solve_schur_dense(schur_factors, lambd) + elif self.linear_solver == "cholmod": + local_delta = solve_schur_cholmod(schur_factors, lambd) + else: + raise AssertionError( + f"Unexpected elimination plan for {self.linear_solver}." + ) + elif ( isinstance(self.linear_solver, ConjugateGradientConfig) or self.linear_solver == "conjugate_gradient" ): - cg_config = ( - ConjugateGradientConfig() - if self.linear_solver == "conjugate_gradient" - else self.linear_solver - ) + cg_config = self._resolve_cg_config() assert isinstance(sol_prev.cg_state, _ConjugateGradientState) local_delta, cg_state = cg_config._solve( problem, @@ -354,14 +490,9 @@ def lm_inner_step( if self.trust_region is None: accepted = jnp.array(True) else: - cost_predicted = jnp.sum( - ( - A_blocksparse.multiply(scaled_local_delta) - + sol_prev.cost_info.residual_vector - ) - ** 2 + predicted_reduction = 2.0 * jnp.dot(local_delta, ATb) - jnp.sum( + A_blocksparse.multiply(local_delta) ** 2 ) - predicted_reduction = sol_prev.cost_info.cost_total - cost_predicted actual_reduction = ( sol_prev.cost_info.cost_total - proposed_cost_info.cost_total ) @@ -378,9 +509,12 @@ def lm_inner_step( tangent_ordering=problem._tangent_ordering, ATb=ATb, iterations=iterations, + accepted=accepted, ) with jdc.copy_and_mutate(inner_state) as next: next.lambd = lambd + + next.lambda_growth = inner_state.lambda_growth * 2.0 next.accepted = accepted next.sol_proposed = _SolutionState( vals=proposed_vals, @@ -388,6 +522,8 @@ def lm_inner_step( cg_state=cg_state, ) next.local_delta = local_delta + + _record_iteration_time(proposed_cost_info.cost_total) next.summary = SolveSummary( iterations=iterations, termination_criteria=term_criteria, @@ -414,12 +550,6 @@ def lm_outer_step( sol_prev.vals, sol_prev.cost_info.jac_cache ) - with jdc.copy_and_mutate(state, validate=False) as state: - state.jacobian_scaler = jnp.where( - state.summary.iterations == 0, - 1.0 / (1.0 + A_blocksparse.compute_column_norms()) + 1.0, - state.jacobian_scaler, - ) A_blocksparse = A_blocksparse.scale_columns(state.jacobian_scaler) jac_values = jnp.concatenate( @@ -457,6 +587,12 @@ def lm_outer_step( ATb = -AT_multiply(sol_prev.cost_info.residual_vector) + schur_factors: Any = None + if self.elimination is not None: + schur_factors = prepare_schur( + self.elimination, A_blocksparse, ATb, linear_solver=self.linear_solver + ) + if self.trust_region is not None: init_lambd = state.lambd / self.trust_region.lambda_factor lambda_max = self.trust_region.lambda_max @@ -469,6 +605,11 @@ def lm_outer_step( init_inner_state = _LmInnerState( lambd=init_lambd, + lambda_growth=jnp.asarray( + self.trust_region.lambda_factor + if self.trust_region is not None + else 2.0 + ), accepted=jnp.array(False), sol_proposed=sol_prev, local_delta=jnp.zeros_like(ATb), @@ -490,6 +631,7 @@ def lm_outer_step( A_multiply, AT_multiply, ATb, + schur_factors, ), init_val=init_inner_state, ) @@ -513,6 +655,7 @@ def lm_outer_step( sol_prev, ) state_next.lambd = lambd_next + state_next.last_step_accepted = inner_state_final.accepted if self.verbose: jax_log( @@ -633,12 +776,14 @@ def _check_convergence( tangent_ordering: Any, ATb: Any, iterations: Any, + accepted: Any, ) -> Any: + cost_reldelta = ( jnp.abs(cost_nonconstraint_updated - sol_prev.cost_info.cost_nonconstraint) / sol_prev.cost_info.cost_nonconstraint ) - converged_cost = cost_reldelta < self.cost_tolerance + converged_cost = (cost_reldelta < self.cost_tolerance) & accepted flat_vals = jax.flatten_util.ravel_pytree(sol_prev.vals)[0] gradient_mag = jnp.max( @@ -658,7 +803,7 @@ def _check_convergence( param_delta = jnp.linalg.norm(jnp.abs(tangent)) / ( jnp.linalg.norm(flat_vals) + self.parameter_tolerance ) - converged_parameters = param_delta < self.parameter_tolerance + converged_parameters = (param_delta < self.parameter_tolerance) & accepted term_flags = jnp.array( [converged_cost, converged_gradient, converged_parameters] diff --git a/src/jaxls/_py310/_sparse_matrices.py b/src/jaxls/_py310/_sparse_matrices.py index 4d6fff5..f06cd6e 100644 --- a/src/jaxls/_py310/_sparse_matrices.py +++ b/src/jaxls/_py310/_sparse_matrices.py @@ -35,7 +35,7 @@ def to_dense(self) -> Any: out = jax.lax.dynamic_update_slice( out, update=self.blocks_concat[:, start_concat_col:end_concat_col], - start_indices=(0, start_col), + start_indices=(jnp.zeros((), dtype=start_col.dtype), start_col), ) start_concat_col = end_concat_col diff --git a/src/jaxls/_py310/_variables.py b/src/jaxls/_py310/_variables.py index 42611c2..26fddf5 100644 --- a/src/jaxls/_py310/_variables.py +++ b/src/jaxls/_py310/_variables.py @@ -111,6 +111,7 @@ def __init_subclass__( @staticmethod def _euclidean_retract(pytree: Any, delta: Any) -> Any: + flat, unravel = flatten_util.ravel_pytree(pytree) del flat return jax.tree.map(jnp.add, pytree, unravel(delta)) @@ -239,6 +240,7 @@ def _retract(self, tangent: Any, ordering: Any) -> Any: def sort_and_stack_vars(variables: Any, values: Any = None) -> Any: + vars_from_type = dict() vals_from_type = dict() if values is not None else None for i in range(len(variables)): diff --git a/src/jaxls/_py310/_visualization.py b/src/jaxls/_py310/_visualization.py index 0ecc492..a5c746d 100644 --- a/src/jaxls/_py310/_visualization.py +++ b/src/jaxls/_py310/_visualization.py @@ -226,6 +226,7 @@ def problem_show( max_costs: Any = 1000, max_variables: Any = 500, ) -> Any: + total_costs = _count_total_costs(problem) total_vars_by_type = _count_total_variables(problem) total_vars = sum(total_vars_by_type.values()) diff --git a/src/jaxls/_py310/utils.py b/src/jaxls/_py310/utils.py index 00711f4..494bb6c 100644 --- a/src/jaxls/_py310/utils.py +++ b/src/jaxls/_py310/utils.py @@ -6,9 +6,22 @@ import jax import termcolor +from jax import numpy as jnp from loguru import logger +def _batched_gram(a: Any, b: Any) -> Any: + return jnp.sum(a[..., :, :, None] * b[..., :, None, :], axis=-3) + + +def _batched_outer_last(a: Any, b: Any) -> Any: + return jnp.sum(a[..., :, None, :] * b[..., None, :, :], axis=-1) + + +def _batched_matmul(a: Any, b: Any) -> Any: + return jnp.sum(a[..., :, :, None] * b[..., None, :, :], axis=-2) + + @contextlib.contextmanager def stopwatch(label: Any = "unlabeled block") -> Any: start_time = time.time() @@ -30,6 +43,7 @@ def jax_log(fmt: Any, *args, **kwargs) -> Any: def print_deprecation_warning( message0: Any, message1: Any = None, stack_level: Any = 2 ) -> Any: + frame = inspect.currentframe() for _ in range(stack_level): if not frame: diff --git a/src/jaxls/_schur.py b/src/jaxls/_schur.py new file mode 100644 index 0000000..85f48ae --- /dev/null +++ b/src/jaxls/_schur.py @@ -0,0 +1,1169 @@ +"""Variable elimination (Schur complement) for nonlinear least squares. + +Bundle adjustment-style problems have Hessians where the landmark-landmark +block is block-diagonal: each landmark couples only to cameras, never to +other landmarks. Eliminating the landmarks analytically yields a small, much +better-conditioned "reduced camera system", which we solve directly or with +conjugate gradient before back-substituting the landmarks. + +The damped normal equations, partitioned into kept (c) and eliminated (l) +variables: + + [ H_cc W ] [ dc ] [ b_c ] + [ W^T V ] [ dl ] = [ b_l ] + +with V block-diagonal. The Schur complement eliminates dl: + + S dc = b_c - W V^{-1} b_l, S = H_cc - W V^{-1} W^T + dl = V^{-1} (b_l - W^T dc) + +Work is split into three tiers so that each solver iteration does as little +as possible: + +1. `build_elimination_plan` runs once per solve, on the host, with concrete + index arrays. It validates block-diagonality of the eliminated variables + and precomputes all index structure: per-slot variable indices, and the + observation-pair lists needed to assemble the off-diagonal blocks of a + dense S. +2. `prepare_schur` runs once per outer (Levenberg-Marquardt) iteration. It + computes everything that depends on the linearization point but not on + the damping factor: Gram blocks of the eliminated variables, kept-side + Gram diagonal/off-diagonal blocks, camera-landmark cross blocks, and the + gradient slices. +3. `solve_schur_dense` / `solve_schur_cg` run once per inner (lambda) + iteration and only redo the damping-dependent work. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, Any, assert_never, cast + +import jax +import jax_dataclasses as jdc +import numpy as onp +from jax import numpy as jnp + +from ._variables import Var +from .utils import _batched_gram, _batched_matmul, _batched_outer_last + +if TYPE_CHECKING: + from ._problem import AnalyzedLeastSquaresProblem + from ._solvers import ConjugateGradientConfig, _ConjugateGradientState + from ._sparse_matrices import BlockRowSparseMatrix + + +class _TracedVariableIdsError(ValueError): + """Raised when an elimination plan is requested but the problem's + variable IDs are tracers (the problem itself is being traced under + jax.jit), so the host-side index structure cannot be built.""" + + +@dataclasses.dataclass(frozen=True) +class _TypeInfo: + """Static layout info for one variable type (kept or eliminated).""" + + count: int + """Number of variables of this type.""" + dim: int + """Tangent dimension of one variable.""" + orig_start: int + """Start of this type's slice in the full tangent vector.""" + start: int + """Start of this type's slice in the reduced (kept) tangent vector. + Unused for eliminated types.""" + + +@jdc.pytree_dataclass +class _Slot: + """One variable slot of a stacked cost group. A cost group with N + variables per cost has N slots; each slot covers one column-block of the + group's Jacobian blocks.""" + + type_index: jdc.Static[int] + """Index into `kept_types` or `elim_types`, depending on `is_kept`.""" + is_kept: jdc.Static[bool] + col_start: jdc.Static[int] + """Start column of this slot in the group's `blocks_concat`.""" + index: jax.Array + """(num_costs,) global index of the variable within its type.""" + + +@jdc.pytree_dataclass +class _Combo: + """All crossings between one kept type and one eliminated type, + concatenated across cost groups and slots. Each element is one + "observation": a single cost row coupling one kept variable to one + eliminated variable.""" + + kept_type_index: jdc.Static[int] + elim_type_index: jdc.Static[int] + entries: jdc.Static[tuple[tuple[int, int, int], ...]] + """(group index, kept slot position, elim slot position) per entry; the + per-step cross blocks are concatenated in this order.""" + kept_index: jax.Array + """(num_obs,) kept-variable index per observation.""" + elim_index: jax.Array + """(num_obs,) eliminated-variable index per observation.""" + + +@jdc.pytree_dataclass +class _CrossPairs: + """Pairs of observations that share an eliminated variable. These define + the nonzero blocks of W V^{-1} W^T for the dense/sparse-direct reduced + solve. Built on the host once per solve. + + For `combo_a == combo_b` the pair list is canonical and strict + (obs_a < obs_b): scattering each pair block and its transpose covers both + orderings, and the obs_a == obs_b diagonal terms are exactly the + single-observation blocks already produced by `_wvwt_terms`, so they are + handled there instead of inflating the pair list. + """ + + combo_a: jdc.Static[int] + combo_b: jdc.Static[int] + obs_a: jax.Array + """(num_pairs,) index into combo_a's observation list.""" + obs_b: jax.Array + """(num_pairs,) index into combo_b's observation list.""" + block_id: jax.Array + """(num_pairs,) index of the (kept_a, kept_b) block this pair adds to.""" + block_rows: jax.Array + """(num_blocks,) kept-variable index (within combo_a's kept type) of each block.""" + block_cols: jax.Array + """(num_blocks,) kept-variable index (within combo_b's kept type) of each block.""" + num_blocks: jdc.Static[int] + + +@jdc.pytree_dataclass +class _SparseSPattern: + """Host-precomputed COO structure of the reduced matrix S, for the sparse + (CHOLMOD) reduced solve. The off-diagonal kept-kept blocks of S sit at + fixed positions determined only by the problem's incidence structure, so + the scalar (row, col) coordinates can be built once on the host. The + matching values are produced per inner iteration by `_assemble_S_values`, + in the same block order used here. + + We emit the FULL symmetric matrix (both triangles). CHOLMOD reads only the + lower triangle, but a consistent full matrix keeps the scipy COO->CSC + conversion and the cache key unambiguous.""" + + rows: jax.Array + """(nnz,) scalar row indices, concatenated over all block sources.""" + cols: jax.Array + """(nnz,) scalar column indices, in the same order as `rows`.""" + + +@jdc.pytree_dataclass +class EliminationPlan: + """Static structure for Schur-complement variable elimination. Built once + per solve by `build_elimination_plan`; safe to reuse across solves with + the same problem structure.""" + + kept_types: jdc.Static[tuple[_TypeInfo, ...]] + elim_types: jdc.Static[tuple[_TypeInfo, ...]] + group_slots: tuple[tuple[_Slot, ...], ...] + combos: tuple[_Combo, ...] + pairs: tuple[_CrossPairs, ...] + reduced_dim: jdc.Static[int] + """Total tangent dimension of the kept variables.""" + tangent_dim: jdc.Static[int] + """Total tangent dimension of the full problem.""" + sparse_s_pattern: _SparseSPattern | None = None + """COO coordinates of S for the CHOLMOD reduced solve; None unless the + sparse path is requested.""" + + +def infer_eliminate( + problem: AnalyzedLeastSquaresProblem, +) -> tuple[type[Var[Any]], ...]: + """Choose variable types to eliminate automatically. + + A set of types is *eligible* when no single cost touches more than one + variable from it (the eliminated Hessian block is then block-diagonal). + Types are added greedily by total tangent size — in bundle adjustment + the landmarks, which dominate the problem — preferring many small + variables over few large ones on ties, and never eliminating every type. + + Elimination only pays off when the eliminated block is large: the chosen + set is returned only if it covers at least half of the total tangent + dimension, otherwise this returns `()` (no elimination). + + Only static structure is inspected, so this is safe to call under + `jax.jit` tracing. + """ + # Slots per cost group for each variable type. Shapes are static. + slots_per_group = [ + { + var_type: ids.shape[-1] + for var_type, ids in cost.sorted_ids_from_var_type.items() + } + for cost in problem._stacked_costs + ] + + def total_elim_slots_ok(candidate: list[type[Var[Any]]]) -> bool: + return all( + sum(slots.get(var_type, 0) for var_type in candidate) <= 1 + for slots in slots_per_group + ) + + sizes = { + var_type: (int(ids.shape[0]) * var_type.tangent_dim, int(ids.shape[0])) + for var_type, ids in problem._sorted_ids_from_var_type.items() + } + # Largest tangent block first; prefer many small variables on ties + # (finer block diagonal -> cheaper to invert). + candidates = sorted(sizes, key=lambda t: sizes[t], reverse=True) + + chosen = list[type[Var[Any]]]() + for var_type in candidates: + if len(chosen) + 1 == len(candidates): + break # At least one type must be kept. + if total_elim_slots_ok(chosen + [var_type]): + chosen.append(var_type) + + eliminated_dim = sum(sizes[t][0] for t in chosen) + if eliminated_dim * 2 < problem._tangent_dim: + return () + return tuple(chosen) + + +def build_elimination_plan( + problem: AnalyzedLeastSquaresProblem, + eliminate: tuple[type[Var[Any]], ...], +) -> EliminationPlan: + """Validate `eliminate` and precompute all index structure for the + Schur-complement solve. Must be called outside of `jax.jit`, with + concrete variable IDs.""" + # Deduplicate while preserving order. + elim_set = list[type[Var[Any]]]() + for var_type in eliminate: + if var_type not in elim_set: + elim_set.append(var_type) + + for var_type in elim_set: + if var_type not in problem._sorted_ids_from_var_type: + raise ValueError( + f"eliminate={var_type.__name__} was requested, but no variables " + "of this type exist in the problem." + ) + + def _concrete(array: jax.Array) -> onp.ndarray: + try: + return onp.asarray(array) + except Exception as e: + raise _TracedVariableIdsError( + "Variable elimination requires concrete variable IDs; the " + "problem appears to be traced inside jax.jit." + ) from e + + # Partition variable types into kept and eliminated, with layout offsets. + kept_types = list[_TypeInfo]() + elim_types = list[_TypeInfo]() + kept_pos = dict[type[Var[Any]], int]() + elim_pos = dict[type[Var[Any]], int]() + reduced_dim = 0 + for var_type, ids in problem._tangent_ordering.ordered_dict_items( + problem._sorted_ids_from_var_type + ): + (count,) = ids.shape + if var_type in elim_set: + elim_pos[var_type] = len(elim_types) + elim_types.append( + _TypeInfo( + count=count, + dim=var_type.tangent_dim, + orig_start=problem._tangent_start_from_var_type[var_type], + start=0, + ) + ) + else: + kept_pos[var_type] = len(kept_types) + kept_types.append( + _TypeInfo( + count=count, + dim=var_type.tangent_dim, + orig_start=problem._tangent_start_from_var_type[var_type], + start=reduced_dim, + ) + ) + reduced_dim += count * var_type.tangent_dim + + if len(kept_types) == 0: + raise ValueError( + "All variable types were marked for elimination; at least one " + "type must be kept to form the reduced system." + ) + + # Per-group slots. Iteration order must match how `_compute_jac_values` + # lays out columns in `blocks_concat`: variable types in tangent order, + # then slots in ID-column order. + group_slots = list[tuple[_Slot, ...]]() + slot_indices_onp = list[list[onp.ndarray]]() # Host copies, for pairing. + for cost in problem._stacked_costs: + slots = list[_Slot]() + indices_onp = list[onp.ndarray]() + col = 0 + num_elim_slots = 0 + for var_type, ids in problem._tangent_ordering.ordered_dict_items( + cost.sorted_ids_from_var_type + ): + ids_onp = _concrete(ids) + sorted_ids = _concrete(problem._sorted_ids_from_var_type[var_type]) + assert ids_onp.ndim == 2 + for k in range(ids_onp.shape[-1]): + global_index = onp.searchsorted(sorted_ids, ids_onp[:, k]).astype( + onp.int32 + ) + is_kept = var_type not in elim_set + if not is_kept: + num_elim_slots += 1 + slots.append( + _Slot( + type_index=kept_pos[var_type] + if is_kept + else elim_pos[var_type], + is_kept=is_kept, + col_start=col, + index=jnp.asarray(global_index), + ) + ) + indices_onp.append(global_index) + col += var_type.tangent_dim + if num_elim_slots >= 2: + elim_names = ", ".join( + t.__name__ for t in cost.sorted_ids_from_var_type if t in elim_set + ) + raise ValueError( + f"Cost '{cost._get_name()}' couples multiple eliminated " + f"variables ({elim_names}); the eliminated block would not be " + "block-diagonal. Eliminate fewer variable types." + ) + group_slots.append(tuple(slots)) + slot_indices_onp.append(indices_onp) + + # Combos: one per (kept type, eliminated type) with at least one cost + # row coupling them. + combo_entries = dict[tuple[int, int], list[tuple[int, int, int]]]() + for g, slots in enumerate(group_slots): + elim_positions = [p for p, s in enumerate(slots) if not s.is_kept] + if len(elim_positions) == 0: + continue + (e_pos,) = elim_positions + for p, s in enumerate(slots): + if s.is_kept: + key = (s.type_index, slots[e_pos].type_index) + combo_entries.setdefault(key, []).append((g, p, e_pos)) + + combos = list[_Combo]() + combo_kept_onp = list[onp.ndarray]() + combo_elim_onp = list[onp.ndarray]() + for (kt_i, et_i), entries in sorted(combo_entries.items()): + kept_idx = onp.concatenate( + [slot_indices_onp[g][p] for g, p, _ in entries], axis=0 + ) + elim_idx = onp.concatenate( + [slot_indices_onp[g][e] for g, _, e in entries], axis=0 + ) + combos.append( + _Combo( + kept_type_index=kt_i, + elim_type_index=et_i, + entries=tuple(entries), + kept_index=jnp.asarray(kept_idx), + elim_index=jnp.asarray(elim_idx), + ) + ) + combo_kept_onp.append(kept_idx) + combo_elim_onp.append(elim_idx) + + # Pair lists for the dense reduced system: observations sharing an + # eliminated variable. + pairs = list[_CrossPairs]() + for ci_a in range(len(combos)): + for ci_b in range(ci_a, len(combos)): + if combos[ci_a].elim_type_index != combos[ci_b].elim_type_index: + continue + num_elim = elim_types[combos[ci_a].elim_type_index].count + pair_a, pair_b = _matching_pairs( + combo_elim_onp[ci_a], combo_elim_onp[ci_b], num_elim + ) + if ci_a == ci_b: + # Canonical strict pairs; transposes are handled at scatter + # time, and the obs_a == obs_b diagonal via `_wvwt_terms`. + keep = pair_a < pair_b + pair_a, pair_b = pair_a[keep], pair_b[keep] + if pair_a.shape[0] == 0: + continue + block_rows_all = combo_kept_onp[ci_a][pair_a] + block_cols_all = combo_kept_onp[ci_b][pair_b] + num_cols_type = kept_types[combos[ci_b].kept_type_index].count + keys = block_rows_all.astype(onp.int64) * num_cols_type + block_cols_all + unique_keys, block_id = onp.unique(keys, return_inverse=True) + # Sort pairs by block so the per-iteration segment_sum can take + # the sorted-indices fast path. + order = onp.argsort(block_id, kind="stable") + pair_a, pair_b, block_id = pair_a[order], pair_b[order], block_id[order] + pairs.append( + _CrossPairs( + combo_a=ci_a, + combo_b=ci_b, + obs_a=jnp.asarray(pair_a.astype(onp.int32)), + obs_b=jnp.asarray(pair_b.astype(onp.int32)), + block_id=jnp.asarray(block_id.astype(onp.int32)), + block_rows=jnp.asarray( + (unique_keys // num_cols_type).astype(onp.int32) + ), + block_cols=jnp.asarray( + (unique_keys % num_cols_type).astype(onp.int32) + ), + num_blocks=int(unique_keys.shape[0]), + ) + ) + + plan = EliminationPlan( + kept_types=tuple(kept_types), + elim_types=tuple(elim_types), + group_slots=tuple(group_slots), + combos=tuple(combos), + pairs=tuple(pairs), + reduced_dim=reduced_dim, + tangent_dim=problem._tangent_dim, + ) + # Precompute the COO structure of S for the CHOLMOD reduced solve. This is + # cheap host-side index work and makes the plan solver-agnostic: the dense + # and CG paths simply ignore it. + with jdc.copy_and_mutate(plan, validate=False) as plan: + plan.sparse_s_pattern = build_sparse_s_pattern(plan) + return plan + + +def _matching_pairs( + elim_a: onp.ndarray, elim_b: onp.ndarray, num_elim: int +) -> tuple[onp.ndarray, onp.ndarray]: + """All index pairs (i, j) with elim_a[i] == elim_b[j], vectorized.""" + order_a = onp.argsort(elim_a, kind="stable") + order_b = onp.argsort(elim_b, kind="stable") + counts_a = onp.bincount(elim_a, minlength=num_elim) + counts_b = onp.bincount(elim_b, minlength=num_elim) + starts_a = onp.cumsum(counts_a) - counts_a + starts_b = onp.cumsum(counts_b) - counts_b + pairs_per_value = counts_a * counts_b + total = int(pairs_per_value.sum()) + value_of_pair = onp.repeat(onp.arange(num_elim), pairs_per_value) + offset = onp.arange(total) - onp.repeat( + onp.cumsum(pairs_per_value) - pairs_per_value, pairs_per_value + ) + a_local = offset // counts_b[value_of_pair] + b_local = offset % counts_b[value_of_pair] + return ( + order_a[starts_a[value_of_pair] + a_local], + order_b[starts_b[value_of_pair] + b_local], + ) + + +@jdc.pytree_dataclass +class SchurFactors: + """Damping-independent quantities for one outer iteration's Schur solve. + Built by `prepare_schur` from the (column-scaled) Jacobian and gradient.""" + + plan: EliminationPlan + v_blocks: tuple[jax.Array, ...] + """Per eliminated type: (count, dim, dim) undamped Gram blocks V.""" + cross_blocks: tuple[jax.Array, ...] + """Per combo: (num_obs, kept_dim, elim_dim) blocks of W.""" + keep_diag: tuple[jax.Array, ...] + """Per kept type: (count, dim, dim) block diagonal of H_cc.""" + hcc_dense: jax.Array | None + """(reduced_dim, reduced_dim) dense H_cc; None for the CG and sparse paths.""" + hcc_offdiag: tuple[jax.Array, ...] + """Per off-diagonal kept-kept source (in `_S_block_index_sources` order): + (K, dim_a, dim_b) H_cc coupling blocks, for the sparse (CHOLMOD) path. + Empty for the dense and CG paths.""" + keep_jacs: tuple[jax.Array | None, ...] + """Per group: (num_costs, residual_dim, total_kept_width) kept-slot + Jacobian columns, for the matrix-free H_cc product; None for the + dense path or for groups with no kept slots.""" + b_keep: jax.Array + """(reduced_dim,) kept slice of ATb, in reduced layout.""" + b_elim: tuple[jax.Array, ...] + """Per eliminated type: (count, dim) slice of ATb.""" + + +def prepare_schur( + plan: EliminationPlan, + A_blocksparse: BlockRowSparseMatrix, + ATb: jax.Array, + linear_solver: str, +) -> SchurFactors: + """Compute all damping-independent Schur quantities for one outer step.""" + blocks_by_group = tuple( + block_row.blocks_concat for block_row in A_blocksparse.block_rows + ) + assert len(blocks_by_group) == len(plan.group_slots) + dtype = ATb.dtype + need_dense = linear_solver == "dense_cholesky" + need_sparse = linear_solver == "cholmod" + + def slot_jac(g: int, slot: _Slot) -> jax.Array: + info = (plan.kept_types if slot.is_kept else plan.elim_types)[slot.type_index] + return blocks_by_group[g][:, :, slot.col_start : slot.col_start + info.dim] + + # Gram blocks of the eliminated variables (V) and the kept block + # diagonal of H_cc. Both are index-keyed accumulations: use segment_sum, + # not scatter-add. + v_blocks = [ + jnp.zeros((info.count, info.dim, info.dim), dtype=dtype) + for info in plan.elim_types + ] + keep_diag = [ + jnp.zeros((info.count, info.dim, info.dim), dtype=dtype) + for info in plan.kept_types + ] + for g, slots in enumerate(plan.group_slots): + for slot in slots: + J = slot_jac(g, slot) + gram = _batched_gram(J, J) + target = keep_diag if slot.is_kept else v_blocks + target[slot.type_index] = target[slot.type_index] + jax.ops.segment_sum( + gram, slot.index, num_segments=target[slot.type_index].shape[0] + ) + + # Cross blocks (W), one entry per observation, concatenated per combo. + cross_blocks = list[jax.Array]() + for combo in plan.combos: + parts = [ + _batched_gram( + slot_jac(g, plan.group_slots[g][kp]), + slot_jac(g, plan.group_slots[g][ep]), + ) + for g, kp, ep in combo.entries + ] + cross_blocks.append( + parts[0] if len(parts) == 1 else jnp.concatenate(parts, axis=0) + ) + + # Gradient slices, rearranged into the reduced/eliminated layouts. + b_keep = jnp.concatenate( + [ + ATb[info.orig_start : info.orig_start + info.count * info.dim] + for info in plan.kept_types + ] + ) + b_elim = tuple( + ATb[info.orig_start : info.orig_start + info.count * info.dim].reshape( + (info.count, info.dim) + ) + for info in plan.elim_types + ) + + # H_cc off-diagonal kept-kept couplings within each cost group. Needed by + # both direct paths: the dense path scatters them into `hcc_dense`, the + # sparse (CHOLMOD) path keeps them as blocks in `hcc_offdiag` (in the same + # order as `_S_block_index_sources`). + hcc_offdiag = list[jax.Array]() + if need_dense or need_sparse: + for g, slots in enumerate(plan.group_slots): + kept_slots = [s for s in slots if s.is_kept] + for a in range(len(kept_slots)): + for b in range(a + 1, len(kept_slots)): + slot_a, slot_b = kept_slots[a], kept_slots[b] + hcc_offdiag.append( + _batched_gram(slot_jac(g, slot_a), slot_jac(g, slot_b)) + ) + + # Dense H_cc for the dense direct path: scatter the block diagonal, then + # the off-diagonal couplings computed above. + hcc_dense = None + if need_dense: + hcc_dense = jnp.zeros((plan.reduced_dim, plan.reduced_dim), dtype=dtype) + for kt_i, info in enumerate(plan.kept_types): + rows = ( + info.start + + jnp.arange(info.count)[:, None] * info.dim + + jnp.arange(info.dim)[None, :] + ) + hcc_dense = hcc_dense.at[rows[:, :, None], rows[:, None, :]].add( + keep_diag[kt_i] + ) + i = 0 + for g, slots in enumerate(plan.group_slots): + kept_slots = [s for s in slots if s.is_kept] + for a in range(len(kept_slots)): + for b in range(a + 1, len(kept_slots)): + slot_a, slot_b = kept_slots[a], kept_slots[b] + blk = hcc_offdiag[i] + i += 1 + rows = _block_rows(plan.kept_types[slot_a.type_index], slot_a.index) + cols = _block_rows(plan.kept_types[slot_b.type_index], slot_b.index) + hcc_dense = hcc_dense.at[rows[:, :, None], cols[:, None, :]].add( + blk + ) + hcc_dense = hcc_dense.at[cols[:, :, None], rows[:, None, :]].add( + jnp.swapaxes(blk, 1, 2) + ) + # The dense path scatters off-diagonals into `hcc_dense`; don't also + # carry them as separate blocks. + hcc_offdiag = [] + + # Kept-column Jacobian per group, for the matrix-free H_cc product. + keep_jacs = list[jax.Array | None]() + if linear_solver == "conjugate_gradient": + for g, slots in enumerate(plan.group_slots): + kept = [slot_jac(g, s) for s in slots if s.is_kept] + if len(kept) == 0: + keep_jacs.append(None) + else: + keep_jacs.append( + kept[0] if len(kept) == 1 else jnp.concatenate(kept, axis=2) + ) + else: + keep_jacs = [None] * len(plan.group_slots) + + return SchurFactors( + plan=plan, + v_blocks=tuple(v_blocks), + cross_blocks=tuple(cross_blocks), + keep_diag=tuple(keep_diag), + hcc_dense=hcc_dense, + hcc_offdiag=tuple(hcc_offdiag), + keep_jacs=tuple(keep_jacs), + b_keep=b_keep, + b_elim=b_elim, + ) + + +def _block_rows(info: _TypeInfo, index: jax.Array) -> jax.Array: + """Tangent row indices for a batch of variables of one type. Shape + (len(index), dim), in the reduced layout.""" + return info.start + index[:, None] * info.dim + jnp.arange(info.dim)[None, :] + + +def _type_slice(x: jax.Array, info: _TypeInfo) -> jax.Array: + """One variable type's slice of a reduced-layout vector, as (count, dim).""" + return x[info.start : info.start + info.count * info.dim].reshape( + (info.count, info.dim) + ) + + +def _invert_psd_blocks(blocks: jax.Array) -> jax.Array: + """Invert a batch of small SPD matrices. Closed-form for dim <= 3, which + avoids a slow batched LAPACK loop on the per-iteration hot path.""" + dim = blocks.shape[-1] + if dim == 1: + return 1.0 / blocks + if dim == 2: + a = blocks[..., 0, 0] + b = blocks[..., 0, 1] + c = blocks[..., 1, 0] + d = blocks[..., 1, 1] + det = a * d - b * c + inv = jnp.stack( + [ + jnp.stack([d, -b], axis=-1), + jnp.stack([-c, a], axis=-1), + ], + axis=-2, + ) + return inv / det[..., None, None] + if dim == 3: + m = blocks + c00 = m[..., 1, 1] * m[..., 2, 2] - m[..., 1, 2] * m[..., 2, 1] + c01 = m[..., 1, 2] * m[..., 2, 0] - m[..., 1, 0] * m[..., 2, 2] + c02 = m[..., 1, 0] * m[..., 2, 1] - m[..., 1, 1] * m[..., 2, 0] + c10 = m[..., 0, 2] * m[..., 2, 1] - m[..., 0, 1] * m[..., 2, 2] + c11 = m[..., 0, 0] * m[..., 2, 2] - m[..., 0, 2] * m[..., 2, 0] + c12 = m[..., 0, 1] * m[..., 2, 0] - m[..., 0, 0] * m[..., 2, 1] + c20 = m[..., 0, 1] * m[..., 1, 2] - m[..., 0, 2] * m[..., 1, 1] + c21 = m[..., 0, 2] * m[..., 1, 0] - m[..., 0, 0] * m[..., 1, 2] + c22 = m[..., 0, 0] * m[..., 1, 1] - m[..., 0, 1] * m[..., 1, 0] + det = m[..., 0, 0] * c00 + m[..., 0, 1] * c01 + m[..., 0, 2] * c02 + cof = jnp.stack( + [ + jnp.stack([c00, c10, c20], axis=-1), + jnp.stack([c01, c11, c21], axis=-1), + jnp.stack([c02, c12, c22], axis=-1), + ], + axis=-2, + ) + return cof / det[..., None, None] + return jnp.linalg.inv(blocks) + + +def _damped_vinv(factors: SchurFactors, lambd: jax.Array | float) -> list[jax.Array]: + """(V + lambda I)^{-1} per eliminated type. + + A tiny floor on the damping keeps V invertible even at lambda=0 + (Gauss-Newton, trust_region=None), where an under-observed eliminated + variable — e.g. a 3D landmark seen by a single 2-residual camera — gives + a rank-deficient V whose closed-form inverse would divide by det=0 and + NaN the whole step. Under Levenberg-Marquardt lambda >= lambda_min (1e-5) + so the floor never binds; it only guards the GN edge. + """ + out = list[jax.Array]() + damp = jnp.maximum(lambd, 1e-10) + for v in factors.v_blocks: + dim = v.shape[-1] + out.append(_invert_psd_blocks(v + damp * jnp.eye(dim, dtype=v.dtype))) + return out + + +def _reduced_rhs(factors: SchurFactors, vinv: list[jax.Array]) -> jax.Array: + """b_c - W V^{-1} b_l, in the reduced layout.""" + plan = factors.plan + b_red = factors.b_keep + u = [ + jnp.einsum("nef,nf->ne", vinv_e, b_e) + for vinv_e, b_e in zip(vinv, factors.b_elim) + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + w = jnp.einsum("mte,me->mt", cross, u[combo.elim_type_index][combo.elim_index]) + contrib = jax.ops.segment_sum(w, combo.kept_index, num_segments=info.count) + b_red = b_red.at[info.start : info.start + info.count * info.dim].add( + -contrib.flatten() + ) + return b_red + + +def _wvwt_terms( + factors: SchurFactors, vinv: list[jax.Array] +) -> tuple[list[jax.Array], list[jax.Array]]: + """Per-combo Y = W_obs V^{-1} blocks, and the per-kept-type block-diagonal + of W V^{-1} W^T restricted to single observations (exact when no two + observations share the same kept and eliminated variable pair).""" + plan = factors.plan + y_blocks = list[jax.Array]() + diag = [ + jnp.zeros((info.count, info.dim, info.dim), dtype=factors.b_keep.dtype) + for info in plan.kept_types + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + y = _batched_matmul(cross, vinv[combo.elim_type_index][combo.elim_index]) + y_blocks.append(y) + kt_i = combo.kept_type_index + diag[kt_i] = diag[kt_i] + jax.ops.segment_sum( + _batched_outer_last(y, cross), + combo.kept_index, + num_segments=plan.kept_types[kt_i].count, + ) + return y_blocks, diag + + +def _assemble_dense_S( + factors: SchurFactors, lambd: jax.Array | float, vinv: list[jax.Array] +) -> jax.Array: + """S = H_cc + lambda I - W (V + lambda I)^{-1} W^T, dense.""" + plan = factors.plan + assert factors.hcc_dense is not None + diag_idx = jnp.arange(plan.reduced_dim) + S = factors.hcc_dense.at[diag_idx, diag_idx].add(lambd) + + y_blocks, wvwt_diag = _wvwt_terms(factors, vinv) + for pair in plan.pairs: + combo_a = plan.combos[pair.combo_a] + combo_b = plan.combos[pair.combo_b] + info_a = plan.kept_types[combo_a.kept_type_index] + info_b = plan.kept_types[combo_b.kept_type_index] + blk = _batched_outer_last( + y_blocks[pair.combo_a][pair.obs_a], + factors.cross_blocks[pair.combo_b][pair.obs_b], + ) + summed = jax.ops.segment_sum( + blk, + pair.block_id, + num_segments=pair.num_blocks, + indices_are_sorted=True, + ) + rows = _block_rows(info_a, pair.block_rows) + cols = _block_rows(info_b, pair.block_cols) + # Each pair block and its transpose. Same-combo pair lists are + # strict (obs_a < obs_b); the obs_a == obs_b diagonal terms are the + # `wvwt_diag` blocks subtracted below. + S = S.at[rows[:, :, None], cols[:, None, :]].add(-summed) + S = S.at[cols[:, :, None], rows[:, None, :]].add(-jnp.swapaxes(summed, 1, 2)) + for kt_i, info in enumerate(plan.kept_types): + rows = ( + info.start + + jnp.arange(info.count)[:, None] * info.dim + + jnp.arange(info.dim)[None, :] + ) + S = S.at[rows[:, :, None], rows[:, None, :]].add(-wvwt_diag[kt_i]) + return S + + +def _iter_S_block_sources(plan: EliminationPlan): + """Single source of truth for the block structure of the reduced matrix S, + shared by the index builder (`build_sparse_s_pattern`) and the value builder + (`_assemble_S_values`) so the two cannot drift out of sync. Yields one + descriptor `(kind, idx, geometry)` per dense block, in assembly order: + + - ("diag", kt_i, geometry): the (count, dim, dim) diagonal block of kept + type kt_i (keep_diag - wvwt_diag + lambda*I). + - ("offdiag", j, geometry): the j-th H_cc off-diagonal kept-kept coupling + block, then its transpose. + - ("pair", j, geometry): the j-th -W V^{-1} W^T pair block, then its + transpose. + + Off-diagonal sources are emitted as a (block, transpose) pair: the index + builder swaps (row, col) for the transpose; the value builder swaps axes. + + `geometry` is a zero-arg callable returning (info_row, idx_row, info_col, + idx_col) host numpy index arrays — used ONLY by the index builder. It is a + thunk so the value builder (which runs under jax.jit, where the underlying + `slot.index` / `pair.block_rows` are tracers) never materializes it. + """ + # 1. Per-kept-type diagonal blocks. + for kt_i, info in enumerate(plan.kept_types): + idx = onp.arange(info.count, dtype=onp.int64) + yield "diag", kt_i, (lambda info=info, idx=idx: (info, idx, info, idx)) + + # 2. H_cc off-diagonal kept-kept couplings within each cost group. + j = 0 + for slots in plan.group_slots: + kept_slots = [s for s in slots if s.is_kept] + for a in range(len(kept_slots)): + for b in range(a + 1, len(kept_slots)): + slot_a, slot_b = kept_slots[a], kept_slots[b] + yield ( + "offdiag", + j, + lambda slot_a=slot_a, slot_b=slot_b: ( + plan.kept_types[slot_a.type_index], + onp.asarray(slot_a.index), + plan.kept_types[slot_b.type_index], + onp.asarray(slot_b.index), + ), + ) + j += 1 + + # 3. -W V^{-1} W^T off-diagonal pair blocks. + for j, pair in enumerate(plan.pairs): + yield ( + "pair", + j, + lambda pair=pair: ( + plan.kept_types[plan.combos[pair.combo_a].kept_type_index], + onp.asarray(pair.block_rows), + plan.kept_types[plan.combos[pair.combo_b].kept_type_index], + onp.asarray(pair.block_cols), + ), + ) + + +def _block_rows_onp(info: _TypeInfo, index: onp.ndarray) -> onp.ndarray: + """Host-side `_block_rows`: tangent rows (reduced layout) for a batch of + variables of one type. Shape (len(index), dim).""" + return ( + info.start + + index[:, None] * info.dim + + onp.arange(info.dim, dtype=onp.int64)[None, :] + ) + + +def build_sparse_s_pattern(plan: EliminationPlan) -> _SparseSPattern: + """Build the host-side COO coordinates of S, emitting the full symmetric + matrix (each off-diagonal source adds both the block and its transpose). + Must run outside jax.jit (the block indices come from concrete plan + structure). The block order matches `_assemble_S_values`.""" + rows_parts = list[onp.ndarray]() + cols_parts = list[onp.ndarray]() + + def emit(rows: onp.ndarray, cols: onp.ndarray) -> None: + # Expand each (K, dim_r)/(K, dim_c) block into K*dim_r*dim_c scalar + # (row, col) pairs, row-major within each block: every row index is + # paired with every col index of the same block. + rows_parts.append(onp.repeat(rows, cols.shape[1], axis=1).reshape(-1)) + cols_parts.append(onp.tile(cols, (1, rows.shape[1])).reshape(-1)) + + for kind, _, geometry in _iter_S_block_sources(plan): + info_r, idx_r, info_c, idx_c = geometry() + rows = _block_rows_onp(info_r, idx_r) + cols = _block_rows_onp(info_c, idx_c) + emit(rows, cols) + if kind != "diag": + emit(cols, rows) # transpose + + rows_all = ( + onp.concatenate(rows_parts) if rows_parts else onp.zeros((0,), dtype=onp.int64) + ) + cols_all = ( + onp.concatenate(cols_parts) if cols_parts else onp.zeros((0,), dtype=onp.int64) + ) + return _SparseSPattern( + rows=jnp.asarray(rows_all.astype(onp.int32)), + cols=jnp.asarray(cols_all.astype(onp.int32)), + ) + + +def _assemble_S_values( + factors: SchurFactors, lambd: jax.Array | float, vinv: list[jax.Array] +) -> jax.Array: + """Flat values of S matching `plan.sparse_s_pattern` coordinates, in the + block order of `_iter_S_block_sources`. Reuses the same block quantities as + `_assemble_dense_S`.""" + plan = factors.plan + y_blocks, wvwt_diag = _wvwt_terms(factors, vinv) + parts = list[jax.Array]() + + for kind, idx, _ in _iter_S_block_sources(plan): + if kind == "diag": + info = plan.kept_types[idx] + blk = ( + factors.keep_diag[idx] + - wvwt_diag[idx] + + lambd * jnp.eye(info.dim, dtype=factors.b_keep.dtype) + ) + elif kind == "offdiag": + # Damping-independent; precomputed in `prepare_schur`. + blk = factors.hcc_offdiag[idx] + else: # "pair": -W V^{-1} W^T. + pair = plan.pairs[idx] + blk = -jax.ops.segment_sum( + _batched_outer_last( + y_blocks[pair.combo_a][pair.obs_a], + factors.cross_blocks[pair.combo_b][pair.obs_b], + ), + pair.block_id, + num_segments=pair.num_blocks, + indices_are_sorted=True, + ) + parts.append(blk.reshape(-1)) + if kind != "diag": + parts.append(jnp.swapaxes(blk, 1, 2).reshape(-1)) # transpose + + return ( + jnp.concatenate(parts) if parts else jnp.zeros((0,), dtype=factors.b_keep.dtype) + ) + + +def _solve_spd_scaled(S: jax.Array, b: jax.Array) -> jax.Array: + """Solve S x = b with S symmetric positive definite, robustly. + + Forming S = H_cc - W V^{-1} W^T cancels catastrophically in float32 (the + eliminated variables absorb most of the kept-variable information), which + can leave S numerically indefinite. Jacobi scaling plus a + precision-adaptive Tikhonov floor keeps the factorization finite without + measurably perturbing float64 solves. + """ + diag = jnp.diagonal(S) + # Use |diag|, not max(diag, tiny): float32 cancellation can leave a + # diagonal entry slightly negative, and max(., tiny) would collapse its + # scale to ~sqrt(tiny) and blow that row up by ~1e19. abs keeps the + # Jacobi scaling well-conditioned; the Tikhonov floor below restores PD. + scale = jnp.sqrt(jnp.maximum(jnp.abs(diag), jnp.finfo(S.dtype).tiny)) + S_scaled = S / (scale[:, None] * scale[None, :]) + eps = jnp.finfo(S.dtype).eps + floor = eps * (2e4 if S.dtype == jnp.float32 else 4.0) + diag_idx = jnp.arange(S.shape[0]) + S_scaled = S_scaled.at[diag_idx, diag_idx].add(floor) + factor = jax.scipy.linalg.cho_factor(S_scaled) + return jax.scipy.linalg.cho_solve(factor, b / scale) / scale + + +def _back_substitute( + factors: SchurFactors, vinv: list[jax.Array], dc: jax.Array +) -> jax.Array: + """Recover the eliminated update and scatter both into the full tangent + layout: dl = (V + lambda I)^{-1} (b_l - W^T dc).""" + plan = factors.plan + wt_dc = [ + jnp.zeros((info.count, info.dim), dtype=dc.dtype) for info in plan.elim_types + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + dc_obs = _type_slice(dc, info)[combo.kept_index] + et_i = combo.elim_type_index + wt_dc[et_i] = wt_dc[et_i] + jax.ops.segment_sum( + jnp.einsum("mte,mt->me", cross, dc_obs), + combo.elim_index, + num_segments=plan.elim_types[et_i].count, + ) + + delta = jnp.zeros(plan.tangent_dim, dtype=dc.dtype) + for info in plan.kept_types: + delta = delta.at[info.orig_start : info.orig_start + info.count * info.dim].set( + dc[info.start : info.start + info.count * info.dim] + ) + for et_i, info in enumerate(plan.elim_types): + dl = jnp.einsum("nef,nf->ne", vinv[et_i], factors.b_elim[et_i] - wt_dc[et_i]) + delta = delta.at[info.orig_start : info.orig_start + info.count * info.dim].set( + dl.flatten() + ) + return delta + + +def solve_schur_dense(factors: SchurFactors, lambd: jax.Array | float) -> jax.Array: + """Direct reduced solve: form S densely and Cholesky-factor it. Produces + the exact damped Newton step (identical to a full dense solve in float64; + in float32 the deliberate ~2e-3 Tikhonov floor in `_solve_spd_scaled` + makes the two paths diverge measurably).""" + vinv = _damped_vinv(factors, lambd) + b_red = _reduced_rhs(factors, vinv) + S = _assemble_dense_S(factors, lambd, vinv) + dc = _solve_spd_scaled(S, b_red) + return _back_substitute(factors, vinv, dc) + + +def solve_schur_cholmod(factors: SchurFactors, lambd: jax.Array | float) -> jax.Array: + """Direct reduced solve via CHOLMOD: assemble the reduced system S as a + sparse symmetric matrix and factor it with a fill-reducing ordering, then + back-substitute the eliminated variables. This is the Ceres/g2o-style + "Schur + sparse-direct-on-reduced-system" combination: the cheap, + block-diagonal landmark elimination is done on-device, and only the small, + irregular camera system goes to CHOLMOD. + + Like the full-system CHOLMOD path, S is regularized by lambd + 1e-5 (the + extra 1e-5 is folded into the diagonal so CHOLMOD's `beta` is not needed).""" + from ._solvers import _cholmod_solve_symmetric + + plan = factors.plan + pattern = plan.sparse_s_pattern + assert pattern is not None, ( + "solve_schur_cholmod requires plan.sparse_s_pattern; build the plan " + "with the sparse path enabled." + ) + # Regularize by lambd + 1e-5 everywhere (V, the kept diagonal, and the RHS + # via V), matching the full-system CHOLMOD path which damps ATA + (lambd + + # 1e-5) I. Damping V and the kept block by the same amount keeps the Schur + # complement equal to the Schur complement of that regularized full system. + lam_eff = jnp.asarray(lambd) + 1e-5 + vinv = _damped_vinv(factors, lam_eff) + b_red = _reduced_rhs(factors, vinv) + s_values = _assemble_S_values(factors, lam_eff, vinv) + dc = _cholmod_solve_symmetric( + s_values, pattern.rows, pattern.cols, plan.reduced_dim, b_red + ) + return _back_substitute(factors, vinv, dc) + + +def solve_schur_cg( + factors: SchurFactors, + lambd: jax.Array | float, + cg_config: ConjugateGradientConfig, + prev_state: _ConjugateGradientState, +) -> tuple[jax.Array, _ConjugateGradientState]: + """Matrix-free reduced solve: CG on S without forming it. Each product + S x = (H_cc + lambda I) x - W (V + lambda I)^{-1} (W^T x) costs one pass + over the observations; the reduced system's conditioning keeps the + iteration count low.""" + from ._solvers import _ConjugateGradientState + + plan = factors.plan + # Match the regularization of the full-system CG path. + lam_eff = lambd + 1e-5 + vinv = _damped_vinv(factors, lam_eff) + b_red = _reduced_rhs(factors, vinv) + + # Preconditioner on the (approximate) block diagonal of S, honoring the + # same `preconditioner` config choices as the full-system CG path. + if cg_config.preconditioner is None: + precondition = lambda x: x # noqa: E731 + else: + _, wvwt_diag = _wvwt_terms(factors, vinv) + diag_blocks = [ + factors.keep_diag[kt_i] + + (lam_eff + 1e-6) * jnp.eye(info.dim, dtype=b_red.dtype) + - wvwt_diag[kt_i] + for kt_i, info in enumerate(plan.kept_types) + ] + if cg_config.preconditioner == "block_jacobi": + precond_inv = [jnp.linalg.inv(block) for block in diag_blocks] + + def precondition(x: jax.Array) -> jax.Array: + parts = list[jax.Array]() + for kt_i, info in enumerate(plan.kept_types): + xs = _type_slice(x, info) + parts.append( + jnp.einsum("nij,nj->ni", precond_inv[kt_i], xs).flatten() + ) + return jnp.concatenate(parts) + elif cg_config.preconditioner == "point_jacobi": + inv_diag = [ + 1.0 / jnp.diagonal(block, axis1=1, axis2=2) for block in diag_blocks + ] + + def precondition(x: jax.Array) -> jax.Array: + parts = list[jax.Array]() + for kt_i, info in enumerate(plan.kept_types): + parts.append((_type_slice(x, info) * inv_diag[kt_i]).flatten()) + return jnp.concatenate(parts) + else: + assert_never(cg_config.preconditioner) + + def matvec(x: jax.Array) -> jax.Array: + out = lam_eff * x + # H_cc x, group by group through residual space. + for g, slots in enumerate(plan.group_slots): + keep_jac = factors.keep_jacs[g] + if keep_jac is None: + continue + kept_slots = [s for s in slots if s.is_kept] + gathered = list[jax.Array]() + for slot in kept_slots: + info = plan.kept_types[slot.type_index] + gathered.append(_type_slice(x, info)[slot.index]) + x_cat = ( + gathered[0] if len(gathered) == 1 else jnp.concatenate(gathered, axis=1) + ) + residual_space = jnp.einsum("crk,ck->cr", keep_jac, x_cat) + back = jnp.einsum("crk,cr->ck", keep_jac, residual_space) + col = 0 + for slot in kept_slots: + info = plan.kept_types[slot.type_index] + contrib = jax.ops.segment_sum( + back[:, col : col + info.dim], + slot.index, + num_segments=info.count, + ) + out = out.at[info.start : info.start + info.count * info.dim].add( + contrib.flatten() + ) + col += info.dim + # - W (V + lambda I)^{-1} W^T x. + wt_x = [ + jnp.zeros((info.count, info.dim), dtype=x.dtype) for info in plan.elim_types + ] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + xs = _type_slice(x, info)[combo.kept_index] + et_i = combo.elim_type_index + wt_x[et_i] = wt_x[et_i] + jax.ops.segment_sum( + jnp.einsum("mte,mt->me", cross, xs), + combo.elim_index, + num_segments=plan.elim_types[et_i].count, + ) + z = [jnp.einsum("nef,nf->ne", vinv_e, wt_e) for vinv_e, wt_e in zip(vinv, wt_x)] + for combo, cross in zip(plan.combos, factors.cross_blocks): + info = plan.kept_types[combo.kept_type_index] + back = jnp.einsum( + "mte,me->mt", cross, z[combo.elim_type_index][combo.elim_index] + ) + contrib = jax.ops.segment_sum( + back, combo.kept_index, num_segments=info.count + ) + out = out.at[info.start : info.start + info.count * info.dim].add( + -contrib.flatten() + ) + return out + + # Eisenstat-Walker forcing term, matching the full-system CG path. + b_norm = jnp.linalg.norm(b_red) + eta = jnp.minimum( + cg_config.eisenstat_walker_gamma + * (b_norm / (prev_state.ATb_norm_prev + 1e-7)) + ** cg_config.eisenstat_walker_alpha, + cg_config.tolerance_max, + ) + eta = jnp.maximum(cg_config.tolerance_min, jnp.minimum(eta, prev_state.eta)) + + dc, _ = jax.scipy.sparse.linalg.cg( + A=matvec, + b=b_red, + x0=jnp.zeros_like(b_red), + maxiter=plan.reduced_dim, + tol=cast(float, eta), + M=precondition, + ) + delta = _back_substitute(factors, vinv, cast(jax.Array, dc)) + return delta, _ConjugateGradientState(ATb_norm_prev=b_norm, eta=eta) diff --git a/src/jaxls/_solvers.py b/src/jaxls/_solvers.py index efd282a..8024a92 100644 --- a/src/jaxls/_solvers.py +++ b/src/jaxls/_solvers.py @@ -1,9 +1,14 @@ from __future__ import annotations +import contextlib +import dataclasses +import functools +import time from typing import ( TYPE_CHECKING, Callable, Hashable, + Iterator, Literal, assert_never, cast, @@ -13,6 +18,7 @@ import jax import jax.flatten_util import jax_dataclasses as jdc +import numpy as onp import scipy import scipy.sparse from jax import numpy as jnp @@ -30,6 +36,14 @@ update_al_state, update_problem_al_params, ) +from ._schur import ( + EliminationPlan, + SchurFactors, + prepare_schur, + solve_schur_cg, + solve_schur_cholmod, + solve_schur_dense, +) from ._sparse_matrices import BlockRowSparseMatrix, SparseCooMatrix, SparseCsrMatrix from ._variables import VarTypeOrdering, VarValues from .utils import jax_log @@ -94,6 +108,165 @@ def _cholmod_solve_on_host( return cost.solve_A(ATb) +_cholmod_symmetric_analyze_cache: dict[Hashable, sksparse.cholmod.Factor] = {} + + +def _cholmod_solve_symmetric( + s_values: jax.Array, + rows: jax.Array, + cols: jax.Array, + reduced_dim: int, + b: jax.Array, +) -> jax.Array: + """JIT-friendly CHOLMOD solve of a symmetric system S x = b, where S is + given in COO form (s_values at (rows, cols), duplicates summed). Used by + the Schur + CHOLMOD reduced solve.""" + return jax.pure_callback( + functools.partial(_cholmod_solve_symmetric_on_host, reduced_dim=reduced_dim), + b, # Result shape/dtype. + s_values, + rows, + cols, + b, + vmap_method="sequential", + ) + + +def _cholmod_solve_symmetric_on_host( + s_values: jax.Array, + rows: jax.Array, + cols: jax.Array, + b: jax.Array, + *, + reduced_dim: int, +) -> jax.Array: + """Factor a symmetric sparse S (COO: s_values at (rows, cols)) with CHOLMOD + and solve S x = b. Runs on the host.""" + import sksparse.cholmod + + rows_onp = onp.asarray(rows) + cols_onp = onp.asarray(cols) + S = scipy.sparse.coo_matrix( + (onp.asarray(s_values), (rows_onp, cols_onp)), + shape=(reduced_dim, reduced_dim), + ).tocsc() # Sums duplicate (i, j) entries. + + # Cache the symbolic analysis. The COO coordinates are fixed across solves + # for a given problem, so the summed CSC pattern is stable. + cache_key = (rows_onp.tobytes(), cols_onp.tobytes(), reduced_dim) + factor = _cholmod_symmetric_analyze_cache.get(cache_key, None) + if factor is None: + factor = sksparse.cholmod.analyze(S) + _cholmod_symmetric_analyze_cache[cache_key] = factor + + max_cache_size = 512 + if len(_cholmod_symmetric_analyze_cache) > max_cache_size: + _cholmod_symmetric_analyze_cache.pop( + next(iter(_cholmod_symmetric_analyze_cache)) + ) + + # Numeric refactorization reusing the cached symbolic analysis. The 1e-5 + # diagonal floor is already folded into s_values by the caller. + factor = factor.cholesky(S) + return factor.solve_A(onp.asarray(b)) + + +_active_iteration_time_recorder: list[float] | None = None + + +@contextlib.contextmanager +def record_iteration_times() -> Iterator[list[float]]: + """Record a host wall-clock timestamp at each outer LM iteration of any + `solve()` running in this block, for cost-vs-time benchmarking. + + Returns a list that is filled with `time.perf_counter()` values (one per + outer iteration, in order); `times[i] - times[0]` is the elapsed time to + reach iteration `i`. Off the normal solve path entirely — the timestamps + never enter the jitted computation or `SolveSummary`, so there is zero + overhead and no dtype dependence when not recording, and the values are + always host float64 (the in-array alternative would be float32 without + jax_enable_x64, too coarse to resolve millisecond steps). + + Timestamps include async-dispatch latency, so read differences, not + absolute values, and wrap the solve in `block_until_ready` for precise + totals. Not reentrant / not for concurrent solves from multiple threads + (one active recorder at a time). + + with jaxls.record_iteration_times() as times: + sol = problem.solve(init) + # times[1:] - times[0] -> per-iteration elapsed seconds + + The recording callback is decided at trace time, so whether timestamps are + captured is baked into the compiled solve. Entering and leaving this block + therefore clears the solver's JIT cache, so the first solve inside it + recompiles *with* the callback (even if an identical solve was already + compiled without one) and the first solve after it recompiles *without*. + The compile inside the block also acts as the warmup; for steady-state + timing, run the solve twice in the block and read the second: + + with jaxls.record_iteration_times() as times: + jax.block_until_ready(problem.solve(init)) # warmup, compiles + times.clear() + sol = jax.block_until_ready(problem.solve(init)) + """ + global _active_iteration_time_recorder + prev = _active_iteration_time_recorder + times: list[float] = [] + _active_iteration_time_recorder = times + # The callback's presence is fixed at trace time and not part of the JIT + # cache key, so a cached (callback-free) executable would otherwise be + # reused and silently record nothing. Clear the cache on the way in (force + # a callback-bearing recompile) and on the way out (restore the zero- + # overhead executable for normal solves). + _clear_solve_cache() + try: + yield times + finally: + _active_iteration_time_recorder = prev + _clear_solve_cache() + + +def _clear_solve_cache() -> None: + # NonlinearSolver.solve is wrapped by @jdc.jit, whose .clear_cache() the + # type checker can't see through the FunctionType it infers. + NonlinearSolver.solve.clear_cache() # type: ignore[attr-defined] + + +def _record_iteration_time(anchor: jax.Array) -> None: + """Append `time.perf_counter()` to the active recorder, if any, via + `jax.debug.callback` so it works inside the jitted `lax.while_loop`. + + `anchor` is a value that depends on the just-computed iterate (its cost); + the callback consumes it so the timestamp is ordered after that + iteration's work and the callbacks are not merged/hoisted by XLA. A no-op + (no callback emitted) when no `record_iteration_times()` block is active, + so the default solve path is untouched.""" + if _active_iteration_time_recorder is None: + return + + def _stamp(_anchor: object) -> None: + if _active_iteration_time_recorder is not None: + _active_iteration_time_recorder.append(time.perf_counter()) + + jax.debug.callback(_stamp, anchor) + + +def _compute_jacobian_scaler(column_norms: jax.Array) -> jax.Array: + """Column scaling for the Levenberg-Marquardt normal equations. + + This is the historical jaxls scaler: bounded gain in (1, 2], so weak + directions are never amplified and lambdas tuned against it keep their + meaning. A Marquardt-style scale-invariant splice + (``min(legacy, 1.5/n)``) was evaluated and rejected: it silently + re-denominates every tuned lambda for problems whose column norms + exceed unit (IK, pose graphs), regressing downstream workloads, and its + bundle-adjustment motivation disappeared once rejected-step termination + handling was fixed and lambda escalation made adaptive. Measurements and + derivation: benchmarks/results.md, "The column-scaling math". + """ + return 1.0 / (1.0 + column_norms) + 1.0 + + @jdc.pytree_dataclass class _ConjugateGradientState: """State used for Eisenstat-Walker criterion in ConjugateGradientLinearSolver.""" @@ -213,6 +386,14 @@ class _LmInnerState: sol_proposed: _SolutionState local_delta: jax.Array summary: SolveSummary + lambda_growth: jax.Array = dataclasses.field(default_factory=lambda: jnp.array(2.0)) + """Lambda escalation factor; doubled after each consecutive rejection + within one outer step, so a grossly under-damped lambda recovers in + O(log log) tries instead of O(log). This is the rejection-side half of + the strategy in H.B. Nielsen (1999), "Damping Parameter in Marquardt's + Method", IMM-REP-1999-05, DTU; the decrease-on-acceptance rule remains + jaxls's plain halving. A no-op whenever the first proposal is accepted, + so well-tuned solves are unaffected.""" @jdc.pytree_dataclass @@ -225,6 +406,11 @@ class _LmOuterState: lambd: float | jax.Array jacobian_scaler: jax.Array al_state: AugmentedLagrangianState | None # AL state when constraints present. + last_step_accepted: jax.Array = dataclasses.field( + default_factory=lambda: jnp.array(True) + ) + """Whether the last outer step found an acceptable proposal. When False + with lambda at its maximum, no further progress is possible.""" @jdc.pytree_dataclass @@ -241,6 +427,22 @@ class NonlinearSolver: verbose: jdc.Static[bool] augmented_lagrangian: AugmentedLagrangianConfig | None = None """Configuration for Augmented Lagrangian method. Set when constraints are present.""" + elimination: EliminationPlan | None = None + """Schur-complement variable elimination plan. Set when `eliminate` is + passed to `solve()`.""" + + def _resolve_cg_config(self) -> ConjugateGradientConfig: + """CG settings for this solve. A `ConjugateGradientConfig` passed to + `solve(linear_solver=...)` reaches this class as the string + "conjugate_gradient" plus the separate `conjugate_gradient_config` + field (see `AnalyzedLeastSquaresProblem.solve`), so both fields must + be consulted; reading only `linear_solver` silently drops user + settings.""" + if isinstance(self.linear_solver, ConjugateGradientConfig): + return self.linear_solver + if self.conjugate_gradient_config is not None: + return self.conjugate_gradient_config + return ConjugateGradientConfig() @overload def solve( @@ -281,6 +483,9 @@ def solve( lambda_history = jnp.zeros(self.termination.max_iterations) if self.trust_region is not None: lambda_history = lambda_history.at[0].set(self.trust_region.lambda_initial) + # Timestamp the initial point (no-op unless a record_iteration_times() + # block is active); kept off the solver state entirely. + _record_iteration_time(cost_info.cost_total) # Initialize AL state if constraints are present. al_state: AugmentedLagrangianState | None = None @@ -297,6 +502,16 @@ def solve( cost_info = problem._compute_cost_info(vals) cost_history = cost_history.at[0].set(cost_info.cost_nonconstraint) + # Jacobian column scaler: depends only on the linearization at the + # initial point, so compute it once here rather than every outer step + # (compute_column_norms is a full-Jacobian scatter — wasted on every + # iteration past the first when the scaler is frozen anyway). + jacobian_scaler = _compute_jacobian_scaler( + problem._compute_jac_values( + vals, cost_info.jac_cache + ).compute_column_norms() + ) + state = _LmOuterState( solution=_SolutionState( vals=vals, @@ -305,11 +520,7 @@ def solve( if self.linear_solver != "conjugate_gradient" else _ConjugateGradientState( ATb_norm_prev=0.0, - eta=( - ConjugateGradientConfig() - if self.conjugate_gradient_config is None - else self.conjugate_gradient_config - ).tolerance_max, + eta=self._resolve_cg_config().tolerance_max, ), ), summary=SolveSummary( @@ -322,7 +533,7 @@ def solve( lambd=self.trust_region.lambda_initial if self.trust_region is not None else 0.0, - jacobian_scaler=jnp.ones(problem._tangent_dim), + jacobian_scaler=jacobian_scaler, al_state=al_state, ) @@ -334,6 +545,17 @@ def should_continue(state: _LmOuterState) -> jax.Array: basic_checks = ~jnp.isnan(state.solution.cost_info.cost_total) & ( state.summary.iterations < self.termination.max_iterations ) + # No acceptable step exists even at maximum damping: further + # outer steps would re-reject the same proposals until the + # iteration cap. Unconstrained only: under augmented + # Lagrangian, a penalty update reshapes the cost surface and + # can restore progress, so a maxed-out lambda sweep must not + # end the solve before the AL machinery gets that chance. + if self.trust_region is not None and self.augmented_lagrangian is None: + basic_checks = basic_checks & ~( + ~state.last_step_accepted + & (state.lambd >= self.trust_region.lambda_max) + ) # For unconstrained: stop when termination criteria met or step rejected. if self.augmented_lagrangian is None: @@ -397,6 +619,7 @@ def lm_inner_step( A_multiply: Callable[[jax.Array], jax.Array], AT_multiply: Callable[[jax.Array], jax.Array], ATb: jax.Array, + schur_factors: SchurFactors | None = None, ) -> _LmInnerState: """Levenberg-Marquardt inner step. Tries one lambda value.""" # Compute lambda for this iteration. @@ -404,7 +627,7 @@ def lm_inner_step( # - For GN: keep lambda at 0 if self.trust_region is not None: lambd = jnp.minimum( - inner_state.lambd * self.trust_region.lambda_factor, + inner_state.lambd * inner_state.lambda_growth, self.trust_region.lambda_max, ) else: @@ -416,23 +639,37 @@ def lm_inner_step( # Solve the linear system. cg_state: _ConjugateGradientState | None = None - if ( + if schur_factors is not None: + # Variable elimination: solve the reduced (Schur complement) + # system, then back-substitute the eliminated variables. + if self.linear_solver == "conjugate_gradient": + assert isinstance(sol_prev.cg_state, _ConjugateGradientState) + local_delta, cg_state = solve_schur_cg( + schur_factors, lambd, self._resolve_cg_config(), sol_prev.cg_state + ) + elif self.linear_solver == "dense_cholesky": + local_delta = solve_schur_dense(schur_factors, lambd) + elif self.linear_solver == "cholmod": + local_delta = solve_schur_cholmod(schur_factors, lambd) + else: + raise AssertionError( + f"Unexpected elimination plan for {self.linear_solver}." + ) + elif ( isinstance(self.linear_solver, ConjugateGradientConfig) or self.linear_solver == "conjugate_gradient" ): - # Use default CG config if specified as a string, otherwise use the provided config. - cg_config = ( - ConjugateGradientConfig() - if self.linear_solver == "conjugate_gradient" - else self.linear_solver - ) + cg_config = self._resolve_cg_config() assert isinstance(sol_prev.cg_state, _ConjugateGradientState) local_delta, cg_state = cg_config._solve( problem, A_blocksparse, # We could also use (lambd * ATA_diagonals * vec) for # scale-invariant damping. But this is hard to match with CHOLMOD. - # Add 1e-5 regularization to match dense_cholesky behavior. + # The extra 1e-5 keeps simple/underdetermined problems from blowing + # up, matching the CHOLMOD path (see beta=lambd + 1e-5 above); the + # Schur CG path in solve_schur_cg adds the same term. Only + # dense_cholesky omits it, using lambd alone. lambda vec: AT_multiply(A_multiply(vec)) + (lambd + 1e-5) * vec, ATb=ATb, prev_linear_state=sol_prev.cg_state, @@ -472,14 +709,22 @@ def lm_inner_step( if self.trust_region is None: accepted = jnp.array(True) else: - cost_predicted = jnp.sum( - ( - A_blocksparse.multiply(scaled_local_delta) - + sol_prev.cost_info.residual_vector - ) - ** 2 + # Predicted reduction of the linear model, expanded analytically: + # |r|^2 - |r + J dx|^2 = -2 dx^T J^T r - |J dx|^2 + # = 2 dx^T ATb - |J dx|^2. + # The expanded form avoids the catastrophic cancellation of + # subtracting two O(cost) sums over ~1e5+ residual terms, which + # otherwise reduces `step_quality` to rounding noise once the + # true per-step reduction approaches eps * cost. + # + # `A_blocksparse` is the column-scaled Jacobian and `local_delta` + # / `ATb` live in the same scaled coordinates, so the products + # are exactly J dx and dx^T J^T r in the original space. + # (Multiplying by `scaled_local_delta` here would apply the + # column scaling twice.) + predicted_reduction = 2.0 * jnp.dot(local_delta, ATb) - jnp.sum( + A_blocksparse.multiply(local_delta) ** 2 ) - predicted_reduction = sol_prev.cost_info.cost_total - cost_predicted actual_reduction = ( sol_prev.cost_info.cost_total - proposed_cost_info.cost_total ) @@ -496,9 +741,12 @@ def lm_inner_step( tangent_ordering=problem._tangent_ordering, ATb=ATb, iterations=iterations, + accepted=accepted, ) with jdc.copy_and_mutate(inner_state) as next: next.lambd = lambd + # Accelerate the next escalation step (see `lambda_growth`). + next.lambda_growth = inner_state.lambda_growth * 2.0 next.accepted = accepted next.sol_proposed = _SolutionState( vals=proposed_vals, @@ -506,6 +754,10 @@ def lm_inner_step( cg_state=cg_state, ) next.local_delta = local_delta + # Timestamp this iterate (no-op unless a record_iteration_times() + # block is active). Anchored on this step's cost so the callback + # can't be reordered/hoisted ahead of the work it measures. + _record_iteration_time(proposed_cost_info.cost_total) next.summary = SolveSummary( iterations=iterations, termination_criteria=term_criteria, @@ -535,14 +787,9 @@ def lm_outer_step( sol_prev.vals, sol_prev.cost_info.jac_cache ) - # Compute Jacobian scaler on first iteration only. We use jnp.where to - # avoid double JIT compilation (one trace for first=True, one for first=False). - with jdc.copy_and_mutate(state, validate=False) as state: - state.jacobian_scaler = jnp.where( - state.summary.iterations == 0, - 1.0 / (1.0 + A_blocksparse.compute_column_norms()) + 1.0, - state.jacobian_scaler, - ) + # The Jacobian scaler is frozen at the initial linearization (computed + # once in `solve`), so just apply it — no per-iteration column-norm + # scatter. A_blocksparse = A_blocksparse.scale_columns(state.jacobian_scaler) # Get flattened version for COO/CSR matrices. @@ -583,6 +830,15 @@ def lm_outer_step( # Compute right-hand side of normal equation. ATb = -AT_multiply(sol_prev.cost_info.residual_vector) + # Variable elimination: precompute the damping-independent parts of + # the Schur-complement system once per outer step. The inner (lambda + # search) loop below only redoes the damping-dependent work. + schur_factors: SchurFactors | None = None + if self.elimination is not None: + schur_factors = prepare_schur( + self.elimination, A_blocksparse, ATb, linear_solver=self.linear_solver + ) + # Inner loop: search for a good lambda value. # # For Levenberg-Marquardt, we try increasing lambdas until the step is @@ -607,6 +863,11 @@ def lm_outer_step( init_inner_state = _LmInnerState( lambd=init_lambd, + lambda_growth=jnp.asarray( + self.trust_region.lambda_factor + if self.trust_region is not None + else 2.0 + ), accepted=jnp.array(False), # Dummy values - will be overwritten by first lm_inner_step call. sol_proposed=sol_prev, @@ -630,6 +891,7 @@ def lm_outer_step( A_multiply, AT_multiply, ATb, + schur_factors, ), init_val=init_inner_state, ) @@ -655,6 +917,7 @@ def lm_outer_step( sol_prev, ) state_next.lambd = lambd_next + state_next.last_step_accepted = inner_state_final.accepted # Debug: log step acceptance and ATb_norm. if self.verbose: @@ -807,17 +1070,20 @@ def _check_convergence( tangent_ordering: VarTypeOrdering, ATb: jax.Array, iterations: jax.Array, + accepted: jax.Array, ) -> tuple[jax.Array, jax.Array]: """Check for convergence!""" # Cost tolerance: use cost_nonconstraint for meaningful convergence check. # For constrained problems, the total cost changes with each AL update, # but cost_nonconstraint (original objective) provides stable convergence. + # Only accepted proposals count: the cost delta of a rejected step says + # nothing about progress, since the step is not taken. cost_reldelta = ( jnp.abs(cost_nonconstraint_updated - sol_prev.cost_info.cost_nonconstraint) / sol_prev.cost_info.cost_nonconstraint ) - converged_cost = cost_reldelta < self.cost_tolerance + converged_cost = (cost_reldelta < self.cost_tolerance) & accepted # Gradient tolerance: infinity norm of the gradient step. flat_vals = jax.flatten_util.ravel_pytree(sol_prev.vals)[0] @@ -835,11 +1101,17 @@ def _check_convergence( False, ) - # Parameter tolerance. + # Parameter tolerance. Like the cost criterion above, only accepted + # proposals count: a rejected step is not taken, so its size says + # nothing about convergence — and treating it as converged exits the + # lambda-escalation loop early, freezing LM into re-proposing the + # same rejected step forever (measurements: benchmarks/results.md). + # When no acceptable step exists at any damping, the lambda_max + # check in the outer loop terminates instead. param_delta = jnp.linalg.norm(jnp.abs(tangent)) / ( jnp.linalg.norm(flat_vals) + self.parameter_tolerance ) - converged_parameters = param_delta < self.parameter_tolerance + converged_parameters = (param_delta < self.parameter_tolerance) & accepted # Check termination flags. We'll terminate if any of the conditions are met. term_flags = jnp.array( diff --git a/src/jaxls/_sparse_matrices.py b/src/jaxls/_sparse_matrices.py index ca37d11..f002c2c 100644 --- a/src/jaxls/_sparse_matrices.py +++ b/src/jaxls/_sparse_matrices.py @@ -60,7 +60,9 @@ def to_dense(self) -> jax.Array: out = jax.lax.dynamic_update_slice( out, update=self.blocks_concat[:, start_concat_col:end_concat_col], - start_indices=(0, start_col), + # Match index dtypes; `0` and `start_col` can otherwise + # disagree (int64 vs int32) when x64 is enabled. + start_indices=(jnp.zeros((), dtype=start_col.dtype), start_col), ) start_concat_col = end_concat_col diff --git a/src/jaxls/utils.py b/src/jaxls/utils.py index 3709f95..30cee61 100644 --- a/src/jaxls/utils.py +++ b/src/jaxls/utils.py @@ -6,8 +6,36 @@ import jax import termcolor +from jax import numpy as jnp from loguru import logger +# Batched products over a tiny contraction axis, written as explicit +# broadcast-multiply-sums rather than `einsum` / `dot_general`. When the +# contraction dimension is tiny (a residual dim of 2, a landmark dim of 3), +# XLA lowers the batched-GEMM form to a kernel that is ~5-30x slower than +# the elementwise form on GPU (and modestly slower on CPU). These products +# dominate Schur-complement assembly and block-Jacobi preconditioner +# construction, so the form matters. Measurements: benchmarks/results.md, +# "Where the GPU time went: batched einsum vs broadcast". + + +def _batched_gram(a: jax.Array, b: jax.Array) -> jax.Array: + """Per-row blocks summed over the *middle* (contraction) axis: + ``a[...,r,i], b[...,r,j] -> out[...,i,j] = sum_r a[...,r,i] b[...,r,j]``.""" + return jnp.sum(a[..., :, :, None] * b[..., :, None, :], axis=-3) + + +def _batched_outer_last(a: jax.Array, b: jax.Array) -> jax.Array: + """Per-row blocks summed over the *last* axis: + ``a[...,t,f], b[...,s,f] -> out[...,t,s] = sum_f a[...,t,f] b[...,s,f]``.""" + return jnp.sum(a[..., :, None, :] * b[..., None, :, :], axis=-1) + + +def _batched_matmul(a: jax.Array, b: jax.Array) -> jax.Array: + """Per-row matrix products ``a[n] @ b[n]``: + ``a[...,t,e], b[...,e,f] -> out[...,t,f] = sum_e a[...,t,e] b[...,e,f]``.""" + return jnp.sum(a[..., :, :, None] * b[..., None, :, :], axis=-2) + @contextlib.contextmanager def stopwatch(label: str = "unlabeled block") -> Generator[None, None, None]: diff --git a/tests/test_schur.py b/tests/test_schur.py new file mode 100644 index 0000000..20ef101 --- /dev/null +++ b/tests/test_schur.py @@ -0,0 +1,707 @@ +"""Tests for variable elimination (Schur complement). + +Elimination is automatic: `solve()` eliminates dominant block-diagonal +variable types for all three reduced-solve paths (dense Cholesky, CG, and +CHOLMOD sparse-direct on the reduced system). Covers: +- Exactness: the (automatic) reduced solve reproduces the full dense step. +- Multiple kept types, and multiple eliminated types. +- Automatic inference: what gets eliminated, and graceful fallbacks. +- Plan-builder validation of structurally invalid eliminations. +- float32 robustness of the reduced SPD solve. +""" + +import sys +from typing import Literal + +import jax +import jax.numpy as jnp +import numpy as onp +import pytest + +import jaxls + +if sys.version_info >= (3, 12): + from jaxls import _problem, _schur +else: + # On older interpreters the package runs from the transpiled sources. + from jaxls._py310 import _problem, _schur # type: ignore[no-redef] + +jax.config.update("jax_enable_x64", True) + + +class CamVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(9)): ... + + +class BiasVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(2)): ... + + +class PointVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(3)): ... + + +class ColorVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(1)): ... + + +def _rodrigues(aa: jax.Array, x: jax.Array) -> jax.Array: + theta = jnp.linalg.norm(aa) + 1e-12 + k = aa / theta + return ( + x * jnp.cos(theta) + + jnp.cross(k, x) * jnp.sin(theta) + + k * jnp.dot(k, x) * (1.0 - jnp.cos(theta)) + ) + + +def _reproject(cam: jax.Array, point: jax.Array) -> jax.Array: + p = _rodrigues(cam[:3], point) + cam[3:6] + proj = -p[:2] / p[2] + f, k1, k2 = cam[6], cam[7], cam[8] + r2 = jnp.sum(proj**2) + return f * (1.0 + k1 * r2 + k2 * r2**2) * proj + + +def _make_ba_problem( + seed: int = 0, + n_cams: int = 4, + n_pts: int = 12, + with_bias: bool = False, + extra_costs: list[jaxls.Cost] | None = None, + extra_variables: list[jaxls.Var] | None = None, + schur_elimination: Literal["auto", "off"] | tuple[type[jaxls.Var], ...] = "auto", +) -> tuple[jaxls.AnalyzedLeastSquaresProblem, jaxls.VarValues]: + """Small synthetic bundle adjustment problem. Includes a kept-only cost + group (camera priors) and an eliminated-only cost group (point priors). + With `with_bias=True`, each observation also couples a second kept + variable type, exercising multi-kept-slot groups and cross-type pairs in + the reduced system. `extra_costs`/`extra_variables` are appended to the + problem; extra variables are initialized to their type defaults.""" + rng = onp.random.default_rng(seed) + cams_gt = onp.concatenate( + [ + rng.normal(0, 0.1, (n_cams, 3)), + rng.normal(0, 0.5, (n_cams, 2)), + rng.normal(5.0, 0.2, (n_cams, 1)), + onp.full((n_cams, 1), 500.0), + onp.zeros((n_cams, 2)), + ], + axis=1, + ) + pts_gt = rng.normal(0, 1.0, (n_pts, 3)) + + cam_idx_list = [] + pt_idx_list = [] + for j in range(n_pts): + for i in rng.choice(n_cams, size=3, replace=False): + cam_idx_list.append(int(i)) + pt_idx_list.append(j) + cam_idx = onp.array(cam_idx_list) + pt_idx = onp.array(pt_idx_list) + + obs_list = [] + for i, j in zip(cam_idx, pt_idx): + z = _reproject(jnp.array(cams_gt[i]), jnp.array(pts_gt[j])) + obs_list.append(onp.asarray(z) + rng.normal(0, 1.0, 2)) + obs = jnp.array(onp.array(obs_list)) + + costs = list[jaxls.Cost]() + if with_bias: + + def residual_bias( + vals: jaxls.VarValues, + cam: CamVar, + bias: BiasVar, + pt: PointVar, + z: jax.Array, + ) -> jax.Array: + return _reproject(vals[cam], vals[pt]) + vals[bias] - z + + costs.append( + jaxls.Cost( + residual_bias, + ( + CamVar(jnp.array(cam_idx)), + BiasVar(jnp.array(cam_idx)), + PointVar(jnp.array(pt_idx)), + obs, + ), + ) + ) + costs.append( + jaxls.Cost( + lambda vals, bias: 1e-2 * vals[bias], + (BiasVar(jnp.arange(n_cams)),), + ) + ) + else: + + def residual( + vals: jaxls.VarValues, cam: CamVar, pt: PointVar, z: jax.Array + ) -> jax.Array: + return _reproject(vals[cam], vals[pt]) - z + + costs.append( + jaxls.Cost( + residual, + (CamVar(jnp.array(cam_idx)), PointVar(jnp.array(pt_idx)), obs), + ) + ) + # Kept-only cost group. + costs.append( + jaxls.Cost( + lambda vals, cam, target: 1e-3 * (vals[cam] - target), + (CamVar(jnp.arange(n_cams)), jnp.array(cams_gt)), + ) + ) + # Eliminated-only cost group. + costs.append( + jaxls.Cost( + lambda vals, pt, target: 1e-3 * (vals[pt] - target), + (PointVar(jnp.arange(n_pts)), jnp.array(pts_gt)), + ) + ) + + if extra_costs is not None: + costs.extend(extra_costs) + + variables: list[jaxls.Var] = [ + CamVar(jnp.arange(n_cams)), + PointVar(jnp.arange(n_pts)), + ] + if with_bias: + variables.append(BiasVar(jnp.arange(n_cams))) + if extra_variables is not None: + variables.extend(extra_variables) + problem = jaxls.LeastSquaresProblem(costs, variables).analyze( + schur_elimination=schur_elimination + ) + + init_list: list = [ + CamVar(jnp.arange(n_cams)).with_value( + jnp.array(cams_gt + rng.normal(0, 0.02, cams_gt.shape)) + ), + PointVar(jnp.arange(n_pts)).with_value( + jnp.array(pts_gt + rng.normal(0, 0.05, pts_gt.shape)) + ), + ] + if with_bias: + init_list.append( + BiasVar(jnp.arange(n_cams)).with_value( + jnp.array(rng.normal(0, 0.1, (n_cams, 2))) + ) + ) + if extra_variables is not None: + init_list.extend(extra_variables) + return problem, jaxls.VarValues.make(init_list) + + +def _solve_kwargs(max_iterations: int = 10) -> dict: + return dict( + termination=jaxls.TerminationConfig( + max_iterations=max_iterations, early_termination=False + ), + verbose=False, + return_summary=True, + ) + + +def _solve_without_elimination( + linear_solver: str, + max_iterations: int = 10, + **fixture_kwargs, +) -> tuple[jaxls.VarValues, jaxls.SolveSummary]: + """Reference solve on the full system: rebuilds the (deterministic) + fixture with `analyze(schur_elimination="off")` and solves it.""" + problem, init = _make_ba_problem(schur_elimination="off", **fixture_kwargs) + assert problem._elimination is None + return problem.solve( + init, + linear_solver=linear_solver, # type: ignore[arg-type] + **_solve_kwargs(max_iterations), + ) + + +def _cost_history(summary: jaxls.SolveSummary, n: int) -> onp.ndarray: + return onp.asarray(summary.cost_history[:n]) + + +def test_schur_dense_exactness() -> None: + """dense_cholesky automatically eliminates the landmark block, and the + reduced solve must trace the same LM path as the full dense solve.""" + problem, init = _make_ba_problem() + assert _schur.infer_eliminate(problem) == (PointVar,) + _, summary_full = _solve_without_elimination("dense_cholesky") + _, summary_schur = problem.solve( + init, linear_solver="dense_cholesky", **_solve_kwargs() + ) + c_full = _cost_history(summary_full, 10) + c_schur = _cost_history(summary_schur, 10) + rel = onp.abs(c_full - c_schur) / onp.abs(c_full) + assert rel.max() < 1e-6, f"max relative cost difference {rel.max()}" + + +def test_schur_dense_exactness_multiple_kept_types() -> None: + """Exactness with two kept types per observation (cross-type blocks in + the reduced system). Compared over the pre-convergence horizon; once LM + starts rejecting steps near the optimum, last-bit differences can flip an + accept decision and the trajectories separate (the single-step test below + is the sharp check).""" + problem, init = _make_ba_problem(with_bias=True, n_pts=18) + assert _schur.infer_eliminate(problem) == (PointVar,) + _, summary_full = _solve_without_elimination( + "dense_cholesky", max_iterations=6, with_bias=True, n_pts=18 + ) + _, summary_schur = problem.solve( + init, linear_solver="dense_cholesky", **_solve_kwargs(6) + ) + c_full = _cost_history(summary_full, 6) + c_schur = _cost_history(summary_schur, 6) + rel = onp.abs(c_full - c_schur) / onp.abs(c_full) + assert rel.max() < 1e-6, f"max relative cost difference {rel.max()}" + + +def test_schur_single_step_exactness() -> None: + """The damped reduced step equals the explicit full dense solve, for a + sweep of damping values, with multiple kept types in play.""" + problem, init = _make_ba_problem(with_bias=True) + cost_info = problem._compute_cost_info(init) + A = problem._compute_jac_values(init, cost_info.jac_cache) + A_dense = A.to_dense() + ATb = -A_dense.T @ cost_info.residual_vector + + plan = _schur.build_elimination_plan(problem, (PointVar,)) + factors = _schur.prepare_schur(plan, A, ATb, linear_solver="dense_cholesky") + for lambd in (1e-4, 1e-2, 1.0): + ATA = A_dense.T @ A_dense + lambd * jnp.eye(A_dense.shape[1]) + ref = jnp.linalg.solve(ATA, ATb) + step = _schur.solve_schur_dense(factors, lambd) + rel = float(jnp.linalg.norm(step - ref) / jnp.linalg.norm(ref)) + assert rel < 1e-8, f"lambda={lambd}: relative step difference {rel}" + + +try: + import sksparse.cholmod # noqa: F401 + + _HAS_CHOLMOD = True +except Exception: + _HAS_CHOLMOD = False + +requires_cholmod = pytest.mark.skipif( + not _HAS_CHOLMOD, reason="sksparse.cholmod not available" +) + + +def test_sparse_s_equals_dense_s() -> None: + """The sparse COO assembly of S (summed) reproduces the dense S, for a + sweep of damping values. Pure assembly check, no CHOLMOD needed.""" + import scipy.sparse + + problem, init = _make_ba_problem(with_bias=True) + cost_info = problem._compute_cost_info(init) + A = problem._compute_jac_values(init, cost_info.jac_cache) + ATb = -A.to_dense().T @ cost_info.residual_vector + + plan = _schur.build_elimination_plan(problem, (PointVar,)) + assert plan.sparse_s_pattern is not None + factors_dense = _schur.prepare_schur(plan, A, ATb, linear_solver="dense_cholesky") + factors_sparse = _schur.prepare_schur(plan, A, ATb, linear_solver="cholmod") + for lambd in (1e-4, 1e-2, 1.0): + vinv = _schur._damped_vinv(factors_dense, lambd) + S_dense = _schur._assemble_dense_S(factors_dense, lambd, vinv) + + vinv_s = _schur._damped_vinv(factors_sparse, lambd) + s_values = _schur._assemble_S_values(factors_sparse, lambd, vinv_s) + pat = plan.sparse_s_pattern + S_sparse = scipy.sparse.coo_matrix( + (onp.asarray(s_values), (onp.asarray(pat.rows), onp.asarray(pat.cols))), + shape=(plan.reduced_dim, plan.reduced_dim), + ).toarray() + + rel = float( + onp.linalg.norm(S_sparse - onp.asarray(S_dense)) + / onp.linalg.norm(onp.asarray(S_dense)) + ) + assert rel < 1e-10, f"lambda={lambd}: sparse vs dense S relative diff {rel}" + + +@requires_cholmod +def test_schur_cholmod_single_step_exactness() -> None: + """The CHOLMOD reduced step equals the explicit full dense solve. The + CHOLMOD path regularizes by lambd + 1e-5, so the reference matches.""" + problem, init = _make_ba_problem(with_bias=True) + cost_info = problem._compute_cost_info(init) + A = problem._compute_jac_values(init, cost_info.jac_cache) + A_dense = A.to_dense() + ATb = -A_dense.T @ cost_info.residual_vector + + plan = _schur.build_elimination_plan(problem, (PointVar,)) + factors = _schur.prepare_schur(plan, A, ATb, linear_solver="cholmod") + for lambd in (1e-4, 1e-2, 1.0): + ATA = A_dense.T @ A_dense + (lambd + 1e-5) * jnp.eye(A_dense.shape[1]) + ref = jnp.linalg.solve(ATA, ATb) + step = _schur.solve_schur_cholmod(factors, lambd) + rel = float(jnp.linalg.norm(step - ref) / jnp.linalg.norm(ref)) + assert rel < 1e-8, f"lambda={lambd}: relative step difference {rel}" + + +@requires_cholmod +def test_schur_cholmod_converges() -> None: + """End-to-end: Schur + CHOLMOD reaches the same optimum as Schur + dense + over a full LM trajectory.""" + problem, init = _make_ba_problem() + _, summary_dense = problem.solve( + init, linear_solver="dense_cholesky", **_solve_kwargs(20) + ) + _, summary_cholmod = problem.solve( + init, linear_solver="cholmod", **_solve_kwargs(20) + ) + cost_dense = _cost_history(summary_dense, 20)[-1] + cost_cholmod = _cost_history(summary_cholmod, 20)[-1] + rel = abs(cost_cholmod - cost_dense) / (abs(cost_dense) + 1e-12) + assert rel < 1e-5, f"final cost mismatch: dense={cost_dense} cholmod={cost_cholmod}" + + +def test_schur_dense_exactness_multiple_elim_types() -> None: + """Two eliminated types (in different costs) at once.""" + n_pts = 12 + rng = onp.random.default_rng(1) + + # Per-point color variables, observed independently (a cost may not + # couple two eliminated variables, so colors get their own costs). + color_obs = jnp.array(rng.normal(0.5, 0.1, (n_pts, 1))) + color_costs = [ + jaxls.Cost( + lambda vals, color, z: vals[color] - z, + (ColorVar(jnp.arange(n_pts)), color_obs), + ) + ] + problem, init = _make_ba_problem( + extra_costs=color_costs, + extra_variables=[ColorVar(jnp.arange(n_pts))], + ) + assert _schur.infer_eliminate(problem) == (PointVar, ColorVar) + _, summary_full = _solve_without_elimination( + "dense_cholesky", + extra_costs=color_costs, + extra_variables=[ColorVar(jnp.arange(n_pts))], + ) + _, summary_schur = problem.solve( + init, linear_solver="dense_cholesky", **_solve_kwargs() + ) + c_full = _cost_history(summary_full, 10) + c_schur = _cost_history(summary_schur, 10) + rel = onp.abs(c_full - c_schur) / onp.abs(c_full) + assert rel.max() < 1e-6, f"max relative cost difference {rel.max()}" + + +def test_schur_cg_converges() -> None: + """The matrix-free reduced CG path reaches the same optimum. Compared on + the cost of the returned (accepted) solutions, not the proposal + history.""" + problem, init = _make_ba_problem() + # CG trails the exact dense steps by a few outer iterations (its + # Eisenstat-Walker tolerance tightens over the lambda trajectory); 20 + # iterations is enough for both to reach the optimum. + sol_dense, _ = problem.solve( + init, linear_solver="dense_cholesky", **_solve_kwargs(20) + ) + sol_cg, _ = problem.solve( + init, linear_solver="conjugate_gradient", **_solve_kwargs(20) + ) + final_dense = float(onp.asarray(problem._compute_cost_info(sol_dense).cost_total)) + final_cg = float(onp.asarray(problem._compute_cost_info(sol_cg).cost_total)) + assert abs(final_cg - final_dense) / final_dense < 1e-3 + + +def test_schur_cg_preconditioner_options() -> None: + """The Schur CG path must honor `ConjugateGradientConfig.preconditioner` + and converge with every option.""" + problem, init = _make_ba_problem() + assert _schur.infer_eliminate(problem) == (PointVar,) + # 20 iterations: see test_schur_cg_converges. + sol_ref, _ = problem.solve( + init, linear_solver="dense_cholesky", **_solve_kwargs(20) + ) + ref = float(onp.asarray(problem._compute_cost_info(sol_ref).cost_total)) + # Unpreconditioned CG's convergence rate is bound by the raw conditioning + # of the reduced system, so it lags the preconditioned options at any + # fixed iteration budget; it gets a looser parity bound. + preconditioners: tuple[ + tuple[Literal["block_jacobi", "point_jacobi"] | None, float], ... + ] = ( + ("block_jacobi", 1e-2), + ("point_jacobi", 1e-2), + (None, 1e-1), + ) + for preconditioner, tol in preconditioners: + sol, _ = problem.solve( + init, + linear_solver=jaxls.ConjugateGradientConfig(preconditioner=preconditioner), + **_solve_kwargs(20), + ) + cost = float(onp.asarray(problem._compute_cost_info(sol).cost_total)) + assert abs(cost - ref) / ref < tol, f"preconditioner={preconditioner}" + + +def test_no_elimination_when_ineligible() -> None: + """When the dominant type self-couples (pose-graph style), solves run on + the full system without elimination.""" + costs = [ + jaxls.Cost( + lambda vals, c0, c1: vals[c0] - vals[c1], + (CamVar(jnp.arange(3)), CamVar(jnp.arange(1, 4))), + ), + jaxls.Cost( + lambda vals, c0: vals[c0] - 1.0, + (CamVar(jnp.array([0])),), + ), + ] + problem = jaxls.LeastSquaresProblem(costs, [CamVar(jnp.arange(4))]).analyze() + assert _schur.infer_eliminate(problem) == () + _, summary = problem.solve(None, linear_solver="dense_cholesky", **_solve_kwargs(5)) + costs_history = _cost_history(summary, 5) + assert costs_history[-1] < costs_history[0] + + +def test_cholmod_receives_plan(monkeypatch: pytest.MonkeyPatch) -> None: + """cholmod now solves the reduced (Schur) system sparse-directly, so the + precomputed elimination plan IS handed to the solver.""" + problem, init = _make_ba_problem() + assert problem._elimination is not None # Plan was prebuilt by analyze(). + captured = {} + real_solver = _problem.NonlinearSolver + + def spy_solver(*args, **kwargs): + solver = real_solver(*args, **kwargs) + captured["elimination"] = solver.elimination + return solver + + monkeypatch.setattr(_problem, "NonlinearSolver", spy_solver) + try: + problem.solve(init, linear_solver="cholmod", **_solve_kwargs(2)) + except Exception: + # CHOLMOD itself may be unavailable in this environment; the solver + # is constructed (and the plan selected) before factorization. + pass + assert captured["elimination"] is not None, "cholmod did not receive the plan" + + +def test_solver_receives_prebuilt_plan(monkeypatch: pytest.MonkeyPatch) -> None: + """By default the solver receives the plan that analyze() prebuilt.""" + problem, init = _make_ba_problem() + assert problem._elimination is not None # Plan was prebuilt by analyze(). + captured = {} + real_solver = _problem.NonlinearSolver + + def spy_solver(*args, **kwargs): + solver = real_solver(*args, **kwargs) + captured["elimination"] = solver.elimination + return solver + + monkeypatch.setattr(_problem, "NonlinearSolver", spy_solver) + problem.solve(init, linear_solver="dense_cholesky", **_solve_kwargs(2)) + assert captured["elimination"] is not None + + +def test_analyze_flag_skips_plan() -> None: + """analyze(schur_elimination="off") must not build a plan, and solves of + the resulting problem run on the full system.""" + costs = [ + jaxls.Cost( + lambda vals, cam, pt: vals[cam][:3] - vals[pt], + (CamVar(jnp.array([0, 0])), PointVar(jnp.arange(2))), + ) + ] + variables: list[jaxls.Var] = [CamVar(jnp.array([0])), PointVar(jnp.arange(2))] + with_plan = jaxls.LeastSquaresProblem(costs, variables).analyze() + without_plan = jaxls.LeastSquaresProblem(costs, variables).analyze( + schur_elimination="off" + ) + assert with_plan._elimination is not None + assert without_plan._elimination is None + without_plan.solve(None, linear_solver="dense_cholesky", verbose=False) + + +def test_analyze_explicit_eliminate_tuple() -> None: + """A tuple of variable types eliminates exactly those types, bypassing + inference. Eliminating PointVar here matches what "auto" would pick, and the + resulting solve still converges.""" + problem, init = _make_ba_problem(schur_elimination=(PointVar,)) + assert problem._elimination is not None + assert [t.count for t in problem._elimination.elim_types] # non-empty + # The eliminated type is PointVar; cameras are kept. + elim_dim = problem._tangent_dim - problem._elimination.reduced_dim + assert elim_dim == 12 * PointVar.tangent_dim # n_pts default = 12 + problem.solve(init, linear_solver="dense_cholesky", **_solve_kwargs(5)) + + +def test_analyze_explicit_eliminate_rejects_bad_type() -> None: + """An explicit tuple naming a non-block-diagonal type is rejected at + analyze() time (build_elimination_plan raises). A cost couples two + PointVars, so the point block is not block-diagonal; CamVar is kept.""" + costs = [ + jaxls.Cost( + lambda vals, p0, p1: vals[p0] - vals[p1], + (PointVar(jnp.array([0])), PointVar(jnp.array([1]))), + ), + jaxls.Cost( + lambda vals, cam, pt: vals[cam][:3] - vals[pt], + (CamVar(jnp.array([0])), PointVar(jnp.array([0]))), + ), + ] + variables: list[jaxls.Var] = [CamVar(jnp.array([0])), PointVar(jnp.arange(2))] + with pytest.raises(ValueError, match="couples multiple eliminated"): + jaxls.LeastSquaresProblem(costs, variables).analyze( + schur_elimination=(PointVar,) + ) + + +def test_analyze_rejects_invalid_schur_arg() -> None: + """schur_elimination must be 'auto', 'off', or a tuple of types.""" + costs = [ + jaxls.Cost( + lambda vals, cam, pt: vals[cam][:3] - vals[pt], + (CamVar(jnp.array([0])), PointVar(jnp.array([0]))), + ) + ] + prob = jaxls.LeastSquaresProblem( + costs, [CamVar(jnp.array([0])), PointVar(jnp.arange(1))] + ) + with pytest.raises(ValueError, match="must be 'auto', 'off'"): + prob.analyze(schur_elimination="sometimes") # type: ignore[arg-type] + + +def test_plan_rejects_coupled_same_type() -> None: + """A cost coupling two eliminated variables of the same type must be + rejected by the plan builder (automatic inference never selects such a + set; this guards the internal API).""" + costs = [ + jaxls.Cost( + lambda vals, p0, p1: vals[p0] - vals[p1], + (PointVar(jnp.array([0])), PointVar(jnp.array([1]))), + ), + jaxls.Cost( + lambda vals, cam, pt: vals[cam][:3] - vals[pt], + (CamVar(jnp.array([0])), PointVar(jnp.array([0]))), + ), + ] + analyzed = jaxls.LeastSquaresProblem( + costs, [CamVar(jnp.array([0])), PointVar(jnp.arange(2))] + ).analyze() + # Automatic inference avoids the self-coupled point block (it falls back + # to the camera block, which is eligible here). + assert PointVar not in _schur.infer_eliminate(analyzed) + with pytest.raises(ValueError, match="couples multiple eliminated"): + _schur.build_elimination_plan(analyzed, (PointVar,)) + + +def test_plan_rejects_coupled_different_types() -> None: + """A cost coupling two different eliminated types must be rejected.""" + costs = [ + jaxls.Cost( + lambda vals, pt, color: vals[pt][:1] - vals[color], + (PointVar(jnp.array([0])), ColorVar(jnp.array([0]))), + ), + jaxls.Cost( + lambda vals, cam, pt: vals[cam][:3] - vals[pt], + (CamVar(jnp.array([0])), PointVar(jnp.array([0]))), + ), + ] + analyzed = jaxls.LeastSquaresProblem( + costs, + [CamVar(jnp.array([0])), PointVar(jnp.array([0])), ColorVar(jnp.array([0]))], + ).analyze() + with pytest.raises(ValueError, match="couples multiple eliminated"): + _schur.build_elimination_plan(analyzed, (PointVar, ColorVar)) + + +def test_plan_rejects_absent_type() -> None: + """Eliminating a type with no variables in the problem must raise.""" + problem, _ = _make_ba_problem() + with pytest.raises(ValueError, match="no variables"): + _schur.build_elimination_plan(problem, (ColorVar,)) + + +def test_plan_rejects_eliminating_everything() -> None: + """Eliminating every variable type must raise.""" + problem, _ = _make_ba_problem() + with pytest.raises(ValueError, match="at least one"): + _schur.build_elimination_plan(problem, (CamVar, PointVar)) + + +def test_traced_problem_uses_prebuilt_plan() -> None: + """The elimination plan is built at analyze() time, so solving a traced + problem inside jax.jit works (the plan's index arrays are ordinary + pytree leaves) instead of raising.""" + problem, init = _make_ba_problem() + + @jax.jit + def solve_traced_problem( + problem_traced: jaxls.AnalyzedLeastSquaresProblem, + init_vals: jaxls.VarValues, + ) -> jaxls.VarValues: + return problem_traced.solve( + init_vals, + linear_solver="dense_cholesky", + termination=jaxls.TerminationConfig( + max_iterations=5, early_termination=False + ), + verbose=False, + ) + + sol = solve_traced_problem(problem, init) + cost_before = float(onp.asarray(problem._compute_cost_info(init).cost_total)) + cost_after = float(onp.asarray(problem._compute_cost_info(sol).cost_total)) + assert cost_after < cost_before + + +def test_solve_spd_scaled_float64_exact() -> None: + """The robust SPD solve must not perturb well-conditioned float64 solves.""" + rng = onp.random.default_rng(0) + A = rng.normal(size=(20, 20)) + S = jnp.array(A @ A.T + 20 * onp.eye(20)) + x_gt = jnp.array(rng.normal(size=20)) + b = S @ x_gt + x = _schur._solve_spd_scaled(S, b) + assert float(jnp.linalg.norm(x - x_gt) / jnp.linalg.norm(x_gt)) < 1e-10 + + +def test_solve_spd_scaled_float32_robust() -> None: + """A numerically indefinite (catastrophically cancelled) float32 reduced + system must produce finite output, and a well-conditioned float32 system + must stay accurate.""" + rng = onp.random.default_rng(0) + # Near-singular: a rank-deficient PSD matrix plus tiny noise that makes it + # slightly indefinite in float32, as catastrophic cancellation does. + u = rng.normal(size=(20, 3)) + S_bad = (u @ u.T).astype(onp.float32) + S_bad += rng.normal(0, 1e-6, S_bad.shape).astype(onp.float32) + S_bad = ((S_bad + S_bad.T) / 2).astype(onp.float32) + b = rng.normal(size=20).astype(onp.float32) + x = _schur._solve_spd_scaled(jnp.array(S_bad), jnp.array(b)) + assert bool(jnp.all(jnp.isfinite(x))), "robust SPD solve produced non-finite values" + + # Well-conditioned float32: still accurate to ~1e-3. + A = rng.normal(size=(20, 20)).astype(onp.float32) + S_good = jnp.array(A @ A.T + 20 * onp.eye(20, dtype=onp.float32)) + x_gt = jnp.array(rng.normal(size=20).astype(onp.float32)) + b_good = S_good @ x_gt + x_good = _schur._solve_spd_scaled(S_good, b_good) + rel = float(jnp.linalg.norm(x_good - x_gt) / jnp.linalg.norm(x_gt)) + # The float32 Tikhonov floor (~2e-3 relative) bounds the attainable + # accuracy; this is the robustness/accuracy trade-off, by design. + assert rel < 1e-2 + + +def test_float32_end_to_end_no_nans() -> None: + """Full float32 solve through the (automatic) direct Schur path: zero + NaNs.""" + jax.config.update("jax_enable_x64", False) + try: + problem, init = _make_ba_problem() + _, summary = problem.solve( + init, linear_solver="dense_cholesky", **_solve_kwargs() + ) + costs = onp.asarray(summary.cost_history[:10]) + assert not onp.any(onp.isnan(costs)), f"NaNs in float32 cost history: {costs}" + assert costs[-1] < costs[0] + finally: + jax.config.update("jax_enable_x64", True) diff --git a/uv.lock b/uv.lock index f9ebaef..ce62b60 100644 --- a/uv.lock +++ b/uv.lock @@ -3299,28 +3299,27 @@ wheels = [ [[package]] name = "ruff" -version = "0.14.6" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/52/f0/62b5a1a723fe183650109407fa56abb433b00aa1c0b9ba555f9c4efec2c6/ruff-0.14.6.tar.gz", hash = "sha256:6f0c742ca6a7783a736b867a263b9a7a80a45ce9bee391eeda296895f1b4e1cc", size = 5669501, upload-time = "2025-11-21T14:26:17.903Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/67/d2/7dd544116d107fffb24a0064d41a5d2ed1c9d6372d142f9ba108c8e39207/ruff-0.14.6-py3-none-linux_armv6l.whl", hash = "sha256:d724ac2f1c240dbd01a2ae98db5d1d9a5e1d9e96eba999d1c48e30062df578a3", size = 13326119, upload-time = "2025-11-21T14:25:24.2Z" }, - { url = "https://files.pythonhosted.org/packages/36/6a/ad66d0a3315d6327ed6b01f759d83df3c4d5f86c30462121024361137b6a/ruff-0.14.6-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9f7539ea257aa4d07b7ce87aed580e485c40143f2473ff2f2b75aee003186004", size = 13526007, upload-time = "2025-11-21T14:25:26.906Z" }, - { url = "https://files.pythonhosted.org/packages/a3/9d/dae6db96df28e0a15dea8e986ee393af70fc97fd57669808728080529c37/ruff-0.14.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:7f6007e55b90a2a7e93083ba48a9f23c3158c433591c33ee2e99a49b889c6332", size = 12676572, upload-time = "2025-11-21T14:25:29.826Z" }, - { url = "https://files.pythonhosted.org/packages/76/a4/f319e87759949062cfee1b26245048e92e2acce900ad3a909285f9db1859/ruff-0.14.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0a8e7b9d73d8728b68f632aa8e824ef041d068d231d8dbc7808532d3629a6bef", size = 13140745, upload-time = "2025-11-21T14:25:32.788Z" }, - { url = "https://files.pythonhosted.org/packages/95/d3/248c1efc71a0a8ed4e8e10b4b2266845d7dfc7a0ab64354afe049eaa1310/ruff-0.14.6-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d50d45d4553a3ebcbd33e7c5e0fe6ca4aafd9a9122492de357205c2c48f00775", size = 13076486, upload-time = "2025-11-21T14:25:35.601Z" }, - { url = "https://files.pythonhosted.org/packages/a5/19/b68d4563fe50eba4b8c92aa842149bb56dd24d198389c0ed12e7faff4f7d/ruff-0.14.6-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:118548dd121f8a21bfa8ab2c5b80e5b4aed67ead4b7567790962554f38e598ce", size = 13727563, upload-time = "2025-11-21T14:25:38.514Z" }, - { url = "https://files.pythonhosted.org/packages/47/ac/943169436832d4b0e867235abbdb57ce3a82367b47e0280fa7b4eabb7593/ruff-0.14.6-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:57256efafbfefcb8748df9d1d766062f62b20150691021f8ab79e2d919f7c11f", size = 15199755, upload-time = "2025-11-21T14:25:41.516Z" }, - { url = "https://files.pythonhosted.org/packages/c9/b9/288bb2399860a36d4bb0541cb66cce3c0f4156aaff009dc8499be0c24bf2/ruff-0.14.6-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ff18134841e5c68f8e5df1999a64429a02d5549036b394fafbe410f886e1989d", size = 14850608, upload-time = "2025-11-21T14:25:44.428Z" }, - { url = "https://files.pythonhosted.org/packages/ee/b1/a0d549dd4364e240f37e7d2907e97ee80587480d98c7799d2d8dc7a2f605/ruff-0.14.6-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:29c4b7ec1e66a105d5c27bd57fa93203637d66a26d10ca9809dc7fc18ec58440", size = 14118754, upload-time = "2025-11-21T14:25:47.214Z" }, - { url = "https://files.pythonhosted.org/packages/13/ac/9b9fe63716af8bdfddfacd0882bc1586f29985d3b988b3c62ddce2e202c3/ruff-0.14.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:167843a6f78680746d7e226f255d920aeed5e4ad9c03258094a2d49d3028b105", size = 13949214, upload-time = "2025-11-21T14:25:50.002Z" }, - { url = "https://files.pythonhosted.org/packages/12/27/4dad6c6a77fede9560b7df6802b1b697e97e49ceabe1f12baf3ea20862e9/ruff-0.14.6-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:16a33af621c9c523b1ae006b1b99b159bf5ac7e4b1f20b85b2572455018e0821", size = 14106112, upload-time = "2025-11-21T14:25:52.841Z" }, - { url = "https://files.pythonhosted.org/packages/6a/db/23e322d7177873eaedea59a7932ca5084ec5b7e20cb30f341ab594130a71/ruff-0.14.6-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:1432ab6e1ae2dc565a7eea707d3b03a0c234ef401482a6f1621bc1f427c2ff55", size = 13035010, upload-time = "2025-11-21T14:25:55.536Z" }, - { url = "https://files.pythonhosted.org/packages/a8/9c/20e21d4d69dbb35e6a1df7691e02f363423658a20a2afacf2a2c011800dc/ruff-0.14.6-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:4c55cfbbe7abb61eb914bfd20683d14cdfb38a6d56c6c66efa55ec6570ee4e71", size = 13054082, upload-time = "2025-11-21T14:25:58.625Z" }, - { url = "https://files.pythonhosted.org/packages/66/25/906ee6a0464c3125c8d673c589771a974965c2be1a1e28b5c3b96cb6ef88/ruff-0.14.6-py3-none-musllinux_1_2_i686.whl", hash = "sha256:efea3c0f21901a685fff4befda6d61a1bf4cb43de16da87e8226a281d614350b", size = 13303354, upload-time = "2025-11-21T14:26:01.816Z" }, - { url = "https://files.pythonhosted.org/packages/4c/58/60577569e198d56922b7ead07b465f559002b7b11d53f40937e95067ca1c/ruff-0.14.6-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:344d97172576d75dc6afc0e9243376dbe1668559c72de1864439c4fc95f78185", size = 14054487, upload-time = "2025-11-21T14:26:05.058Z" }, - { url = "https://files.pythonhosted.org/packages/67/0b/8e4e0639e4cc12547f41cb771b0b44ec8225b6b6a93393176d75fe6f7d40/ruff-0.14.6-py3-none-win32.whl", hash = "sha256:00169c0c8b85396516fdd9ce3446c7ca20c2a8f90a77aa945ba6b8f2bfe99e85", size = 13013361, upload-time = "2025-11-21T14:26:08.152Z" }, - { url = "https://files.pythonhosted.org/packages/fb/02/82240553b77fd1341f80ebb3eaae43ba011c7a91b4224a9f317d8e6591af/ruff-0.14.6-py3-none-win_amd64.whl", hash = "sha256:390e6480c5e3659f8a4c8d6a0373027820419ac14fa0d2713bd8e6c3e125b8b9", size = 14432087, upload-time = "2025-11-21T14:26:10.891Z" }, - { url = "https://files.pythonhosted.org/packages/a5/1f/93f9b0fad9470e4c829a5bb678da4012f0c710d09331b860ee555216f4ea/ruff-0.14.6-py3-none-win_arm64.whl", hash = "sha256:d43c81fbeae52cfa8728d8766bbf46ee4298c888072105815b392da70ca836b2", size = 13520930, upload-time = "2025-11-21T14:26:13.951Z" }, +version = "0.15.17" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/a9/3abdf488f1bf3d24c699415e454ed554a6350d5d89ce183be1ee0a3361ac/ruff-0.15.17.tar.gz", hash = "sha256:2ec446937fd16c8c4de2674a209cc5af64d9c6f17d21fbf1151054fa0bcf5219", size = 4743346, upload-time = "2026-06-11T17:54:47.663Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/4d/e11259f5da07cb6afb2d074c31bf09da9671993f7329d4f15d2fdc458301/ruff-0.15.17-py3-none-linux_armv6l.whl", hash = "sha256:d9feddb927fc68bd295f5eebc587a7e42cfaf9b65f60ca4a2386febff575da8f", size = 10856677, upload-time = "2026-06-11T17:54:49.533Z" }, + { url = "https://files.pythonhosted.org/packages/29/3e/772d679e1a0dc058e58875bd2c0cb713a0530877b4a76fee3c7966df0d49/ruff-0.15.17-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:25805a226d741c47d274a35ad5c10a7dde175fcddfa511d7cf3da0a21eb3eab7", size = 11223443, upload-time = "2026-06-11T17:55:00.573Z" }, + { url = "https://files.pythonhosted.org/packages/68/58/bd41f7688b2fd5623012605130ed70e60aa7f2244baa3d5066bdd61530c8/ruff-0.15.17-py3-none-macosx_11_0_arm64.whl", hash = "sha256:f6ad73b14c2d18a3bf8ad7cb6974294d7f613a7898604826058e6ac64918ef4d", size = 10566458, upload-time = "2026-06-11T17:55:07.52Z" }, + { url = "https://files.pythonhosted.org/packages/d8/5b/733371013fcf1ec339e477ece6ab42bfe10bdd9bba8ee88a9516aa56bfc0/ruff-0.15.17-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ba0c1e4f95bcb3869d0d30cbd5917071ef2e28665abfec970cdab0492c713ed", size = 10914483, upload-time = "2026-06-11T17:55:05.501Z" }, + { url = "https://files.pythonhosted.org/packages/bd/cc/6f24251cc0252f7239391ccb85833f320efad14ebe5b443943f37ced6332/ruff-0.15.17-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:81647960f10bff57d2e51cadd0c3950fe598400c852863a038720ef5b8cca91e", size = 10647497, upload-time = "2026-06-11T17:54:57.733Z" }, + { url = "https://files.pythonhosted.org/packages/68/dd/0d10c17ce1a1624d6fc3156309c3f834fdb5dfaad026ec90c85684f3990e/ruff-0.15.17-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e01a84ddbc8c16c23055ba3924476850f1bbc1917cebbb9376665a63e74260d", size = 11416967, upload-time = "2026-06-11T17:54:51.461Z" }, + { url = "https://files.pythonhosted.org/packages/2f/91/556bfb156f6144f355e831c23db00b2fc4120f86b3ce81cc5f7fd2df51f3/ruff-0.15.17-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:84fe9f653152f8f294f9f7e03bf3a453d8b4a27f7a59c78c8666167f2b17b96c", size = 12335770, upload-time = "2026-06-11T17:54:45.793Z" }, + { url = "https://files.pythonhosted.org/packages/88/82/8b5999aa13355e926f06d9f42a32dcca862f623bf0363785ff89d607dffd/ruff-0.15.17-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8c0fe88a7676e7a05b73174d4d4a59cb2ac21ff8263583f87a81a6018475a978", size = 11575441, upload-time = "2026-06-11T17:54:32.661Z" }, + { url = "https://files.pythonhosted.org/packages/11/93/f10377bb04109ca0e8cbc483ff1982c54b6d418210041776f93e8cdc7fa9/ruff-0.15.17-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ecfc3c7878fff94633ab0348524e093f9ce3243080416dd7d14f8ba400174719", size = 11557614, upload-time = "2026-06-11T17:54:34.698Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a6/eeeae7f7d5493df41649ab3db92f086b2d0a30199e4efdf8e3dd7a033f24/ruff-0.15.17-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:b8461180b22420b1bdc289909410930761629fddf2a5aaf60fae1ab26cedc4c4", size = 11544450, upload-time = "2026-06-11T17:54:39.042Z" }, + { url = "https://files.pythonhosted.org/packages/32/88/5991ce565129a24dd4a00db1254b3b5db2e53018cbe4018ea5a89738e727/ruff-0.15.17-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6eccbe50a038b503e7140b441aa9c7fc8c1f36edf23ebef9f4165c2f28f568b7", size = 10892524, upload-time = "2026-06-11T17:55:09.432Z" }, + { url = "https://files.pythonhosted.org/packages/f5/1d/0fdd248313425f55223968af04b0a42125466a8d88d21c1d99c6af0a51e8/ruff-0.15.17-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:382fc0521025f5a8ad447d8bdd523545d0d7646adb718eb1c2dac5065ec27c0f", size = 10659573, upload-time = "2026-06-11T17:54:36.824Z" }, + { url = "https://files.pythonhosted.org/packages/9e/0e/072e8260deb9461062ce9311ced27a8e541229a6ffd483013dd37661e43e/ruff-0.15.17-py3-none-musllinux_1_2_i686.whl", hash = "sha256:456d41fcd1b2777ad63f09a6e7121d43f7b688bbc76a800c10f7f8fb1f912c3f", size = 11127818, upload-time = "2026-06-11T17:55:03.124Z" }, + { url = "https://files.pythonhosted.org/packages/ab/b4/55060a34163121498014696b5f656db5b8c6963768f227dbf0d76b311073/ruff-0.15.17-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:b1a04bcc94ae6194e9db05d16ad31f298a7194bfbcb08258bbe589cee1d587b8", size = 11655901, upload-time = "2026-06-11T17:54:53.562Z" }, + { url = "https://files.pythonhosted.org/packages/49/71/9b29d6b87cef468d697f43c6a91e3fae4a80185779d7d5a4ef27d173439f/ruff-0.15.17-py3-none-win32.whl", hash = "sha256:596065960ab1ff593f744220c9fe6580eda00a95003cffa9f4048bb5b1bf0392", size = 10925574, upload-time = "2026-06-11T17:54:55.723Z" }, + { url = "https://files.pythonhosted.org/packages/3d/b2/8fc77f3723228836fa5d12497eb71c808f83782e10d058d2b15cfa14640b/ruff-0.15.17-py3-none-win_amd64.whl", hash = "sha256:6769e5fa1710b179b92e0bfa5a51735b35baea9013dadb06d5f44cbcf9547084", size = 12058788, upload-time = "2026-06-11T17:54:41.042Z" }, + { url = "https://files.pythonhosted.org/packages/2d/c7/c53e8dbff9c9dc4b7928773421ae294a5d28fcb8dcda1a089579d3a7e510/ruff-0.15.17-py3-none-win_arm64.whl", hash = "sha256:f3be1fbb34bcdfd146240d8fb92a709d4c2c8191348580a3c044ec60fa0b4456", size = 11355275, upload-time = "2026-06-11T17:54:43.635Z" }, ] [[package]]