feat: cluster test workspace — likelihood sanity + JAX likelihood functions

[Jammy2211]

Overview

Stand up dedicated cluster functionality in autolens_workspace_test to stress-test the cluster likelihood function end-to-end. Today autolens_workspace/scripts/cluster/ has the production csv_api / simulator / modeling scripts (shipped via #187), but there's no test-workspace counterpart that perturbs the truth model and validates likelihood behaviour. This task adds eight deliverables: a paired csv_api.py + simulator.py, the existing visualization_cluster.py moved to its permanent home, two likelihood sanity-check scripts (source-plane + image-plane chi², multi-redshift sensitivity), an imaging-mode counterpart, end-to-end visualization, and two new JAX likelihood functions with numerical assertions.

The goal is explicit: find bugs, edge cases, and anything wrong in the cluster likelihood path. Surprises become follow-up issues; this PR ships the test scaffolding that surfaces them.

Plan

• New scripts/cluster/csv_api.py paired to autolens_workspace's csv_api guide. Includes the full galaxy population — main + (optional) extra + scaling + source — so workspace_test can exercise every tier the cluster code path supports.
• New scripts/cluster/simulator.py that loads the model CSVs csv_api.py produces and emits the full dataset (data.fits, point_datasets.csv, source-light CSV, scaling_galaxies.csv, etc.) into dataset/cluster/<name>/.
• Move scripts/imaging/visualization_cluster.py → scripts/cluster/visualization.py, run against the new simulator outputs.
• New scripts/cluster/likelihood_sanity.py: take the truth mass model from mass.csv, perturb each numeric mass parameter by 0.1%, 1%, 5%, and assert:
  • At 0% perturbation: chi² ≈ 0 for both FitPositionsSource (source-plane) and FitPositionsImagePair (image-plane) likelihoods.
  • Monotonic log-likelihood decrease as perturbation grows for both chi² flavours.
• New scripts/cluster/likelihood_redshift_sensitivity.py: perturb the source redshifts (small multi-plane displacements) and assert the max log-likelihood lands at the truth redshift; small redshift errors produce small likelihood decreases.
• New scripts/cluster/likelihood_imaging.py: same perturbation tests but with AnalysisImaging + light profiles (the truth values from light.csv) producing arcs. Assert the likelihood patterns hold for the imaging-mode likelihood too.
• Visualization for the redshift-sensitivity and imaging-mode scripts: image-plane positions overlaid on data, residual maps of the light-profile fit, and any other plots that help a human spot pathology.
• New scripts/jax_likelihood_functions/cluster/single_plane.py + multi_plane.py: JAX-jit'd cluster likelihood functions with numerical assertions matching the perturbation test results above. Modelled on the existing scripts/jax_likelihood_functions/point_source/{image_plane,source_plane,point}.py.

Detailed implementation plan

Affected Repositories

• autolens_workspace_test (primary, single repo)

Work Classification

Workspace.

Branch Survey

Repository	Current Branch	Dirty?
autolens_workspace_test	main	clean

Suggested branch: feature/cluster-test-workspace
Worktree root: ~/Code/PyAutoLabs-wt/cluster-test-workspace/

Implementation Steps

1. scripts/cluster/csv_api.py — port autolens_workspace/scripts/cluster/csv_api.py's round-trip demonstration into a workspace_test-style guide. Same three family CSVs (mass / light / point) plus scaling_galaxies. The "extras" tier (per-galaxy individually modelled, NOT shared scaling relation) is supported by the same mass.csv schema — include a small example extra_0 row in this guide to demonstrate.
2. scripts/cluster/simulator.py — mirror autolens_workspace/scripts/cluster/simulator.py under workspace_test naming conventions (deterministic seed, outputs in dataset/cluster/test/). Drop the heavy JAX/PointSolver setup if not needed for the downstream tests; the simulator should always be runnable under PYAUTO_TEST_MODE=2.
3. Move scripts/imaging/visualization_cluster.py → scripts/cluster/visualization.py. Update the dataset_path to point at the new test dataset. Verify the three reference PNGs (overlaid positions, per-source grid, critical curves) still render.
4. scripts/cluster/likelihood_sanity.py — load the truth mass model from mass.csv. For each (galaxy_name, param_name) in the dPIE / NFW parameter set, perturb by ε ∈ {0, 0.001, 0.01, 0.05}, build a Tracer, evaluate both FitPositionsSource and FitPositionsImagePair against the simulator's point_datasets.csv. Assert chi² ≈ 0 at ε=0 (numerical tolerance ~1e-3); assert log-likelihood is monotone decreasing in |ε|.
5. scripts/cluster/likelihood_redshift_sensitivity.py — perturb each source's redshift by small multi-plane displacements (δz ∈ {0, 0.01, 0.05, 0.1}). Build the same fit objects, assert the truth redshift maximises log-likelihood and that small δz produces small likelihood decrements.
6. scripts/cluster/likelihood_imaging.py — same perturbation tests but with al.AnalysisImaging + the light profiles from light.csv. The truth arcs are simulated from data.fits; perturbing the mass profile should produce visibly wrong arcs and a worse figure_of_merit. Assert the imaging-mode log-likelihood pattern matches the point-source pattern.
7. Visualization — figures next to each likelihood script: chi²-vs-perturbation curves; image-plane positions overlaid on data for the perturbed runs; residual maps for the imaging fit. Save into dataset/cluster/test/visualization/.
8. scripts/jax_likelihood_functions/cluster/single_plane.py + multi_plane.py — JAX-jit'd cluster likelihood functions. Model on the existing scripts/jax_likelihood_functions/point_source/image_plane.py (single-source, single-plane) and the multi-source / multi-plane case respectively. Each carries a numerical assertion block referencing the perturbation results from step 4/5 — e.g. assert chi_squared < tolerance for epsilon == 0. Verify both backends (numpy + JAX) produce matching values.

Key Files

• New: scripts/cluster/{csv_api.py, simulator.py, visualization.py, likelihood_sanity.py, likelihood_redshift_sensitivity.py, likelihood_imaging.py}
• New: scripts/jax_likelihood_functions/cluster/{__init__.py, single_plane.py, multi_plane.py}
• Move: scripts/imaging/visualization_cluster.py → scripts/cluster/visualization.py
• Reference (read-only): scripts/jax_likelihood_functions/point_source/{image_plane.py, source_plane.py, point.py}
• Reference (read-only): autolens_workspace/scripts/cluster/{csv_api.py, simulator.py, modeling.py}

Smoke Tests

Add new scripts to smoke_tests.txt once they reliably pass under PYAUTO_TEST_MODE=2. Likelihood scripts are pure-Python (no Nautilus search) so they should smoke quickly.

Open Questions (resolve during implementation)

• Should the "extras" tier (individually modelled middle tier) be exercised explicitly in csv_api.py, or is the existing main / scaling / source surface sufficient for cluster tests? The user's prompt mentions "mains, extra, scaling, source" so probably include an extra_0 example in the guide and a row in mass.csv.
• How to encode source redshift in a way likelihood_redshift_sensitivity.py can override at fit time without re-running the simulator. Likely just constructs perturbed PointDataset copies with mutated redshift attrs.
• For the imaging-mode test (step 6), how much of the modeling pipeline does the script actually need? Just analysis.log_likelihood_function(instance) on perturbed instances, probably — no search.

Stress-Test Mandate

Per the prompt: "be persistent, be nit-picking, look for obvious issues with our implementation." Bugs found during implementation become separate follow-up issues; the test scaffolding ships regardless.

Original Prompt

Click to expand starting prompt

Lets set up autolens_workspace_test to have dediciated cluster functionality, in particular:

1. Write a scripts/cluster/csv_api.py, which follows the autolens_workspace/scripts/cluster/csv_api.py and therefore includes a full suite of mains, extra, scaling, source galaxies/
2. Write a scripts/cluster/simulator.py, which uses the autolens_workspace/scripts/cluster/simulator.py and uses the csv exa,ple but outputs its point dataset csv, source light profile .csv and all other data.
3. move scripts/imaging/visualization_cluster.py to scripts/cluster/visualization.py, run against the simulator.py output.
4. Set up a likelihood sanity check, which uses the simulated cluster (will need a scritps/cluster/simulator file) and the does sanity checks on its source plane chi-squared (using Point objects with FitPositionsSource) and image plane chi squared (using FitPositionsImagePair). I think the sensible way to do this is to perturb the mass model parameter inputs to the simulator script by a small amount (e.g. 0.1%) and medium amount (e.g. 1%) and large amount (e.g. 5%) and make sure the chi-squared is near zero for the first and that ther log likelihood decreases. The point is that we havent really done full end to end tests tthat the likelihood function gives sensible results to build up confidence cluster modeling is ok. This should ue the simulator made in the test workspace.
5. Build on 3. to do an example where the simulated cluster has sources at different redshifts, and we only get the maximum likelihood soltion when the redshifts are right, and test therefore for that we must get redshifts right. For small changes in redshift (that are still multiplane) make sure we get small likelihood decreases.
6. Can you pair the example in 4 with fits which uses Imaging and light profiles (the values used in the simulator to make images with arfcs and again make sure the likelihood patterns hold.
7. Produce good visualization of 4 and 5 so we can see where the positions end up compared to the data and what the residuals of the light profiles for the cluster set up look like.
8. Now make scripts/jax_likelihood_functions/cluster/single_plane.py and multi_plane.py that has numerical assertions against it using the above test case.

The whole idea behind this task is to stress test cluster likelihood function as much as possible and look for bugs, edge cases and really anything that can (and will go wrong). So be persistent, be nit-picking and look for obvious issues with our implementation.

[Jammy2211]

Smoke Test Results — 2026-05-19

Workspace	Passed	Failed	Skipped	Total
autolens_workspace_test	11	0	1	12

All 11 smoke tests passed. 1 script skipped (database/scrape/general.py — NEEDS_FIX per no_run.yaml).

New cluster scripts (csv_api.py, simulator.py, visualization.py, likelihood_sanity.py) are not in smoke_tests.txt yet (deliberately — they are slow and involve JAX compilation). No existing smoke test regressed.

[Jammy2211]

Shipped (items 1-4) — 2026-05-19

PR

• feat(cluster): test-workspace scaffolding + likelihood sanity diagnostic #105

Summary

Shipped the cluster test-workspace dataset scaffolding (csv_api.py + simulator.py + visualization.py move) and the source-plane likelihood_sanity.py diagnostic. The full galaxy population — 2 main lenses + 1 extra galaxy + 1 host halo + 10 scaling members + 2 multi-plane sources — is built in plain Python by csv_api.py, written to the family CSVs, and consumed by the simulator. Items 5-8 (multi-redshift sensitivity, imaging-mode counterpart, paired visualization, JAX cluster likelihood functions) deferred to a follow-up issue.

Session Notes

• Headline finding from likelihood_sanity.py: cluster source-plane chi² (FitPositionsSource) is dominated by the PointSolver precision floor amplified by image-plane magnification (~100x at multi-image positions, ÷ 0.005" position noise → ~8e7 baseline). Perturbations below ~10% can produce sub-1% chi² variations in either direction. The chi² is not a reliable sensitivity probe at small parameter offsets. Worth investigating in the follow-up.
• FitPositionsImagePair skipped in likelihood_sanity.py. Each call triggers a PointSolver forward-solve; the full sweep (10 params × 7 epsilons × 2 datasets) timed out at 1500s. Added a RUN_IMAGE_PLANE = False toggle so the image-plane variant can be enabled when budget allows. Source-plane (FitPositionsSource) sweep runs in ~30s on CPU.
• Source-plane reference choice. Tried both profile=point_profile (truth Point centre as reference) and profile=None (barycenter of traced positions). Same precision-floor dominance under both — fundamental to the chi² formula residual² × magnification² / noise², not the reference choice.
• *_centres.json files dropped from the upstream cluster surface (PR #189 / cluster-csv-api). Updated visualization.py during the move to read centres from mass.csv via al.galaxy_models_from_csv.
• Scope cut at 4/8 deliverables vs the original 8. Decision: ship the dataset scaffolding + the diagnostic finding now; items 5-8 in a follow-up. The likelihood_sanity.py precision-floor finding alone is a meaningful contribution — pushing 5-8 in the same PR risked the same insensitivity blocking the multi-redshift / imaging tests too.

Follow-up

Filed below — items 5-8 plus an explicit "investigate the precision-floor finding" handle.