autolens_jax_joss: JOSS benchmark suite + real-data start_here pairing

[Jammy2211]

Overview

The PyAutoLens-JAX paper draft (PyAutoLens/paper_jax/paper.md) promises 10 end-to-end benchmarks (6 single-dataset, 4 joint) that do not yet exist. This task creates a new public repo PyAutoLabs/autolens_jax_joss containing one runnable script per benchmark plus a timing/report harness, and pairs each benchmark to a real-data start_here.py (or multi feature script) in autolens_workspace — upgrading the simulated ones to real data. Because these land in start_here.py, they are user-facing far beyond the paper: each rewrite keeps teaching-quality prose, runs on a normal laptop, and auto-downloads its data.

Per the paper, every benchmark reports: total wall-clock, JAX compile time, number of likelihood evaluations, post-compile runtime (+ hardware/versions). Official numbers on an NVIDIA A100; scripts reproducible on any GPU/CPU.

Dataset decisions (user-approved): ALMA = SDP.81 · point source = RXJ1131-1231 · cluster + weak + joint strong/weak anchor = Abell 2744 · repo = PyAutoLabs org, public.

Plan

• Phase 1 — birth autolens_jax_joss (public) + timing harness + benchmarks already on real data: galaxy-scale imaging (COSMOS-Web Ring F150W) and multi-band (4 bands), paired to imaging/start_here.py and multi/start_here.py.
• Phase 2 — interferometry: public SDP.81 ALMA visibilities (>1M), download-on-first-run; rewrite interferometer/start_here.py; benchmark.
• Phase 3 — point source: RXJ1131-1231 published positions/time delays; rewrite point_source/start_here.py; new multi/features/imaging_and_point_source/ joint script; two benchmarks.
• Phase 4 — group (verify existing dataset is real, make fast) + cluster: A2744 public multiple-image catalogues; rewrite cluster/start_here.py; two benchmarks.
• Phase 5 — weak: new weak/start_here.py from the real A2744 shape catalogue; joint strong+weak (A2744) and joint imaging+interferometer (SDP.81); three benchmarks; final RESULTS.md sweep.

Benchmark ↔ workspace pairing table

Benchmark	Paired example	Data
Galaxy-scale imaging	scripts/imaging/start_here.py	cosmos_web_ring (real, JWST F150W)
Interferometry	scripts/interferometer/start_here.py	SDP.81 (ALMA, >1M vis)
Point source	scripts/point_source/start_here.py	RXJ1131-1231 (positions + delays)
Group	scripts/group/start_here.py	existing real group lens (verify)
Cluster	scripts/cluster/start_here.py	Abell 2744 (multiple images)
Weak	scripts/weak/start_here.py (new)	Abell 2744 JWST shape catalogue
Multi-band	scripts/multi/start_here.py	cosmos_web_ring 4 wavebands
Strong + weak	scripts/weak/features/strong_lensing/	Abell 2744
Imaging + point source	scripts/multi/features/imaging_and_point_source/ (new)	RXJ1131-1231
Imaging + interferometer	scripts/multi/features/imaging_and_interferometer/	SDP.81

Detailed implementation plan

Affected Repositories

• autolens_jax_joss (new, primary deliverable)
• autolens_workspace (primary tracked repo)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace	main	clean

Suggested branch: feature/jax-joss-benchmarks
Worktree: ~/Code/PyAutoLabs-wt/jax-joss-benchmarks/

Data hosting policy

Large real datasets are never committed: sizeable data (SDP.81 visibilities, A2744 catalogues/imaging, RXJ1131 cutouts) is fetched automatically on first run from a stable public URL (Zenodo preferred; weak/real_data/a2744.py download-and-cache pattern is the template) into gitignored dataset/ paths, with checksum + printed provenance. Small derived products (position/delay JSONs, model CSVs, kB–few-MB) may live in dataset/ per existing convention.

Implementation Steps

1. gh repo create PyAutoLabs/autolens_jax_joss --public; skeleton: README (pairing table + how to run), LICENSE, .gitignore, requirements.txt, benchmarks/, results/.
2. harness.py: runner recording wall-clock, compile time (first-eval − steady-state), likelihood evals (search counter), post-compile runtime, GPU/CPU + jax/autolens versions → results/<benchmark>.json; regenerates RESULTS.md.
3. benchmarks/imaging.py + benchmarks/multi_band.py — same data/model family as their paired start_here scripts (MGE lens light + SIE+shear + pixelized source, gradient-based search), tuned toward ~5 min on A100.
4. Phases 2–5 per the checklist above: each = real-data acquisition (download-on-first-run) + start_here rewrite (teaching prose, laptop-runnable) + JOSS benchmark script(s).
5. A100 timing runs via the RAL HPC pipeline; commit results/*.json + regenerated RESULTS.md.

Key Files

• autolens_jax_joss/harness.py — timing/report harness (new)
• autolens_jax_joss/benchmarks/*.py — the 10 benchmarks (new)
• autolens_workspace/scripts/{interferometer,point_source,cluster}/start_here.py — real-data rewrites
• autolens_workspace/scripts/weak/start_here.py — new front door for weak lensing
• autolens_workspace/scripts/multi/features/imaging_and_point_source/ — new joint feature script
• PyAutoLens/paper_jax/paper.md — benchmark section these numbers feed (not edited here)

Risks

• ~5 min A100 targets are aspirational until measured — report honestly, tune (e.g. lax.map batch_size) rather than fake.
• SDP.81 visibility extraction may need CASA; fallback = pre-exported uv table on Zenodo.
• Cluster/group JAX point-solver performance at scale less battle-tested than imaging.

Original Prompt

Click to expand starting prompt

Create a repo autolens_jax_joss which provides examples to run each JAX benchmark from the PyAutoLens-JAX paper draft and reports the run time and info. These benchmarks currently do not exist.

Single-dataset and single-regime benchmarks (establish GPU acceleration and autodiff across lensing scales and data types):

• Galaxy-scale CCD imaging: Model JWST COSMOS-Web Ring F150W imaging, including lens-light subtraction and a pixelized source reconstruction, in approximately five minutes.
• Interferometry: Model a real ALMA strong-lensing dataset containing more than one million interferometer visibilities in approximately five minutes.
• Point-source lensing: Model a real multiply imaged quasar or supernova using point-source observables, including image positions and, where available, time delays or flux information, in under five minutes.
• Group-scale strong lensing: Model a real group-scale lens containing multiple deflecting galaxies in under five minutes, demonstrating that PyAutoLens-JAX is not restricted to isolated galaxy-scale lenses.
• Cluster-scale strong lensing: Model a real cluster lens with multiple mass components, multiple images, and potentially multiple source planes in under five minutes.
• Weak lensing: Fit a weak-lensing shear catalogue using a differentiable JAX likelihood in under five minutes, demonstrating that PyAutoLens-JAX is not restricted to strong-lensing data.

Joint and multi-dataset benchmarks (different datasets, lensing regimes, and physical scales combined in a single differentiable, GPU-accelerated probabilistic model):

• Multi-band imaging: Jointly model the four available JWST COSMOS-Web Ring bands, constraining a common lens mass model while fitting the wavelength-dependent lens and source emission in each dataset.
• Joint strong and weak lensing: Constrain a single group- or cluster-scale mass model using both strong-lensing and weak-lensing observables.
• Imaging and point-source lensing: Jointly model extended arcs and point-source constraints from a lensed quasar or supernova within the same lens model.
• Imaging and interferometry: Jointly fit optical or infrared imaging and radio or submillimetre interferometer visibilities, constraining a common mass model using complementary observations of the lensed source.

Pairing requirement: pair each benchmark to the start_here.py file in each autolens_workspace package. For example the JWST COSMOS-Web Ring F150W example, which is made fast, is paired to autolens_workspace/scripts/imaging/start_here.py. All 6 single-dataset start_here.py examples should use real data, and each JOSS example should use the same real data.

The four multi examples pair to examples in autolens_workspace/scripts/multi or scripts/weak/features/strong_lensing. A new multi script will likely be needed for the imaging + point-source lensing joint-modelling benchmark.

[Jammy2211]

Progress — repo born, all 10 benchmarks authored, workspace real-data rewrites in place

autolens_jax_joss is live: https://github.com/PyAutoLabs/autolens_jax_joss (public, PyAutoLabs org).

Benchmark repo (pushed, 9c66b02):

• harness.py — records the paper's four numbers (total wall-clock, JAX compile time — measured by compiling the same batched value_and_grad program the search builds, likelihood evals, post-compile runtime) + hardware/versions into results/*.json, regenerates RESULTS.md. Datasets always fetched on first run (workspace raw URLs / pyRRG / VizieR / hips2fits), never committed.
• All 10 benchmark scripts written; imaging and point_source validated end-to-end locally in --quick mode (point_source: 11,077 evals through positions+delays likelihood).

Workspace branch feature/jax-joss-benchmarks (worktree, not yet pushed):

• point_source/start_here.py → real RXJ1131-1231 (Suyu+13 positions, Tewes+13 COSMOGRAIL delays; committed JSONs) — passes TEST_MODE end-to-end including the time-delay fit.
• weak/start_here.py → new front door on the real A2744 pyRRG catalogue.
• cluster/start_here.py → real Abell 2744: Bergamini+23 gold systems (7 multi-plane sources, 25 images, z=1.69–5.66) + 188 members from VizieR J/A+A/670/A60, CSVs committed with prep.py + provenance README; HST image via hips2fits on first run.
• New multi/features/imaging_and_point_source/ — real RXJ1131 HST H-band arcs + positions/delays joint fit (the missing pairing).
• New weak/features/strong_lensing/a2744.py — real cluster-scale joint strong+weak on A2744, both datasets in one shared coordinate frame.
• imaging/start_here.py, multi/start_here.py, group/start_here.py already used real data (COSMOS-Web Ring; Euclid group) — no changes needed, benchmarks pair directly.

Known blocked leg (Phase 2): SDP.81 ALMA visibilities require a one-off CASA export from the public SV measurement sets (tens of GB) — shipped as data_prep/sdp81/{README.md,export_uv.py}; benchmarks/interferometer.py + imaging_and_interferometer.py exit with instructions until the products exist (then: Zenodo-host + rewrite interferometer/start_here.py).

Next: smoke-validate remaining workspace scripts, push branch + PR, then A100 timing runs.

[Jammy2211]

End-of-day state (overnight runs in progress)

Shipped today:

• PR Real-data start_here suite for the PyAutoLens-JAX JOSS benchmarks #282 open (pending-release): 2 real-data start_here rewrites (RXJ1131 point source, A2744 cluster), new weak/start_here.py, new multi/features/imaging_and_point_source/, new weak/features/strong_lensing/a2744.py, datasets + notebooks.
• autolens_jax_joss: harness + 10 benchmarks; fixes landed for visualization-free timing and never-resume semantics.
• Filed PyAutoLens#614 (weak JAX-path visualization crash).

Official A100 rows so far: point_source 5.95 min (13,100 evals) · group 8.44 min · imaging 10.0 min (timing good; cold-start multi-start Adam missed the pixelized basin — Nautilus comparison queued) · weak 0.75 min (3.9 s compile).

Running overnight on RAL: job 330501 (cluster — 7h+ already; clear finding that the A2744 multi-plane benchmark needs tuning against its 5-min target) and job 330527 (strong_and_weak → imaging_and_point_source → multi_band → imaging Nautilus comparison).

Tomorrow: pull all result JSONs → commit RESULTS.md; clean viz-off re-timings of the three shakedown rows; merge decision on #282 (human); SDP.81 CASA export remains the one data-blocked leg (data_prep/sdp81/).

[Jammy2211]

First complete A100 sweep — every runnable benchmark has an official row

results/RESULTS.md now carries all 8 runnable benchmarks (SDP.81's two remain data-blocked):

Benchmark	Total	Compile	Evals	Verdict
weak	0.75 min	3.9 s	4,800	✅ well under target
point_source	5.18 min	27 s	12,900	✅ at target
group	6.42 min	49 s	4,800	✅ near target, deterministic
imaging (Nautilus)	16.3 min	42 s	44,586	✅ finds the basin (logL +7357); tune toward 5
imaging (multi-start Adam)	10.2 min	45 s	4,800	❌ cold starts miss the pixelized basin — the open search question, now with a head-to-head row
multi_band (4 bands)	27.8 min	203 s	4,800	✅ converged (logL +48,137); tune
imaging_and_point_source	129.6 min	39 s	389,421	runs end-to-end; large sampling budget + demo-grade noise-map caveat documented in the paired script
cluster	442 min	90 min	271,216	converged; compile-dominated
strong_and_weak	515 min	97 min	231,401	converged; compile-dominated

Two systematic findings for the paper:

1. The cluster-scale rows are compile-dominated (~1.5 h factor-graph XLA fusion compile) — the exact pathology autolens_profiling#71 just characterised, with a persistent-cache mitigation measured at ~6× on warm repeats. Tuning prompt filed: draft/feature/autolens_workspace/joss_cluster_benchmark_tuning.md.
2. Pixelized-source inference: nested sampling on the GPU-batched JAX likelihood is currently the robust choice — Nautilus reaches logL +7357 in 16 min where 16-start cold Adam lands at −10⁷. Feeds autolens_workspace_developer#100.

Remaining human actions: merge PR #282; SDP.81 CASA export (autolens_jax_joss/data_prep/sdp81/) to unblock the two interferometry benchmarks + the interferometer/start_here.py rewrite.