research: 4-lens + 4-source MGE search benchmark (JAX gradient optimizers vs Nautilus)

[Jammy2211]

Overview

searches/ currently runs a JAX gradient MAP optimizer (af.MultiStartAdam) only on the single-lens/single-MGE-source cell. This task extends that proven MGE cell to a 4-deflector + 4-source group-scale model (~50–55 non-linear params) to answer: do the JAX gradient max-likelihood optimizers scale to a higher-dimensional, harder lens model? We benchmark the MultiStart gradient family (fixed-step and auto-convergence) against Nautilus and check truth recovery. If cold-start optimizers fail, a careful-initialization arm follows.

Plan

• Write simulators/group4_mge.py — simulate a 4-lens + 4-source MGE imaging dataset with known truth (truth.json).
• Add a group/mge cell to the searches/ harness (_setup.py dataset+model builders, AnalysisImaging).
• Add searches/_recovery.py — compare max_log_likelihood_instance to truth; emit a per-galaxy recovery block into the summary JSON.
• Register the MultiStart gradient family in two modes — fixed-step (convergence=None) and automatic-convergence (af.MultiStartGradientConvergence), including a first-class multi_start_prodigy_autoconv cell — plus Nautilus as the reference.
• Run the benchmark, aggregate a recovery + wall-time comparison.
• Contingency: if cold-start fails recovery, add narrow-prior / warm-start initialization and re-benchmark.

Detailed implementation plan

Affected Repositories

• autolens_profiling (primary) — workspace/research only, no library source changes

Branch Survey

Repository	Current Branch	Dirty?
./autolens_profiling	research/multiband-compile-ab	dirty (untracked dataset/imaging/jwst_lw/)

Suggested branch: feature/group4-mge-search-benchmark (fresh worktree off main via /start_workspace)

Model (settled)

• 4 lenses: al.mp.Isothermal mass + MGE light (al.model_util.mge_model_from), ExternalShear on lens_0; per-galaxy centre prior seeded near truth.
• 4 sources: per-galaxy MGE light at distinct source-plane positions. Amplitudes stay linear (inversion); non-linear count ~50–55.

Samplers

• nautilus (reference), and MultiStart JAX family run fixed-step (n_steps=300) and auto-convergence (MultiStartGradientConvergence(check_for_convergence=True, window=50, rtol=1e-4, atol=1e-3, min_steps=100)). First-class multi_start_prodigy_autoconv cell; summary surfaces actual steps-taken for the early-stop-vs-fixed-300 comparison.

Implementation Steps

1. simulators/group4_mge.py (mirror simulators/imaging.py CLI + simulators/group.py geometry) -> dataset + truth.json under dataset/imaging/group4_mge/<instrument>/.
2. searches/_setup.py: dataset_class="group" -> _build_group_imaging (auto-simulate via dataset_type="group4_mge"), _group_mge_model(), AnalysisImaging branch, _mask_radius_for + format_best_fit updates.
3. searches/_recovery.py + _runner._build_summary "recovery" block loading truth.json.
4. searches/_samplers.py: ("group","mge"): 200 n_live; fixed + autoconv builders; register in SAMPLER_BUILDERS.
5. Leaf scripts searches/<sampler>/group/mge.py; add cell to searches/sweep.py.
6. Run cells (CPU JAX first); searches/aggregate.py; record findings.
7. Contingency (Phase 4): narrow-prior / warm-start init arm if cold-start fails.

Key Files

• New: simulators/group4_mge.py, searches/_recovery.py, searches/{nautilus,multi_start_adam,multi_start_prodigy,multi_start_prodigy_autoconv,multi_start_lion,multi_start_adabelief}/group/mge.py
• Edit: searches/{_setup,_samplers,_runner,sweep}.py, searches/README.md

Verification

ruff check . && ruff format --check .; AUTOLENS_PROFILING_SMOKE=1 import-smoke; simulate -> 4 lensed sources; Nautilus recovery anchor; then gradient sweep + aggregate.

Original Prompt

Click to expand starting prompt

Research profiling experiment in the autolens_profiling repo. We currently have autolens_profiling examples that run JAX gradient max-likelihood optimizers on a single lens galaxy with a single MGE source. Extend this to a much higher-dimensional, harder model: 4 lens galaxies + 4 source galaxies. Write a simulator.py that generates the dataset from known input truth, then run the existing JAX gradient optimizers (max-likelihood samplers) alongside Nautilus and record whether any of them scale to this dimensionality and recover the input truth. If none of the optimizers succeed, investigate more careful initialization strategies. This is an exploratory benchmark, not a library change.

[Jammy2211]

Progress checkpoint (phases 1–3 code complete; benchmark runs pending)

Branch feature/group4-mge-search-benchmark pushed (bddedc5). No PR yet.

Done + verified

• Simulator simulators/group4_mge.py — 4 deflectors (Isothermal mass + Sersic light, shear on primary) lensing 4 SersicCore sources; single-source-of-truth GROUP4_TRUTH; writes data/psf/noise_map.fits + tracer.json + truth.json + preview PNG to dataset/imaging/group4_mge/hst/. Ran; preview shows 4 bright deflector cores + multiple lensed arcs (a genuinely rich group).
• Model cell searches/_setup.py dataset_class="group" → _group_mge_model() = 54 free params. Every galaxy's light+mass centre seeded near truth (Gaussian σ 0.1 lens / 0.3 source) to break the 4-lens/4-source permutation symmetry; Einstein radii, ellipticities, shear keep broad default priors. Auto-simulates via dataset_type="group4_mge". AnalysisImaging. Likelihood evaluates on both numpy and JAX paths.
• Recovery searches/_recovery.py scores max_log_likelihood_instance vs truth.json (per-lens θ_E frac-error + mass-centre distance, primary shear, per-source centre; overall_pass); _runner embeds a "recovery" block in the summary JSON. Wiring proven on a prior-median instance.
• Samplers — generic build_multi_start covers the whole JAX family: multi_start_{adam,prodigy,lion,adabelief} fixed-step + multi_start_prodigy_autoconv (Prodigy + af.MultiStartGradientConvergence, window 50 / min_steps 100). ("group","mge") n_live=200. Leaf scripts for nautilus + all 5 gradient samplers. sweep.py gains 6 group cells; README documents the benchmark. ruff clean, import-smoke green.

The open blocker → resume tomorrow

The first real run (multi_start_adam group, local CPU fp64) sat 20+ min in JAX compile / pre-fit with zero step output. The 54-param, 8-galaxy MGE-inversion value_and_grad vmap graph is too heavy for a CPU XLA compile (matches known compile-time behaviour on large graphs).

Resume plan:

1. First check whether the detached overnight CPU run (PID 25434) ever produced output/searches/multi_start_adam/group/mge/hst/... or a results JSON — if still hung, kill it.
2. Re-run the benchmark on the laptop GPU (~/venv/PyAutoGPU, JAX_PLATFORM_NAME=cuda JAX_PLATFORMS=cuda,cpu XLA_PYTHON_CLIENT_MEM_FRACTION=0.5, --config-name local_gpu_fp64) — start with multi_start_adam, then the rest of the family, then the nautilus anchor.
3. If GPU compile is still too heavy, dial the graph down: _GROUP4_MGE_TOTAL_GAUSSIANS (10→6) and/or _MULTI_START_N_STARTS (64→32) in searches/_setup.py / searches/_samplers.py.
4. Record which optimizers recovered truth + wall-time vs Nautilus; run searches/aggregate.py.
5. Phase 4 (contingency) only if cold-start fails recovery: narrow-prior / warm-start init arm.

[Jammy2211]

Worktree wrapped via merge (PR #83, human decision 2026-07-24). Code phases 1-3 + the first GPU results (multi_start_adam) are on main; the worktree and claim are released. REMAINING on this issue: the gradient-family sweep (prodigy/lion/adabelief/prodigy_autoconv) + Nautilus anchor runs on GPU, then recovery/walltime aggregation. NOTE: the repo-wide scripts/// restructure is now unblocked and imminent — after it lands, the group4 cells live under scripts/cluster/ (human taxonomy decision: group-scale = cluster profiling use case), so resume the runs from the NEW paths.

[Jammy2211]

Closing as superseded by the 2026-08-17 inference-methods programme (#134), which now owns sampler-benchmark direction.

What shipped: the simulator, the group/mge cell, _recovery.py, the fixed-step + auto-convergence MultiStart registrations and the first GPU results all merged in #83. The scripts then moved in the scripts/<dataset>/<task>/ restructure (#84).

What did not: the gradient-family sweep (prodigy / lion / adabelief / prodigy_autoconv + the Nautilus anchor) and the recovery/wall-time aggregation, plus the Phase-4 careful-initialization contingency. Those are not carried forward as scoped here — the programme in #134 decides whether and how they get run.

The resume recipe is preserved in PyAutoMind at complete/archive/shelved/research_profiling_experiment_in_the_autolens_pr.md (PyAutoLabs/PyAutoMind#205): the laptop-GPU env, the post-#84 script paths (scripts/cluster/searches/<sampler>/mge.py), the sweep.py --only <sampler>/group/mge invocation and its internal group-cell keying, and the compile-cost knobs (_GROUP4_MGE_TOTAL_GAUSSIANS 10→6, _MULTI_START_N_STARTS 64→32). Warm output is still in the main checkout under output/searches/.

Closing as not planned rather than completed — the benchmark question this issue asked was never fully answered, so it should not read as delivered. Reopen if the sweep is wanted under this issue rather than folded into #134.

---

Generated by Claude Code