feat: cluster test-workspace likelihood scripts (items 5-8 follow-up to #104)

[Jammy2211]

Overview

Follow-up to #104 (PR #105). Ships items 5-8 of the original eight-deliverable cluster test-workspace plan:

5. scripts/cluster/likelihood_redshift_sensitivity.py — perturb source redshifts δz ∈ {±0.01, ±0.05, ±0.1} (small multi-plane displacements) and assert the truth redshift maximises log-likelihood. Small δz produces small likelihood decreases.
6. scripts/cluster/likelihood_imaging.py — same perturbation tests but with al.AnalysisImaging + light profiles loaded from light.csv. Assert the imaging-mode log-likelihood pattern matches the point-source pattern.
7. Paired visualization for (5) and (6): chi²-vs-perturbation curves, image-plane positions overlaid on data, 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 with numerical assertions referencing the perturbation results from feat: cluster test workspace — likelihood sanity + JAX likelihood functions #104. Modelled on the existing scripts/jax_likelihood_functions/point_source/{image_plane,source_plane,point}.py.

PR #105 shipped the dataset scaffolding (csv_api.py + simulator.py + visualization.py move) and the source-plane likelihood_sanity.py diagnostic — the foundation these four scripts build on.

Investigate alongside (from PR #105's surprise)

The likelihood_sanity.py diagnostic shipped in PR #105 surfaced a real finding worth a closer look while writing the redshift / imaging variants:

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.

Two possibilities worth ruling out during this follow-up:

• PointSolver precision sets the floor. Tighter pixel_scale_precision (e.g. 1e-4) should pull the floor down. If chi² becomes more sensitive at 0.1% perturbations, the issue is just precision and we should document the recommended precision for sensitivity studies.
• Magnification computation is bugged. If magnifications_at_positions in FitPositionsSource is computed from the data positions and not the model, perturbations to the lens model won't propagate into the chi² weighting. Worth checking the cache path in AbstractFitPositions.

The redshift / imaging-mode scripts in this follow-up may surface the same insensitivity or a different pathology — either is informative.

Plan

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 (assumed at /start_workspace time)

Suggested branch: feature/cluster-test-likelihood-redshift-imaging-jax

Implementation Steps

1. scripts/cluster/likelihood_redshift_sensitivity.py — load the truth tracer from mass.csv + light.csv + point.csv, perturb each source's redshift via dataset.redshift = z_truth + δz (constructed in a fresh PointDataset copy to avoid mutating the on-disk CSV). Re-evaluate FitPositionsSource.chi_squared at every (source, δz) pair. Assert truth redshift maximises log-likelihood; small δz → small likelihood decreases.
2. scripts/cluster/likelihood_imaging.py — al.AnalysisImaging + light profiles from light.csv produce arcs in the truth dataset. Perturb each mass param on the dPIE / NFW tiers (same ε grid as likelihood_sanity.py); evaluate analysis.log_likelihood_function(instance) against the truth data.fits. Assert imaging-mode log-likelihood maximises at truth.
3. Visualization — chi²-vs-perturbation matplotlib panels per script; image-plane positions overlaid on data.fits at the worst perturbation; residual maps for the imaging fit. Save into dataset/cluster/test/visualization/.
4. JAX cluster likelihood functions — scripts/jax_likelihood_functions/cluster/{__init__.py, single_plane.py, multi_plane.py}. Pattern after point_source/{image_plane,source_plane,point}.py — register tracer pytrees, jit-wrap the likelihood evaluation, run numerical assertions against the truth ε=0 chi².

Key Files

• New: scripts/cluster/{likelihood_redshift_sensitivity.py, likelihood_imaging.py}
• New: scripts/cluster/likelihood_imaging_visualization.py (or add viz to the existing scripts)
• New: scripts/jax_likelihood_functions/cluster/{__init__.py, single_plane.py, multi_plane.py}
• Reference (read-only): PR feat(cluster): test-workspace scaffolding + likelihood sanity diagnostic #105's scripts/cluster/likelihood_sanity.py, scripts/cluster/simulator.py
• Reference (read-only): scripts/jax_likelihood_functions/point_source/{image_plane,source_plane,point}.py

Open Questions

• Should redshift perturbation use the per-PointDataset redshift attribute, or perturb the source Galaxy.redshift (which affects the Tracer's ray-tracing)? The two are nominally pinned together — verify both paths give consistent results before locking the test.
• Does AnalysisImaging even support multi-plane source redshifts? The cluster has 2 sources at z=1.0 and z=2.0 — Tracer handles this natively but the Analysis layer may need a fitting-grid override.
• Smoke-test inclusion: add the new scripts to smoke_tests.txt once they pass under PYAUTO_TEST_MODE=2. The likelihood scripts are pure-Python (no Nautilus) so they should smoke quickly.

Stress-Test Mandate (carried over from #104)

"Be persistent, be nit-picking, look for obvious issues with our implementation." Bugs discovered while writing these scripts become their own follow-up issues; the test scaffolding ships regardless.