feat: add searches_minimal examples for HST MGE likelihood

[Jammy2211]

Overview

Add a searches_minimal/ folder to autolens_workspace_developer that mirrors the equivalent folder in autofit_workspace_developer, but runs each sampler against the real MGE imaging likelihood (HST dataset) instead of the 1D Gaussian toy model. These scripts drive Nautilus, Dynesty, Emcee, scipy.L-BFGS-B, and NSS directly against PyAutoFit's Model / Analysis objects, bypassing the NonLinearSearch wrapper — useful for quickly smoke-testing a search against a realistic likelihood.

Plan

• Create searches_minimal/ with one script per sampler (nautilus, dynesty, emcee, lbfgs, nss_simple, nss_jit, nss_grad).
• Share the HST dataset + MGE model + AnalysisImaging setup via a small _setup.py helper.
• NumPy-path samplers use the standard AnalysisImaging and model.instance_from_vector.
• nss_simple.py wraps the NumPy analysis in jax.pure_callback.
• nss_jit.py / nss_grad.py use AnalysisImaging(use_jax=True) + autofit.jax.register_model for the pure-JAX path — exercising the MGE gradient support.
• Keep run times short (small n_live, short chains) — these are smoke tests, not converged runs.

Detailed implementation plan

Affected Repositories

• autolens_workspace_developer (primary)

Work Classification

Workspace

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_developer	main	clean

Suggested branch: feature/searches-minimal
Worktree root: ~/Code/PyAutoLabs-wt/searches-minimal/ (created later by /start_workspace)

Implementation Steps

1. Create searches_minimal/_setup.py with build_dataset(instrument=\"hst\"), build_model(mask_radius=3.5), build_analysis(dataset, use_jax=False) helpers. Mirrors the setup block in jax_profiling/imaging/mge.py (mask radius 3.5, radial-bin over-sampling, 20-gaussian MGE bulges for lens + source, Isothermal mass + ExternalShear).
2. Create searches_minimal/nautilus_simple.py — load setup, run nautilus.Sampler with n_live=50, print best-fit lens / source parameters + log evidence + wall time / likelihood-call count.
3. Create searches_minimal/dynesty_simple.py — dynesty.NestedSampler with small nlive.
4. Create searches_minimal/emcee_simple.py — emcee.EnsembleSampler, walkers initialised near prior medians with small scatter, short chain (few hundred steps).
5. Create searches_minimal/lbfgs_simple.py — scipy.optimize.minimize(method=\"L-BFGS-B\") with bounds derived from the model's priors via model.physical_upper_limits_for_parameters / physical_lower_limits_for_parameters, started at model.physical_values_from_prior_medians.
6. Create searches_minimal/nss_simple.py — run_nested_sampling using NumPy analysis wrapped in jax.pure_callback (same pattern as autofit reference).
7. Create searches_minimal/nss_jit.py — AnalysisImaging(use_jax=True) + autofit.jax.register_model(model) + params_tree = jax.tree_util.tree_map(jnp.asarray, instance), then run_nested_sampling with the pure-JAX likelihood.
8. Create searches_minimal/nss_grad.py — same pure-JAX setup, then run_hmc_sequential_mc to exercise jax.grad through MGE.
9. Smoke-test each script: NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python searches_minimal/<name>.py, verify a best-fit instance + finite log evidence are printed.

Key Files

• searches_minimal/_setup.py — shared HST dataset + MGE model + analysis setup.
• searches_minimal/nautilus_simple.py — direct Nautilus example.
• searches_minimal/dynesty_simple.py — direct Dynesty example.
• searches_minimal/emcee_simple.py — direct Emcee example.
• searches_minimal/lbfgs_simple.py — direct scipy L-BFGS-B example.
• searches_minimal/nss_simple.py — NSS with NumPy likelihood via jax.pure_callback.
• searches_minimal/nss_jit.py — NSS with pure-JAX MGE likelihood.
• searches_minimal/nss_grad.py — NSS HMC-SMC with gradient-accelerated MGE.

Original Prompt

Click to expand starting prompt

The folder autofit_workspace_developer/searches_minimal has examples of how to set up an likelihood function and
have it so a search runs with a minimal interface, for quick testing of lieklihood functions.

Can you set up all of these on autolens_workspacde_developer, using the same hst data used in jax_profiling/imaging
and using the mge model which now supports gradients.

[Jammy2211]

Shipped — 2026-04-21

PRs

• Add minimal HST MGE sampler smoke tests under searches_minimal/ #28

Summary

Added a searches_minimal/ folder of 7 minimal sampler scripts driving the real HST MGE imaging likelihood (parallel to the 1D Gaussian versions in autofit_workspace_developer/searches_minimal/). Each script is capped to a tight iteration budget (max_steps=1, tiny n_live, maxiter caps) so it runs as a wiring test in under ~10 minutes. This is the foundation for a later nautilus-vs-nss comparison, which was explicitly deferred.

Session Notes

• MGE pytree gotcha (nss_jit.py). First attempt used autofit.jax.register_model to pass a stacked pytree ModelInstance through the JIT, but the MGE Basis class has centre / ell_comps / intensity attributes classified as pytree aux_data. With different values per live point the stacked leaves failed with ValueError: Mismatch custom node data. Fixed by switching every JAX script to the flat-vector entry point model.instance_from_vector(vector=params, xp=jnp) — the same path af.Fitness.call uses in production. This is the robust pattern for multi-point JAX sampling with MGE models.
• run_nested_sampling termination (nss_simple.py). Upstream nss.ns.run_nested_sampling loops until logZ_live - logZ < termination, but with the HST MGE likelihood's dynamic range (|log L| ~ 10^5) no sensible termination value stops after one iteration. Replaced the call with an inlined copy of the upstream loop plus a max_steps=1 guard.
• Gradient-JIT compile cost (nss_grad.py). jax.jit(jax.value_and_grad(log_likelihood)) on the MGE inversion compiles for ~200 s. The blackjax HMC+SMC kernel on top pushed the first full compile past 15 min at the original parameters. Reduced to n_live=8, warmup_steps=2, max_steps=1; now completes in ~9 min (still mostly compile).
• /smoke-test coverage. The skill covers autofit_workspace, autogalaxy_workspace, autolens_workspace, autolens_workspace_test, euclid_strong_lens_modeling_pipeline, and HowToLens — not autolens_workspace_developer. The 7 scripts in this PR were individually validated end-to-end in the session instead; re-running them under /smoke-test is not currently possible.

Future Work

• Run nautilus_simple.py and nss_jit.py with realistic (non-smoke-capped) configs and produce a side-by-side timings / best-fit / log-evidence comparison. Explicitly deferred from this PR.