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//