Goal: make the JAX gradient MAP optimizers (af.MultiStartAdam/MultiStartADABelief/MultiStartLion, PyAutoFit#1369/#1374) a standard, recommended option for autolens users with an MGE source, and likely a Sersic source too — by tuning the settings that control convergence per model complexity against a Nautilus baseline. Results feed the workspace settings + documentation.
Uses the full first-class search (search.fit) here in searches/ — not searches_minimal (prototype tier).
Model cases
- Sersic lens + Sersic source — NEW
sersic model type (lp_linear.Sersic lens + Isothermal + shear + lp_linear.SersicCore source; mirrors the simulator truth).
- MGE lens + MGE source — existing
mge cell (multi_start_adam/imaging/mge already registered).
Pixelized sources are out of scope (compile-pathological — autolens_workspace_developer#100).
Settings array
Per model case: n_starts × n_steps × learning_rate (e.g. 8/16/32/64/128 × 100/300/1000 × 1e-3/1e-2/1e-1).
batch_size is not a tuning knob — numerically inert (verified on A100 across {None,1,4,14,100}); it only bounds VRAM and comes from the vram table.
Baseline + scoring
af.Nautilus on the identical model per case. JAX rows must set force_x1_cpu=True + use_jax_vmap=True (else nautilus.Sampler forks and corrupts JAX state).- Score: basin recovery (
einstein_radius vs truth 1.6), max logL vs Nautilus, per-start basin hit rate, wall (compile vs sample).
What this harness lacks (to build)
- No
sersic model type (_setup.py dispatches mge/pixelization/delaunay/point-source only).
- No truth / correctness scoring — it profiles cost, not whether the search found the right answer.
- No search-hyperparameter axis —
sweep.py sweeps (sampler, dataset, model) × (CPU/GPU × fp64/mp); tuning needs a settings grid.
- Priors are already broad-uniform by design ("the search scripts need the sampler to actually search a realistic prior volume") — correct for tuning; do not switch to truth-centred priors.
Affordable: the MGE benchmark ran 128 starts in ~50 s warm on an A100.
Goal: make the JAX gradient MAP optimizers (
af.MultiStartAdam/MultiStartADABelief/MultiStartLion, PyAutoFit#1369/#1374) a standard, recommended option for autolens users with an MGE source, and likely a Sersic source too — by tuning the settings that control convergence per model complexity against a Nautilus baseline. Results feed the workspace settings + documentation.Uses the full first-class search (
search.fit) here insearches/— notsearches_minimal(prototype tier).Model cases
sersicmodel type (lp_linear.Sersiclens + Isothermal + shear +lp_linear.SersicCoresource; mirrors the simulator truth).mgecell (multi_start_adam/imaging/mgealready registered).Pixelized sources are out of scope (compile-pathological — autolens_workspace_developer#100).
Settings array
Per model case:
n_starts×n_steps×learning_rate(e.g. 8/16/32/64/128 × 100/300/1000 × 1e-3/1e-2/1e-1).batch_sizeis not a tuning knob — numerically inert (verified on A100 across {None,1,4,14,100}); it only bounds VRAM and comes from the vram table.Baseline + scoring
af.Nautiluson the identical model per case. JAX rows must setforce_x1_cpu=True+use_jax_vmap=True(elsenautilus.Samplerforks and corrupts JAX state).einstein_radiusvs truth 1.6), max logL vs Nautilus, per-start basin hit rate, wall (compile vs sample).What this harness lacks (to build)
sersicmodel type (_setup.pydispatches mge/pixelization/delaunay/point-source only).sweep.pysweeps (sampler, dataset, model) × (CPU/GPU × fp64/mp); tuning needs a settings grid.Affordable: the MGE benchmark ran 128 starts in ~50 s warm on an A100.