From a38647cbd233940a6825272076560688ff6e05b5 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 16 May 2026 14:04:38 +0100
Subject: [PATCH] feat: mirror Nautilus profiling scripts and shared search
 helpers (Phase 3)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 3 of the autolens_profiling z_feature roadmap. Stands up
autolens_profiling/searches/ with profiling for the Nautilus sampler
only, mirrored from autolens_workspace_developer/searches_minimal/.
Designs the folder layout so other samplers (Dynesty, Emcee, BlackJAX,
NumPyro, PocoMC, NSS, LBFGS) can slot in cleanly under their own
follow-up prompts.

Files mirrored (4 copies from clean origin/main of _developer):
  searches_minimal/_setup.py           -> searches/_setup.py
  searches_minimal/_metrics.py         -> searches/_metrics.py
  searches_minimal/nautilus_simple.py  -> searches/nautilus/simple.py
  searches_minimal/nautilus_jax.py     -> searches/nautilus/jax.py

Path rewrites in _setup.py:
- _DATASET_SUBPATH:
    Path("jax_profiling") / "dataset" / "imaging" / "hst"
    -> Path("dataset") / "imaging" / "hst"
- Replaced the if al.util.dataset.should_simulate(...) subprocess block
  (which invoked _developer's dataset_setup/imaging.py) with a clean
  FileNotFoundError pointing back at _developer for regeneration —
  matches the Phase 1 pattern.
- Dropped now-unused subprocess + sys imports.
- Docstring example: `from searches_minimal._setup` -> `from searches._setup`
- Cross-reference: `jax_profiling/imaging/mge.py` -> `likelihood/imaging/mge.py`

Import rewrites in nautilus/simple.py and nautilus/jax.py:
- `from searches_minimal._{metrics,setup}` -> `from searches._{metrics,setup}`
- Added sys.path injection at top so scripts work whether run as
  `python searches/nautilus/<name>.py` or `python -m searches.nautilus.<name>`.

Output convention upgrade (both Nautilus scripts):
- Originals wrote a freeform `<script>_summary.txt` to `output/`.
- New: write versioned JSON + bar-chart PNG to
  `results/searches/nautilus/<script>_summary_v<al.__version__>.{json,png}`
  matching Phase 1's convention. JSON carries sampler config, results,
  performance, and convergence sections plus the al.__version__ tag.
  PNG is a horizontal bar chart of headline timings (wall time, time
  per eval, time to ML; plus JIT compile time on jax.py).

READMEs:
- searches/README.md: replaces Phase 0 placeholder. Section overview,
  why bypass af.NonLinearSearch, shared _setup / _metrics roles,
  supported-samplers table (Nautilus ✓, others "planned"), versioned
  artifact convention, run instructions.
- searches/nautilus/README.md: per-sampler narrative grounded in
  _developer/searches_minimal/sweep_findings.md — expected behaviour
  (converged logZ ≈ -169k for this MGE setup), caveats (upstream JAX-
  gradient pathology #514), headline-run-times table placeholder for
  Phase 4 to auto-populate.

Smoke status: ran py_compile across all 4 files (PASS), verified no
residual `searches_minimal` references in the mirrored scripts. Full
runtime smoke (`python searches/nautilus/simple.py` end-to-end) was
intentionally skipped — n_live=200 takes 30+ min on CPU and Phase 5
will add the AUTOLENS_PROFILING_SMOKE=1 short-circuit that makes this
cheap forever. Static checks + matching the Phase 1/2 pattern give
sufficient confidence to ship.

F1 lesson applied: copies came from the worktree's clean origin/main
of _developer (NOT the canonical, which is dirty with ~36 modified
files).

Closes #5.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 searches/README.md          |  71 +++++++++++-
 searches/_metrics.py        |  81 ++++++++++++++
 searches/_setup.py          | 108 ++++++++++++++++++
 searches/nautilus/README.md |  47 ++++++++
 searches/nautilus/jax.py    | 217 ++++++++++++++++++++++++++++++++++++
 searches/nautilus/simple.py | 188 +++++++++++++++++++++++++++++++
 6 files changed, 709 insertions(+), 3 deletions(-)
 create mode 100644 searches/_metrics.py
 create mode 100644 searches/_setup.py
 create mode 100644 searches/nautilus/README.md
 create mode 100644 searches/nautilus/jax.py
 create mode 100644 searches/nautilus/simple.py

diff --git a/searches/README.md b/searches/README.md
index 76a35a5..6ab7d3c 100644
--- a/searches/README.md
+++ b/searches/README.md
@@ -1,7 +1,72 @@
 # searches
 
-Sampler / search profiling — Nautilus first, with the folder layout designed so that other samplers (Dynesty, Emcee, BlackJAX, NumPyro, NSS, LBFGS, PocoMC) can slot in under their own prompts.
+Sampler / search profiling for the PyAutoLens HST MGE lens-modelling likelihood. Each subfolder drives a single sampler family directly against the real likelihood — bypassing `af.NonLinearSearch` — so the per-sampler convergence characteristics (wall time, likelihood evaluations, posterior ESS, evals/time to ML) can be compared on identical footing.
 
-Populated by **Phase 3** of the `autolens_profiling` roadmap. Will mirror `autolens_workspace_developer/searches_minimal/`.
+## Why bypass `af.NonLinearSearch`?
 
-See the top-level [README](../README.md) for the full phase plan.
+`af.NonLinearSearch` adds caching, multi-process forking, output formatting, and result hierarchies that are valuable for production fits but obscure the underlying sampler's cost. The scripts in this section call the sampler library directly and instrument every likelihood evaluation through a shared `MLTracker`. The result is a clean apples-to-apples comparison of:
+
+- Wall time and likelihood-evaluation count to **Nautilus's default convergence** (`n_eff=10000`, `f_live=0.01`).
+- Per-evaluation likelihood cost (NumPy baseline vs JAX-JIT'd path).
+- Evals-to-ML and time-to-ML — the eval index and wall time at which the running max log L first came within 1 nat of the final maximum.
+
+## Shared helpers
+
+| File | Role |
+|------|------|
+| [`_setup.py`](./_setup.py) | Builds the HST imaging dataset, the MGE + Isothermal + ExternalShear lens model with an MGE source bulge, and the `AnalysisImaging` object. The dataset, mask, and model mirror the reference setup in [`likelihood/imaging/mge.py`](../likelihood/imaging/mge.py) so likelihood values are directly comparable across the two sections. |
+| [`_metrics.py`](./_metrics.py) | `MLTracker` — records the log-likelihood and wall time of every evaluation, computes evals-to-ML and time-to-ML headline numbers. Also offers `MLTracker.from_log_l_history` for samplers that JIT their likelihood and only expose log-L per dead/live point post hoc. |
+
+## Supported samplers
+
+| Sampler | Folder | Status | Notes |
+|---------|--------|--------|-------|
+| Nautilus | [`nautilus/`](./nautilus/README.md) | ✓ profiled | Both NumPy (`simple.py`) and JAX-JIT (`jax.py`) variants. |
+| Dynesty | _planned_ | not yet mirrored | Static nested sampling; reference scripts at `autolens_workspace_developer/searches_minimal/dynesty_simple.py`. |
+| Emcee | _planned_ | not yet mirrored | Affine-invariant ensemble MCMC. |
+| BlackJAX (NUTS, SMC) | _planned_ | not yet mirrored | Pure-JAX HMC family. Gradient pathology surfaced in upstream `sweep_findings.md`; HMC viability depends on first fixing NaN-gradient hot spots. |
+| NumPyro (ESS) | _planned_ | not yet mirrored | Ensemble slice sampler under JAX. |
+| PocoMC | _planned_ | not yet mirrored | Preconditioned Monte Carlo. |
+| NSS (simple, jit, grad) | _planned_ | not yet mirrored | Nested slice sampler; `nss_jit.py` shows VRAM ceiling on consumer GPUs (see `sweep_findings.md`). |
+| LBFGS | _planned_ | not yet mirrored | Not a sampler; serves as the maximum-likelihood reference point. |
+
+Each row above corresponds to one or more scripts under `autolens_workspace_developer/searches_minimal/`; the mirror migration here under their own follow-up prompts.
+
+## Versioned artifacts
+
+Each script writes a JSON + PNG pair to:
+
+```
+results/searches/<sampler>/<script>_summary_v<al.__version__>.{json,png}
+```
+
+The JSON carries the structured timings + sampler config + best-fit summary. The PNG is a bar chart of the headline timings (wall time, time per eval, time to ML; plus JIT compile time on JAX scripts).
+
+Old versions are retained alongside new ones; Phase 4's dashboard surfaces the latest per axis.
+
+## Running a script
+
+From the repo root (cwd matters because `_setup.build_dataset()` resolves `dataset/imaging/hst/` relative to the repo root via `Path(__file__).resolve().parent.parent`):
+
+```bash
+cd autolens_profiling
+python searches/nautilus/simple.py
+python searches/nautilus/jax.py
+```
+
+Or as modules:
+
+```bash
+python -m searches.nautilus.simple
+python -m searches.nautilus.jax
+```
+
+Both invocation styles work — each script injects the repo root into `sys.path` before importing `searches._{setup,metrics}` for robustness.
+
+**Requirements:** `nautilus-sampler` for the Nautilus scripts (`pip install nautilus-sampler`). The JAX variant additionally needs a working JAX install.
+
+**Codex / sandboxed runs:**
+
+```bash
+NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python searches/nautilus/simple.py
+```
diff --git a/searches/_metrics.py b/searches/_metrics.py
new file mode 100644
index 0000000..6edb6e4
--- /dev/null
+++ b/searches/_metrics.py
@@ -0,0 +1,81 @@
+"""
+Shared per-evaluation tracker used by every script in this folder.
+
+Wrap a log-likelihood callable with ``MLTracker.wrap`` (or call
+``tracker.record(log_l)`` manually) and the tracker stores the log L and
+wall-clock time of every evaluation. After the run, ``finalise`` returns
+the eval index and wall time at which the running max log L first came
+within ``tolerance`` nats of the final maximum -- the "evals to ML" /
+"time to ML" headline numbers used in the comparison.
+
+For JAX paths where the likelihood runs inside ``jax.jit`` (and a Python
+callback is impossible without forcing a host round-trip), use
+``MLTracker.from_log_l_history`` instead with the full per-eval log L
+sequence reconstructed from the sampler's dead-point + live-point state.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Callable, Optional, Sequence
+
+
+class MLTracker:
+    """Record per-evaluation log L and wall time, compute evals/time to ML."""
+
+    def __init__(self):
+        self.t0 = time.time()
+        self.history_log_l: list[float] = []
+        self.history_wall: list[float] = []
+
+    def record(self, log_l: float) -> None:
+        self.history_log_l.append(float(log_l))
+        self.history_wall.append(time.time() - self.t0)
+
+    def wrap(self, fn: Callable) -> Callable:
+        """Decorate a log-likelihood callable so every call is recorded."""
+
+        def wrapped(*args, **kwargs):
+            log_l = fn(*args, **kwargs)
+            self.record(log_l)
+            return log_l
+
+        return wrapped
+
+    def finalise(
+        self, max_log_l: Optional[float] = None, tolerance: float = 1.0
+    ) -> tuple[Optional[int], Optional[float]]:
+        """Return (evals_to_ml, time_to_ml) — the eval index and wall time at
+        which the running max first came within ``tolerance`` nats of the
+        final maximum. ``(None, None)`` if no evaluations were recorded."""
+        if not self.history_log_l:
+            return None, None
+        if max_log_l is None:
+            max_log_l = max(self.history_log_l)
+        target = max_log_l - tolerance
+        for i, log_l in enumerate(self.history_log_l):
+            if log_l >= target:
+                return i + 1, self.history_wall[i]
+        return None, None
+
+    @staticmethod
+    def from_log_l_history(
+        log_l_history: Sequence[float],
+        total_sampling_time: float,
+        tolerance: float = 1.0,
+    ) -> tuple[Optional[int], Optional[float]]:
+        """Variant for samplers that run their likelihood inside JIT and only
+        expose log L per dead/live point post hoc. ``time_to_ml`` is linearly
+        interpolated from the total sampling time -- evaluations are assumed
+        evenly distributed over the run, which is a reasonable approximation
+        for nested sampling (each step is roughly the same cost)."""
+        if not log_l_history:
+            return None, None
+        max_log_l = max(log_l_history)
+        target = max_log_l - tolerance
+        for i, log_l in enumerate(log_l_history):
+            if log_l >= target:
+                evals_to_ml = i + 1
+                time_to_ml = total_sampling_time * (evals_to_ml / len(log_l_history))
+                return evals_to_ml, time_to_ml
+        return None, None
diff --git a/searches/_setup.py b/searches/_setup.py
new file mode 100644
index 0000000..197f9a6
--- /dev/null
+++ b/searches/_setup.py
@@ -0,0 +1,108 @@
+"""
+Shared setup for the ``searches/`` profiling scripts.
+
+Builds the HST imaging dataset, the MGE + Isothermal + ExternalShear lens model
+with an MGE source bulge, and the ``AnalysisImaging`` object used by every
+sampler in this section. The dataset, mask, and model mirror the reference setup
+in ``likelihood/imaging/mge.py`` so the likelihood value is directly comparable
+across the two sections.
+
+Usage
+-----
+
+    from searches._setup import build_dataset, build_model, build_analysis
+
+    dataset = build_dataset()
+    model = build_model(mask_radius=3.5)
+    analysis = build_analysis(dataset, use_jax=False)
+"""
+
+from pathlib import Path
+
+import numpy as np
+
+import autofit as af
+import autolens as al
+
+_WORKSPACE_ROOT = Path(__file__).resolve().parent.parent
+_DATASET_SUBPATH = Path("dataset") / "imaging" / "hst"
+
+PIXEL_SCALE = 0.05
+MASK_RADIUS = 3.5
+
+
+def build_dataset(mask_radius: float = MASK_RADIUS) -> al.Imaging:
+    """Load the HST imaging dataset with mask + radial-bin over-sampling applied."""
+    dataset_path = _DATASET_SUBPATH
+
+    if al.util.dataset.should_simulate(str(dataset_path)):
+        raise FileNotFoundError(
+            f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+            f"repo mirrors only the curated datasets needed for default smoke "
+            f"runs. To regenerate, use the source-of-truth script at "
+            f"autolens_workspace_developer/jax_profiling/dataset_setup/imaging.py "
+            f"and copy the result into autolens_profiling/dataset/."
+        )
+
+    dataset = al.Imaging.from_fits(
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        pixel_scales=PIXEL_SCALE,
+    )
+
+    mask = al.Mask2D.circular(
+        shape_native=dataset.shape_native,
+        pixel_scales=dataset.pixel_scales,
+        radius=mask_radius,
+    )
+    dataset = dataset.apply_mask(mask=mask)
+    dataset = dataset.apply_over_sampling(over_sample_size_lp=4)
+
+    over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(
+        grid=dataset.grid,
+        sub_size_list=[4, 2, 1],
+        radial_list=[0.3, 0.6],
+        centre_list=[(0.0, 0.0)],
+    )
+    dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)
+    return dataset
+
+
+def build_model(
+    mask_radius: float = MASK_RADIUS, total_gaussians: int = 20
+) -> af.Collection:
+    """Build the lens + source model used in ``jax_profiling/imaging/mge.py``."""
+    lens_bulge = al.model_util.mge_model_from(
+        mask_radius=mask_radius,
+        total_gaussians=total_gaussians,
+        centre_prior_is_uniform=True,
+    )
+    mass = af.Model(al.mp.Isothermal)
+    shear = af.Model(al.mp.ExternalShear)
+    lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear)
+
+    source_bulge = al.model_util.mge_model_from(
+        mask_radius=mask_radius,
+        total_gaussians=total_gaussians,
+        centre_prior_is_uniform=False,
+    )
+    source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)
+
+    return af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+
+def build_analysis(dataset: al.Imaging, use_jax: bool = False) -> al.AnalysisImaging:
+    """Build the analysis object. Set ``use_jax=True`` for the pure-JAX path."""
+    return al.AnalysisImaging(dataset=dataset, use_jax=use_jax)
+
+
+def format_best_fit(instance) -> str:
+    """Terse one-line summary of the lens mass + shear of a best-fit instance."""
+    mass = instance.galaxies.lens.mass
+    shear = instance.galaxies.lens.shear
+    return (
+        f"lens.mass.einstein_radius={mass.einstein_radius:.4f}  "
+        f"lens.mass.centre=({mass.centre[0]:.3f}, {mass.centre[1]:.3f})  "
+        f"shear=({shear.gamma_1:.4f}, {shear.gamma_2:.4f})"
+    )
diff --git a/searches/nautilus/README.md b/searches/nautilus/README.md
new file mode 100644
index 0000000..ad0cfcb
--- /dev/null
+++ b/searches/nautilus/README.md
@@ -0,0 +1,47 @@
+# searches/nautilus
+
+[Nautilus](https://github.com/johannesulf/nautilus) is a production nested-importance sampler that combines neural-network-based importance sampling with traditional nested sampling. It's gradient-free, so it sidesteps the JAX-gradient pathologies that affect HMC-family samplers on this likelihood, and is a strong baseline for end-to-end "what does a real sampler do on this lens model" timing.
+
+These scripts drive Nautilus directly against the HST MGE imaging likelihood, bypassing `af.NonLinearSearch`. They are wiring tests + run-time profiling, **not** converged science fits — `n_live=200` is below what you'd use in production, but enough to see per-evaluation cost and reach the default `n_eff=10000` / `f_live=0.01` termination on this MGE setup.
+
+## Scripts
+
+| Script | Likelihood backend | What it profiles |
+|--------|--------------------|------------------|
+| [`simple.py`](./simple.py) | NumPy (`use_jax=False`) | Baseline: pure-NumPy log-likelihood passed straight to Nautilus. Highest per-evaluation cost; useful as the reference point against which JAX speedup is measured. |
+| [`jax.py`](./jax.py) | JAX JIT (`use_jax=True`, `jax.jit`-compiled likelihood) | Reports JIT compile time separately. Per-evaluation cost is JAX kernel + a single Python ↔ JAX boundary crossing per call (Nautilus itself is NumPy-only). Compare versus a pure-JAX nested sampler like NSS-JIT (not yet mirrored) for the no-boundary-crossing variant. |
+
+Both share the same Nautilus configuration so timings are directly comparable: `n_live=200`, default `n_eff=10000`, default `f_live=0.01`. Both use the shared `_setup` / `_metrics` from one folder up.
+
+## What each script reports
+
+- **Best fit**: max-likelihood lens mass / shear parameters (one-line summary).
+- **Max log L** and **log evidence**.
+- **Wall time** for the sampling phase (excluding JIT compile for `jax.py`).
+- **JIT compile time** (one-shot warmup; `jax.py` only).
+- **Likelihood evaluations** and **time per eval** (ms).
+- **ESS** (effective sample size) and **posterior samples**.
+- **Convergence** indicator (Nautilus's `n_eff` / `f_live` defaults are reached).
+- **Evals to ML** and **time to ML** via the shared `MLTracker`.
+
+The headline JSON+PNG pair is written to `results/searches/nautilus/` per the [section README](../README.md#versioned-artifacts) convention.
+
+## Headline run-times (populated by Phase 4)
+
+| Script | Backend | Wall time | Time / eval | Evals → ML | Time → ML |
+|--------|---------|-----------|-------------|-----------|-----------|
+| `simple.py` | NumPy | _populated_ | _populated_ | _populated_ | _populated_ |
+| `jax.py` | JAX JIT | _populated_ | _populated_ | _populated_ | _populated_ |
+
+Numbers are filled in by Phase 4's `scripts/build_readme.py` from the latest `*_summary_v<version>.json` under `results/searches/nautilus/`.
+
+## Expected behaviour
+
+For reference: prior sweep runs on this exact MGE setup (recorded in `autolens_workspace_developer/searches_minimal/sweep_findings.md`) put converged log-evidence at around **logZ ≈ -169k**. A non-converged early-stop reading of `logZ ≈ -191k` is roughly what you'll see after a few minutes of sampling. The likelihood landscape anneals slowly — fully converged runs at `n_live=100` take ~30–60 minutes on GPU.
+
+The JAX variant's wall time is dominated by the NumPy/JAX boundary crossings, not the JAX kernel. A future NSS-JIT mirror will surface the no-boundary-crossing alternative.
+
+## Caveats
+
+- **`use_jax=True` and JIT compile**: `_setup.build_analysis(dataset, use_jax=True)` returns an analysis object that the Nautilus wrapper feeds via `jax.jit`. If the underlying JAX-jitted likelihood path has an upstream regression (see [PyAutoLens#514](https://github.com/PyAutoLabs/PyAutoLens/issues/514)), the `jax.py` script may produce different log evidence than `simple.py` — that's an upstream issue, not a Nautilus issue.
+- **GPU memory**: on a 6 GB consumer GPU (e.g. RTX 2060), `jax.py` with `n_live=200` fits comfortably for Nautilus (gradient-free, no curvature storage). The same `n_live` causes OOM in NSS-JIT (which stores curvature for HMC-style moves) per the upstream sweep findings.
diff --git a/searches/nautilus/jax.py b/searches/nautilus/jax.py
new file mode 100644
index 0000000..e90a334
--- /dev/null
+++ b/searches/nautilus/jax.py
@@ -0,0 +1,217 @@
+"""
+Minimal Nautilus Example — pure-JAX HST MGE likelihood
+------------------------------------------------------
+
+Drives the Nautilus nested sampler against the HST MGE imaging likelihood
+running fully under ``jax.jit``. The analysis is built with ``use_jax=True``
+and the closure is passed through ``jax.jit`` once, ahead of sampling, so
+the JIT compile cost is reported separately from the sampling wall time.
+
+Nautilus itself is a NumPy sampler, so the wrapper does
+``np.asarray(jit_loglike(jnp.asarray(params)))`` per call -- the JAX kernel
+runs but every evaluation crosses the Python <-> JAX boundary.
+
+``n_live`` is kept at the smoke-test values used by ``simple.py`` — this
+is a wiring test, not a converged posterior.
+
+Compare versus ``simple.py`` (NumPy likelihood under the same sampler) and
+``likelihood/imaging/mge.py`` (single-likelihood JIT profiling of the same
+MGE setup).
+
+Requirements:
+    pip install nautilus-sampler
+"""
+
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autolens as al
+
+# Make ``from searches._{setup,metrics}`` importable regardless of how the
+# script is invoked (``python searches/nautilus/jax.py``, ``python -m
+# searches.nautilus.jax``, or a CI runner).
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from searches._metrics import MLTracker
+from searches._setup import (
+    build_analysis,
+    build_dataset,
+    build_model,
+    format_best_fit,
+)
+
+dataset = build_dataset()
+model = build_model()
+analysis = build_analysis(dataset, use_jax=True)
+
+print(f"Model free parameters: {model.total_free_parameters}")
+
+from nautilus import Sampler
+
+
+def log_likelihood_jax(params):
+    """Pure-JAX log likelihood: flat parameter vector -> scalar log L."""
+    instance = model.instance_from_vector(vector=params, xp=jnp)
+    return analysis.log_likelihood_function(instance=instance)
+
+
+jit_log_likelihood = jax.jit(log_likelihood_jax)
+
+# Warm up the JIT once so the compile cost is measured separately.
+warmup_unit = [0.5] * model.prior_count
+warmup_physical = jnp.asarray(model.vector_from_unit_vector(warmup_unit))
+print("JIT-compiling MGE likelihood (one-shot)...", flush=True)
+t_jit_start = time.time()
+_ = float(jax.block_until_ready(jit_log_likelihood(warmup_physical)))
+t_jit = time.time() - t_jit_start
+print(f"  Compiled in {t_jit:.2f} s", flush=True)
+
+
+def prior_transform(cube):
+    """Map a unit cube to physical parameters via the model's priors."""
+    return np.array(model.vector_from_unit_vector(cube))
+
+
+n_likelihood_calls = 0
+tracker = MLTracker()
+
+
+def log_likelihood(params):
+    """Adapter: NumPy in, JIT'd JAX likelihood, Python float out."""
+    global n_likelihood_calls
+    n_likelihood_calls += 1
+    log_l = float(jit_log_likelihood(jnp.asarray(params)))
+    tracker.record(log_l)
+    return log_l
+
+
+n_live = 200
+
+sampler = Sampler(
+    prior=prior_transform,
+    likelihood=log_likelihood,
+    n_dim=model.prior_count,
+    n_live=n_live,
+)
+
+t_start = time.time()
+# Run to Nautilus's default convergence (n_eff=10000, f_live=0.01) on the
+# JAX-jitted MGE likelihood. JIT compile is paid once above; per-call cost
+# inside sampling is the JAX kernel + Python<->JAX boundary.
+sampler.run(verbose=True)
+t_elapsed = time.time() - t_start
+
+points, log_w, log_l = sampler.posterior()
+best_idx = np.argmax(log_l)
+best_instance = model.instance_from_vector(vector=list(points[best_idx]))
+max_logl = float(np.max(log_l))
+
+evals_to_ml, time_to_ml = tracker.finalise(max_log_l=max_logl, tolerance=1.0)
+
+# ---------------------------------------------------------------------------
+# Print human-readable summary
+# ---------------------------------------------------------------------------
+
+summary = f"""\
+--- Nautilus (JAX JIT) Results ---
+Best fit:        {format_best_fit(best_instance)}
+Max log L:       {max_logl:.4f}
+Log evidence:    {float(sampler.log_z):.4f}
+
+--- Performance ---
+Wall time:           {t_elapsed:.2f} s     (excludes JIT compile, run ahead of time)
+JIT compile time:    {t_jit:.2f} s     (one-shot warm-up before sampling)
+Likelihood evals:    {n_likelihood_calls}
+Time per eval:       {t_elapsed / max(n_likelihood_calls, 1) * 1e3:.3f} ms
+ESS:                 {float(sampler.n_eff):.1f}
+Posterior samples:   {len(points)}
+Sampler config:      n_live={n_live}, default n_eff=10000, f_live=0.01
+
+--- Convergence ---
+Converged:           yes (Nautilus default n_eff / f_live)
+Evals to ML:         {evals_to_ml if evals_to_ml is not None else 'n/a'}     (first eval within 1 nat of max log L)
+Time to ML:          {f'{time_to_ml:.2f} s' if time_to_ml is not None else 'n/a'}
+"""
+
+print()
+print(summary)
+
+# ---------------------------------------------------------------------------
+# Write versioned JSON + PNG to results/searches/nautilus/
+# ---------------------------------------------------------------------------
+
+al_version = al.__version__
+result_dict = {
+    "sampler": "nautilus",
+    "backend": "jax_jit",
+    "instrument": "hst",
+    "model": {
+        "type": "MGE+Isothermal+ExternalShear",
+        "free_parameters": int(model.total_free_parameters),
+    },
+    "sampler_config": {
+        "n_live": n_live,
+        "n_eff_target": 10000,
+        "f_live": 0.01,
+    },
+    "results": {
+        "max_log_likelihood": max_logl,
+        "log_evidence": float(sampler.log_z),
+        "best_fit_summary": format_best_fit(best_instance),
+    },
+    "performance": {
+        "wall_time_s": t_elapsed,
+        "jit_compile_s": t_jit,
+        "likelihood_evals": int(n_likelihood_calls),
+        "time_per_eval_ms": t_elapsed / max(n_likelihood_calls, 1) * 1e3,
+        "ess": float(sampler.n_eff),
+        "posterior_samples": int(len(points)),
+    },
+    "convergence": {
+        "converged": True,
+        "evals_to_ml": int(evals_to_ml) if evals_to_ml is not None else None,
+        "time_to_ml_s": float(time_to_ml) if time_to_ml is not None else None,
+    },
+    "version": al_version,
+}
+
+results_dir = _REPO_ROOT / "results" / "searches" / "nautilus"
+results_dir.mkdir(parents=True, exist_ok=True)
+json_path = results_dir / f"jax_summary_v{al_version}.json"
+json_path.write_text(json.dumps(result_dict, indent=2))
+print(f"  Results JSON saved to: {json_path}")
+
+# Bar chart of the headline timings
+fig, ax = plt.subplots(figsize=(8, 3))
+labels = [
+    "jit_compile (s)",
+    "wall_time (s)",
+    "time_per_eval (ms)",
+    "time_to_ml (s)" if time_to_ml is not None else "time_to_ml (n/a)",
+]
+times = [
+    t_jit,
+    t_elapsed,
+    t_elapsed / max(n_likelihood_calls, 1) * 1e3,
+    float(time_to_ml) if time_to_ml is not None else 0.0,
+]
+ax.barh(labels, times, color=["#8172B2", "#4C72B0", "#55A868", "#C44E52"])
+ax.set_title(f"Nautilus (JAX JIT) — HST MGE — v{al_version}")
+fig.tight_layout()
+png_path = results_dir / f"jax_summary_v{al_version}.png"
+fig.savefig(png_path, dpi=120)
+plt.close(fig)
+print(f"  Bar chart saved to:    {png_path}")
diff --git a/searches/nautilus/simple.py b/searches/nautilus/simple.py
new file mode 100644
index 0000000..d870737
--- /dev/null
+++ b/searches/nautilus/simple.py
@@ -0,0 +1,188 @@
+"""
+Minimal Nautilus Example (HST MGE lens likelihood)
+--------------------------------------------------
+
+Drives the Nautilus nested sampler directly against the HST MGE imaging
+likelihood, bypassing ``af.NonLinearSearch``. Useful as a fast end-to-end
+smoke test of the real PyAutoLens likelihood under a production sampler.
+
+``n_live`` is kept small so the search finishes in a few minutes — this is
+a wiring test, not a converged posterior.
+
+Compare versus ``likelihood/imaging/mge.py`` (single-likelihood profiling
+of the same MGE setup) and ``jax.py`` (JAX-JIT'd likelihood under the same
+sampler).
+
+Requirements:
+    pip install nautilus-sampler
+"""
+
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autolens as al
+
+# Make ``from searches._{setup,metrics}`` importable regardless of how the
+# script is invoked (``python searches/nautilus/simple.py``, ``python -m
+# searches.nautilus.simple``, or a CI runner).
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from searches._metrics import MLTracker
+from searches._setup import (
+    build_analysis,
+    build_dataset,
+    build_model,
+    format_best_fit,
+)
+
+dataset = build_dataset()
+model = build_model()
+analysis = build_analysis(dataset, use_jax=False)
+
+print(f"Model free parameters: {model.total_free_parameters}")
+
+from nautilus import Sampler
+
+
+def prior_transform(cube):
+    """Map a unit cube to physical parameters via the model's priors."""
+    return np.array(model.vector_from_unit_vector(cube))
+
+
+n_likelihood_calls = 0
+tracker = MLTracker()
+
+
+def log_likelihood(params):
+    global n_likelihood_calls
+    n_likelihood_calls += 1
+    instance = model.instance_from_vector(vector=list(params))
+    log_l = float(analysis.log_likelihood_function(instance=instance))
+    tracker.record(log_l)
+    return log_l
+
+
+n_live = 200
+
+sampler = Sampler(
+    prior=prior_transform,
+    likelihood=log_likelihood,
+    n_dim=model.prior_count,
+    n_live=n_live,
+)
+
+t_start = time.time()
+# Run to Nautilus's default convergence (n_eff=10000, f_live=0.01). This
+# may take many thousands of likelihood evaluations against the NumPy MGE
+# -- expect long wall times.
+sampler.run(verbose=True)
+t_elapsed = time.time() - t_start
+
+points, log_w, log_l = sampler.posterior()
+best_idx = np.argmax(log_l)
+best_instance = model.instance_from_vector(vector=list(points[best_idx]))
+max_logl = float(np.max(log_l))
+
+evals_to_ml, time_to_ml = tracker.finalise(max_log_l=max_logl, tolerance=1.0)
+
+# ---------------------------------------------------------------------------
+# Print human-readable summary
+# ---------------------------------------------------------------------------
+
+summary = f"""\
+--- Nautilus (NumPy) Results ---
+Best fit:        {format_best_fit(best_instance)}
+Max log L:       {max_logl:.4f}
+Log evidence:    {float(sampler.log_z):.4f}
+
+--- Performance ---
+Wall time:           {t_elapsed:.2f} s
+Likelihood evals:    {n_likelihood_calls}
+Time per eval:       {t_elapsed / max(n_likelihood_calls, 1) * 1e3:.3f} ms
+ESS:                 {float(sampler.n_eff):.1f}
+Posterior samples:   {len(points)}
+Sampler config:      n_live={n_live}, default n_eff=10000, f_live=0.01
+
+--- Convergence ---
+Converged:           yes (Nautilus default n_eff / f_live)
+Evals to ML:         {evals_to_ml if evals_to_ml is not None else 'n/a'}     (first eval within 1 nat of max log L)
+Time to ML:          {f'{time_to_ml:.2f} s' if time_to_ml is not None else 'n/a'}
+"""
+
+print()
+print(summary)
+
+# ---------------------------------------------------------------------------
+# Write versioned JSON + PNG to results/searches/nautilus/
+# ---------------------------------------------------------------------------
+
+al_version = al.__version__
+result_dict = {
+    "sampler": "nautilus",
+    "backend": "numpy",
+    "instrument": "hst",
+    "model": {
+        "type": "MGE+Isothermal+ExternalShear",
+        "free_parameters": int(model.total_free_parameters),
+    },
+    "sampler_config": {
+        "n_live": n_live,
+        "n_eff_target": 10000,
+        "f_live": 0.01,
+    },
+    "results": {
+        "max_log_likelihood": max_logl,
+        "log_evidence": float(sampler.log_z),
+        "best_fit_summary": format_best_fit(best_instance),
+    },
+    "performance": {
+        "wall_time_s": t_elapsed,
+        "likelihood_evals": int(n_likelihood_calls),
+        "time_per_eval_ms": t_elapsed / max(n_likelihood_calls, 1) * 1e3,
+        "ess": float(sampler.n_eff),
+        "posterior_samples": int(len(points)),
+    },
+    "convergence": {
+        "converged": True,
+        "evals_to_ml": int(evals_to_ml) if evals_to_ml is not None else None,
+        "time_to_ml_s": float(time_to_ml) if time_to_ml is not None else None,
+    },
+    "version": al_version,
+}
+
+results_dir = _REPO_ROOT / "results" / "searches" / "nautilus"
+results_dir.mkdir(parents=True, exist_ok=True)
+json_path = results_dir / f"simple_summary_v{al_version}.json"
+json_path.write_text(json.dumps(result_dict, indent=2))
+print(f"  Results JSON saved to: {json_path}")
+
+# Bar chart of the headline timings
+fig, ax = plt.subplots(figsize=(8, 3))
+labels = [
+    "wall_time (s)",
+    "time_per_eval (ms)",
+    "time_to_ml (s)" if time_to_ml is not None else "time_to_ml (n/a)",
+]
+times = [
+    t_elapsed,
+    t_elapsed / max(n_likelihood_calls, 1) * 1e3,
+    float(time_to_ml) if time_to_ml is not None else 0.0,
+]
+ax.barh(labels, times, color=["#4C72B0", "#55A868", "#C44E52"])
+ax.set_title(f"Nautilus (NumPy) — HST MGE — v{al_version}")
+fig.tight_layout()
+png_path = results_dir / f"simple_summary_v{al_version}.png"
+fig.savefig(png_path, dpi=120)
+plt.close(fig)
+print(f"  Bar chart saved to:    {png_path}")