From e0df04b384fc2651f97f94e6a65f0a07ec0637b0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 14 Oct 2026 10:26:58 +0100 Subject: [PATCH] chore: retire searches_minimal Nautilus profile + findings (re-homed to autofit_profiling) searches_minimal/nautilus_profile.py and nautilus_bottleneck_findings.md now live in autofit_profiling as the search.fit breakdown cell scripts/gaussian_x1/nautilus/search_fit_breakdown.py and wiki/campaigns/search_fit_breakdown.md (autofit_profiling#3, search-extensibility B4a). Committed deletion; pre-delete commit 06d86ac1deed93d3424791e71b47e6c135c7b5f1. Co-Authored-By: Claude Fable 5.1 --- .../nautilus_bottleneck_findings.md | 114 ------- searches_minimal/nautilus_profile.py | 278 ------------------ 2 files changed, 392 deletions(-) delete mode 100644 searches_minimal/nautilus_bottleneck_findings.md delete mode 100644 searches_minimal/nautilus_profile.py diff --git a/searches_minimal/nautilus_bottleneck_findings.md b/searches_minimal/nautilus_bottleneck_findings.md deleted file mode 100644 index 68eb1ba..0000000 --- a/searches_minimal/nautilus_bottleneck_findings.md +++ /dev/null @@ -1,114 +0,0 @@ -# Nautilus NN-training bottleneck — findings and verdict - -Research task: https://github.com/Jammy2211/autofit_workspace_developer/issues/18 -Scripts: `nautilus_profile.py` (anatomy), `nautilus_sweep.py` (zero-code knobs), -`nautilus_jax_mlp.py` (JAX emulator prototype). nautilus 1.0.5, `n_live=200`, -`seed=1`, laptop CPU (WSL2, 8 cores). Wall times vary ±40 % between sessions -(thermal/contention); every comparison below is within-run. - -## 1. Where the time goes (Phase 1) - -With a fast likelihood, Nautilus wall time is almost entirely sampler overhead, -and it is **not only the neural networks**: - -| component | 3-param Gaussian fit | fast 10-d likelihood (5 µs/call) | -|---|---:|---:| -| NN ensemble training | 51.7 % | 40.2 % | -| bound geometry (ellipsoid unions + 1000-pt volume MC) | 25.2 % | 38.2 % | -| proposal sampling | 5.2 % | 18.0 % | -| likelihood (incl. transform) | 14.3 % | **0.6 %** | - -Mechanics (all in nautilus source): each bound update trains one -`NeuralNetworkEmulator` **per ellipsoid** — `n_networks=4` sklearn MLPs -(hidden layers 100/50/20, `tol=0`, early stop after 10 flat epochs), each -pinned to a single thread (`neural.py`, `threadpool_limits(1)`). Training -parallelises only across the ensemble via the sampler pool (`pool_s`). - -**PyAutoFit wrapper finding** (`autofit/non_linear/search/nest/nautilus/search.py`, -read-only this task): `fit_x1_cpu` — the JAX/GPU path — passes `pool=None`, so -in exactly the fast-GPU-likelihood regime every MLP trains serially. - -## 2. Zero-code knobs (Phase 2) - -All configs preserved quality (log Z within 0.02 nats of baseline both -scenarios, ESS ~10 000, reference log Z −57.5 for the Gaussian fit): - -| config | fast 10-d | 3-param fit | note | -|---|---:|---:|---| -| `n_networks=0` (pure ellipsoid bounds) | **3.2×** | **2.0×** | +1–10 % evals; log Z −57.451 vs −57.455 | -| `n_update=2*n_live` (retrain half as often) | 3.0× | — | quality intact | -| `pool=(None, 4)` (parallel MLP training) | 2.3× | 1.1× | `os.fork()` + JAX threads ⇒ deadlock risk | -| `n_networks=2` | 2.2× | — | | -| small MLP `(32,)` | 2.0× | 1.0× | | -| `vectorized=True` (batched callback) | 2.0× | — | the hook a vmapped GPU likelihood needs | -| `n_networks=0` + `n_update=2x` combined | 2.7× | 1.9× | ≈ `n_networks=0` alone; they don't stack | - -`n_networks=0` is an explicitly supported nautilus mode -(`bounds/neural.py:74`): bounds fall back to the ellipsoid union, i.e. -MultiNest-style region sampling. Its cost is sampling efficiency — here only -+1 % (10-d) to +10–27 % (3-param) more likelihood calls — which is precisely -the cheap currency in the fast-likelihood regime. - -**Caveat that decided against it (see §5):** evals-to-ML on the 3-param fit -degraded 4278 → 6481 (+50 %) — the sampler takes materially longer to *find -the peak* without the NN bounds, even on an easy unimodal problem. Bounds -never bias the result (acceptance is always the exact likelihood), but on -real curved/degenerate lens posteriors at 20–30 parameters the -ellipsoid-only efficiency loss is expected to be much larger than these toy -numbers, plausibly erasing the wall-time win entirely. - -## 3. JAX/GPU MLP swap prototype (Phase 3) — negative on CPU - -`nautilus_jax_mlp.py` replaces the sklearn ensemble with a single jitted JAX -program (vmap over the 4 networks, full-batch Adam, padded shapes so the JIT -cache survives across bounds, jitted predict on the proposal hot path). Same -architecture/normalisation; end-to-end quality matched sklearn. - -Result: **4–6× slower** than sklearn on CPU (fast 10-d: 5.3 s → 24–28 s NN -time; 3-param fit: 37 s → 118 s+). sklearn's mini-batch Adam with early -stopping does far less arithmetic than fixed-step full-batch training, and -CPU JAX gains nothing from the ensemble vmap. A GPU would flip the arithmetic -cost, but dispatch/transfer latency on thousands of tiny fits is exactly -where GPUs disappoint, and §1 caps the achievable win: NN training is ≤ 52 % -of wall, so even free training beats it by at most ~2× — less than -`n_networks=0` already delivers for zero code. **Parked**; revisit only if a -GPU-node measurement of a real lens fit contradicts this. - -## 4. Can nautilus be JAX-ified wholesale? No. - -The sampler's control flow is adaptive and host-side by design: -data-dependent while loops (ellipsoid splitting until a volume target, -rejection sampling until `n_batch` accepted), a dynamically growing point -store, per-bound sklearn objects, HDF5 checkpointing. Porting that to XLA's -bounded-shape world is not a port — it is writing a different sampler, which -is what NSS already is (measured 7.5× faster per eval on MGE in the A100 -campaign, with its own quality/OOM caveats). The long-term answer remains a -JAX-native sampler swap; nothing found here changes that, and nothing here -requires waiting for it. - -## 5. Verdict (human decision 2026-07-09) - -The `n_networks=0` route was **rejected by the maintainer** on review: for -real lens-model sampling the NN bounds are what keep Nautilus's sampling -efficiency high, and the measured +50 % evals-to-ML on even a toy unimodal -fit (§2 caveat) supports the expectation that ellipsoid-only bounds would -hurt peak-finding on real curved/degenerate posteriors. No PyAutoFit config -change and no GPU validation follow-up were filed. - -What stands from this investigation: - -1. **The overhead anatomy (§1) is the reference** for any future Nautilus - speed-up work: NN training 40–52 %, bound geometry 25–38 %, and the two - must be attacked together for a win beyond ~2×. -2. **Do not wire a process pool into the JAX path** (`pool=(None, N)`): - measured win is modest and `os.fork()` under a JAX-threaded process risks - deadlock. -3. **The JAX/GPU MLP swap is parked with a measured negative** (§3); the - burden of proof for reopening it is a GPU-node measurement on a real fit. -4. **The long-term answer remains a JAX-native sampler swap** (§4) — none of - the interim knobs measured here beats staying the course on that. -5. If Nautilus overhead work is ever revisited, the next targets in order: - bound geometry (~0.9 s/bound; profile `Union.compute`/split vs the - 1000-pt volume MC first), then early-stopping-preserving training - parallelism that does not use process forks (threads or an async bound - builder). diff --git a/searches_minimal/nautilus_profile.py b/searches_minimal/nautilus_profile.py deleted file mode 100644 index 04db428..0000000 --- a/searches_minimal/nautilus_profile.py +++ /dev/null @@ -1,278 +0,0 @@ -""" -Nautilus Overhead Profile -------------------------- - -Phase 1 of the nautilus-nn-bottleneck research task -(https://github.com/Jammy2211/autofit_workspace_developer/issues/18). - -Instruments nautilus 1.0.5 to break sampler wall time into its overhead -components, on problems where the likelihood itself is fast — the regime -(e.g. MGE-only source on GPU) where sampler overhead dominates: - - - NN training time inside NeuralNetworkEmulator.train - (n_networks sklearn MLPs per ellipsoid, per bound) - - bound geometry NautilusBound.compute minus NN training - (ellipsoid union fitting / splitting / volume MC) - - proposal sampling Sampler.sample_shell (draw + shell association) - - likelihood Sampler.evaluate_likelihood (transform + callback) - - residue everything else (bookkeeping, shell updates, I/O) - -Two scenarios: - gaussian_1d_data the standard minimal 3-parameter Gaussian-fit problem - (identical data/seed to nautilus_simple.py, so rows - are comparable to output/comparison.txt) - fast_gauss_10d a 10-dim correlated Gaussian log-likelihood costing - ~1 us/call — emulates the fast-GPU-likelihood regime - at a dimensionality typical of lens models - -Requirements: - pip install nautilus-sampler -""" - -import time -from pathlib import Path - -import numpy as np -import autofit as af - -from searches_minimal._metrics import MLTracker - -# -------------------------------------------------------------------------- -# Timing instrumentation (monkey-patched wrappers around nautilus internals) -# -------------------------------------------------------------------------- - -import nautilus -from nautilus import Sampler -from nautilus.neural import NeuralNetworkEmulator -from nautilus.bounds import NautilusBound - - -class Timings: - """Accumulate wall time and call counts per instrumented component.""" - - def __init__(self): - self.totals = {} - self.counts = {} - self.per_call = {} - - def add(self, key, dt): - self.totals[key] = self.totals.get(key, 0.0) + dt - self.counts[key] = self.counts.get(key, 0) + 1 - self.per_call.setdefault(key, []).append(dt) - - -TIMINGS = Timings() - -_train_original = NeuralNetworkEmulator.train.__func__ -_compute_original = NautilusBound.compute.__func__ -_sample_shell_original = Sampler.sample_shell -_evaluate_likelihood_original = Sampler.evaluate_likelihood - - -def _train_timed(cls, *args, **kwargs): - t0 = time.time() - result = _train_original(cls, *args, **kwargs) - TIMINGS.add("nn_train", time.time() - t0) - return result - - -def _compute_timed(cls, *args, **kwargs): - t0 = time.time() - result = _compute_original(cls, *args, **kwargs) - TIMINGS.add("bound_compute", time.time() - t0) - n_ellipsoids = len(result.neural_bounds) - TIMINGS.per_call.setdefault("ellipsoids_per_bound", []).append(n_ellipsoids) - return result - - -def _sample_shell_timed(self, *args, **kwargs): - t0 = time.time() - result = _sample_shell_original(self, *args, **kwargs) - TIMINGS.add("sample_shell", time.time() - t0) - return result - - -def _evaluate_likelihood_timed(self, *args, **kwargs): - t0 = time.time() - result = _evaluate_likelihood_original(self, *args, **kwargs) - TIMINGS.add("evaluate_likelihood", time.time() - t0) - return result - - -NeuralNetworkEmulator.train = classmethod(_train_timed) -NautilusBound.compute = classmethod(_compute_timed) -Sampler.sample_shell = _sample_shell_timed -Sampler.evaluate_likelihood = _evaluate_likelihood_timed - -# -------------------------------------------------------------------------- -# Scenario 1 — the standard minimal problem (3 parameters, comparable rows) -# -------------------------------------------------------------------------- - - -class Gaussian: - def __init__(self, centre=30.0, normalization=1.0, sigma=5.0): - self.centre = centre - self.normalization = normalization - self.sigma = sigma - - def model_data_from(self, xvalues, xp=np): - return xp.multiply( - xp.divide(self.normalization, self.sigma * xp.sqrt(2.0 * xp.pi)), - xp.exp(-0.5 * xp.square(xp.divide(xvalues - self.centre, self.sigma))), - ) - - -np.random.seed(1) - -xvalues = np.arange(100) -true_gaussian = Gaussian(centre=50.0, normalization=25.0, sigma=10.0) -data = true_gaussian.model_data_from(xvalues=xvalues) -noise_map = np.full(data.shape, 0.01) -data += np.random.normal(0.0, 0.01, data.shape) - -model = af.Model(Gaussian) -model.centre = af.UniformPrior(lower_limit=0.0, upper_limit=100.0) -model.normalization = af.UniformPrior(lower_limit=0.0, upper_limit=50.0) -model.sigma = af.UniformPrior(lower_limit=0.0, upper_limit=50.0) - - -def gaussian_prior_transform(cube): - return np.array(model.vector_from_unit_vector(cube)) - - -def gaussian_log_likelihood(params): - instance = model.instance_from_vector(vector=params) - model_data = instance.model_data_from(xvalues=xvalues) - residual_map = data - model_data - chi_squared_map = (residual_map / noise_map) ** 2.0 - return float(-0.5 * np.sum(chi_squared_map)) - - -# -------------------------------------------------------------------------- -# Scenario 2 — fast 10-dim correlated Gaussian (fast-GPU-likelihood regime) -# -------------------------------------------------------------------------- - -N_DIM_FAST = 10 - -rng_setup = np.random.default_rng(1) -_a = rng_setup.normal(size=(N_DIM_FAST, N_DIM_FAST)) -_cov = _a @ _a.T + N_DIM_FAST * np.eye(N_DIM_FAST) -_cov_inv = np.linalg.inv(_cov) -_mean = np.full(N_DIM_FAST, 0.5) - - -def fast_prior_transform(cube): - return cube - - -def fast_log_likelihood(params): - d = params - _mean - return float(-0.5 * d @ _cov_inv @ d) - - -# -------------------------------------------------------------------------- -# Runner -# -------------------------------------------------------------------------- - - -def run_scenario(name, prior_transform, log_likelihood, n_dim, n_live=200): - TIMINGS.totals.clear() - TIMINGS.counts.clear() - TIMINGS.per_call.clear() - - tracker = MLTracker() - n_calls = {"n": 0} - - def wrapped(params): - n_calls["n"] += 1 - log_l = log_likelihood(params) - tracker.record(log_l) - return log_l - - # Time the raw likelihood so the callback overhead is separable. - test_point = prior_transform(np.full(n_dim, 0.5)) - t0 = time.time() - for _ in range(1000): - log_likelihood(test_point) - raw_likelihood_ms = time.time() - t0 - - sampler = Sampler( - prior=prior_transform, - likelihood=wrapped, - n_dim=n_dim, - n_live=n_live, - seed=1, - ) - - t_start = time.time() - sampler.run(verbose=True) - wall = time.time() - t_start - - points, log_w, log_l = sampler.posterior() - max_logl = float(np.max(log_l)) - evals_to_ml, time_to_ml = tracker.finalise(max_log_l=max_logl, tolerance=1.0) - - nn = TIMINGS.totals.get("nn_train", 0.0) - bound = TIMINGS.totals.get("bound_compute", 0.0) - shell = TIMINGS.totals.get("sample_shell", 0.0) - like = TIMINGS.totals.get("evaluate_likelihood", 0.0) - geometry = bound - nn - residue = wall - bound - shell - like - ellipsoids = TIMINGS.per_call.get("ellipsoids_per_bound", []) - n_mlp_fits = TIMINGS.counts.get("nn_train", 0) - - def row(label, seconds, extra=""): - return ( - f"{label:<22s} {seconds:>8.2f} s {100 * seconds / wall:>5.1f} % {extra}" - ) - - summary = f"""\ ---- Nautilus Overhead Profile: {name} --- -n_dim / n_live: {n_dim} / {n_live} -Wall time: {wall:.2f} s -Likelihood evals: {n_calls['n']} -Raw likelihood: {raw_likelihood_ms:.3f} ms/call (timed outside sampler, 1000 calls) - ---- Wall-time breakdown --- -{row("nn_train", nn, f"({n_mlp_fits} emulator trainings, {sum(ellipsoids)} ellipsoid-ensembles x 4 MLPs)")} -{row("bound_geometry", geometry, f"({TIMINGS.counts.get('bound_compute', 0)} bounds; ellipsoids/bound: {ellipsoids})")} -{row("sample_shell", shell, f"({TIMINGS.counts.get('sample_shell', 0)} calls)")} -{row("evaluate_likelihood", like, f"({TIMINGS.counts.get('evaluate_likelihood', 0)} batched calls)")} -{row("residue", residue, "(bookkeeping, shell updates)")} - ---- Quality (MLTracker contract) --- -Max log L: {max_logl:.4f} -Log evidence: {float(sampler.log_z):.4f} -ESS: {float(sampler.n_eff):.1f} -Evals to ML: {evals_to_ml if evals_to_ml is not None else 'n/a'} -Time to ML: {f'{time_to_ml:.2f} s' if time_to_ml is not None else 'n/a'} -""" - print() - print(summary) - return summary - - -if __name__ == "__main__": - summaries = [] - summaries.append( - run_scenario( - "gaussian_1d_data (3 params)", - gaussian_prior_transform, - gaussian_log_likelihood, - n_dim=model.prior_count, - ) - ) - summaries.append( - run_scenario( - f"fast_gauss_{N_DIM_FAST}d", - fast_prior_transform, - fast_log_likelihood, - n_dim=N_DIM_FAST, - ) - ) - - output_dir = Path(__file__).parent / "output" - output_dir.mkdir(parents=True, exist_ok=True) - summary_path = output_dir / f"{Path(__file__).stem}_summary.txt" - summary_path.write_text("\n".join(summaries)) - print(f"Summary written to: {summary_path}")