feat: group4 MGE search benchmark — code + first GPU results#83
Merged
Conversation
WIP checkpoint — code complete for phases 1-3, benchmark runs pending.
- simulators/group4_mge.py: simulate 4-deflector + 4-source group imaging
(single-source-of-truth GROUP4_TRUTH) + write truth.json for recovery scoring.
- searches/_setup.py: dataset_class='group' path (_build_group_imaging auto-sim
via dataset_type=group4_mge; _group_mge_model, ~54 free params; centres seeded
near truth to break permutation symmetry, geometry priors broad; AnalysisImaging).
- searches/_recovery.py: score max_lh_instance vs truth.json (per-lens theta_E +
centre, primary shear, per-source centre; overall_pass).
- searches/_runner.py: recovery block in summary JSON for group cells.
- searches/_samplers.py: generic build_multi_start for the whole MultiStart JAX
family (adam/prodigy/lion/adabelief) + prodigy auto-convergence variant via
af.MultiStartGradientConvergence; ('group','mge') n_live=200.
- leaf scripts searches/<sampler>/group/mge.py for nautilus + 5 gradient samplers.
- sweep.py: 6 group cells added; README documents the benchmark + recovery.
Verified: simulate renders 4 lensed sources; group cell builds + evaluates
likelihood (numpy + JAX) at 54 params; recovery scoring wired; ruff clean;
import-smoke green. Nautilus anchor + gradient sweep runs are the remaining step.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JiU4VbBD9X6jPYJMCx3boh
…ll (#82) Two blockers found running the group4 benchmark, both measured: 1. Pre-fit visualization of the 8-galaxy model costs ~1 HOUR of wall-clock before the optimizer takes a single step (18:06->19:06 on the GPU run). attach_viz_timer now takes disable=, swapping the visualize hooks for no-ops (SEARCHES_DISABLE_VIZ=1). The summary records viz_disabled: true and still counts viz_n_calls, so a viz-skipped row can never be mistaken for a real end-to-end visualization measurement. 2. n_starts=64 on the group cell requests a single 4.71 GiB allocation and OOMs a 6 GB laptop GPU (RESOURCE_EXHAUSTED in search.py:301 pulling the per-start foms to host). Each start is a full vmap replica of the likelihood and the group replica is much larger than the single-lens one. n_starts is now per-cell (_MULTI_START_N_STARTS_BY_CELL = {'group': 16}) with a SEARCHES_N_STARTS override. The 16-start default is a LOCAL VRAM accommodation, not a statement about the method — an A100 runs the full 64; set SEARCHES_N_STARTS=64 there. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JiU4VbBD9X6jPYJMCx3boh
… cell (#82) Two defects found by the first real group run (multi_start_adam, 32 starts): 1. convergence=None does NOT mean 'fixed step'. samples_info.json from a plain adam run records check_for_convergence: true, window 50, rtol 1e-4, atol 1e-3, min_steps 100 — the search defaults convergence ON, and those are exactly the values the '*_autoconv' cell was passing. The fixed-step and autoconv arms were therefore the SAME run and the A/B was vacuous. build_multi_start now ALWAYS passes an explicit convergence object: check_for_convergence=False for the fixed-step cells, True for *_autoconv. _sampler_config_dict records it for both arms. 2. n_steps=300 is far too few for this cell. The adam run stopped on 'max_steps' with converged: false while its FOM was still falling 7.2% over the final 50 steps (747335 -> 464003, 37.9% total drop, still descending). Reading 'gradient optimizers cannot fit this model' off that run would be an artefact of the step budget, not a property of the method. n_steps is now per-cell (group: 3000) with a SEARCHES_N_STEPS override. Also noted: likelihood_evals is miscounted for MultiStart (33 = n_starts+1) because the 300 steps run inside a jitted loop that never increments the Python-side counter, which makes search.summary's 'Total Samples' and 'Time Per Sample' meaningless for these searches. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JiU4VbBD9X6jPYJMCx3boh
…ator summary First benchmark row for the group4 (4-deflector) MGE cell: durable results from the laptop-GPU run that the CPU-XLA blocker forced (54-param 8-galaxy vmap graph). Family sweep + Nautilus anchor remain pending (issue #82). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRQH5HrmMPfpWkmfh3ysJb
This was referenced Jul 24, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Wraps the group4-mge-search-benchmark worktree (#82) per human decision: phases 1-3 code (simulator with GROUP4_TRUTH, group dataset_class, recovery scoring, generic multi-start sampler builders, 6 leaf scripts, sweep cells) plus the first durable GPU results (multi_start_adam, local_gpu_fp64) and simulator summary. Remaining benchmark runs (gradient family + Nautilus anchor) continue on main — tracked on #82. Merging now also unblocks the repo-wide scripts/// restructure (group4 cells will land under scripts/cluster/ per the human's taxonomy decision).
🤖 Generated with Claude Code
https://claude.ai/code/session_01MRQH5HrmMPfpWkmfh3ysJb