feat: add ellipse jax_likelihood reference scripts

[Jammy2211]

Overview

Add the numpy-baseline reference scripts under scripts/jax_likelihood_functions/ellipse/ for ellipse likelihood fitting. Step 2 of 7 in the ellipse_fitting_jax feature decomposition (see PyAutoPrompt/z_features/ellipse_fitting_jax.md). The scripts print every component of AnalysisEllipse.fit_list_from(...).figure_of_merit on the numpy path so the JAX round-trip added in prompt 7 has a stable target to assert against with np.testing.assert_allclose(rtol=1e-4).

Plan

• Add scripts/jax_likelihood_functions/ellipse/__init__.py (empty package marker).
• Add a small dedicated simulator (50x50 pixels, single Sersic galaxy at 0.2"/px) writing into dataset/ellipse/jax_test/. Independent of the imaging jax_test simulator so the ellipse reference numbers don't drift if imaging/simulator.py is later retuned.
• Add fit.py: load (or auto-simulate) the dataset, build a single fixed Ellipse model, instantiate AnalysisEllipse(dataset=...), print log_likelihood, chi_squared, noise_normalization, figure_of_merit per FitEllipse. Leave a clearly-marked TODO(7_analysis_ellipse_jax.md) block at the end where the JIT round-trip will go.
• Add multipoles.py: same shape as fit.py but adds an EllipseMultipole(m=4) per ellipse. Locks the multipole code path which has its own JAX-incompatible while loops that prompt 5 will rewrite.
• All four scripts numpy-only — import jax must not appear anywhere.
• The scripts run cleanly under bash run_all_scripts.sh and the printed numbers are reproducible to ~1e-6.

Detailed implementation plan

Affected Repositories

• autogalaxy_workspace_test (primary)
• PyAutoGalaxy (read-only, for understanding AnalysisEllipse)
• autogalaxy_workspace (read-only, for scripts/ellipse/modeling.py composition shape)

Work Classification

Workspace

Branch Survey

Repository	Current Branch	Dirty?
./autogalaxy_workspace_test	main	clean

Suggested branch: feature/ellipse-jax-likelihood-tests
Worktree root: ~/Code/PyAutoLabs-wt/ellipse-jax-likelihood-tests/ (created later by /start_workspace)

Implementation Steps

1. Create scripts/jax_likelihood_functions/ellipse/__init__.py (empty).
2. Create scripts/jax_likelihood_functions/ellipse/simulator.py. Patterned on scripts/jax_likelihood_functions/imaging/simulator.py but trimmed: Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2), single Galaxy with a Sersic bulge centered at (0, 0), noise_seed=1 for reproducibility, writes data.fits, psf.fits (ignored by ellipse but kept for Imaging.from_fits symmetry), noise_map.fits to dataset/ellipse/jax_test/.
3. Create scripts/jax_likelihood_functions/ellipse/fit.py:
   • Auto-simulate the dataset via subprocess.run([sys.executable, "scripts/jax_likelihood_functions/ellipse/simulator.py"]) when data.fits is missing.
   • ag.Imaging.from_fits(...) then apply a generous circular mask (radius=3.0).
   • Compose: ellipse = af.Model(ag.Ellipse) with centre.{0,1} uniform priors, ell_comps.{0,1} uniform priors, major_axis=0.5 fixed. Wrap as model = af.Collection(ellipses=af.Collection(ellipse_0=ellipse)) to match the prior used in the visualization test.
   • analysis = ag.AnalysisEllipse(dataset=dataset) (today defaults to use_jax=False).
   • instance = model.instance_from_prior_medians(), then fit_list = analysis.fit_list_from(instance=instance). Print log_likelihood, chi_squared, noise_normalization, figure_of_merit per fit and the aggregate.
   • Final placeholder block: # TODO(7_analysis_ellipse_jax.md): jax.jit(analysis.fit_from) round-trip + np.testing.assert_allclose(rtol=1e-4) against numpy reference above.
4. Create scripts/jax_likelihood_functions/ellipse/multipoles.py: identical to fit.py but the model has multipoles=af.Collection(ellipse_0=af.Collection(m4=af.Model(ag.EllipseMultipole, m=4, multipole_comps=(0.05, 0.0)))). Same TODO placeholder.
5. Use the workspace docstring style ("""...""" blocks, __Section Name__ headers). No # comments for prose.
6. Numpy-only — no import jax, no xp=jnp. The script imports must match fit.py from the visualization PR: import autofit as af, import autogalaxy as ag, plus numpy as np only if needed for printing reference numbers.

Testing Approach

• Local: python scripts/jax_likelihood_functions/ellipse/simulator.py then python scripts/jax_likelihood_functions/ellipse/fit.py and python scripts/jax_likelihood_functions/ellipse/multipoles.py. All three exit 0 and print stable numbers.
• Reproducibility: run fit.py twice, confirm chi_squared matches byte-for-byte (numpy simulator uses noise_seed=1).
• Smoke: bash run_all_scripts.sh picks the new scripts up automatically (no smoke_tests.txt registration needed yet; mirroring the imaging precedent which is also in smoke_tests.txt — flag in PR if user wants it added).

Key Files

• scripts/jax_likelihood_functions/ellipse/__init__.py — new.
• scripts/jax_likelihood_functions/ellipse/simulator.py — new, dedicated small-grid Sersic dataset.
• scripts/jax_likelihood_functions/ellipse/fit.py — new, numpy reference + TODO placeholder.
• scripts/jax_likelihood_functions/ellipse/multipoles.py — new, numpy reference for the multipole path + TODO placeholder.
• scripts/jax_likelihood_functions/imaging/lp.py — read-only reference for the imaging pattern this mirrors.

Original Prompt

Click to expand starting prompt

Step 2 of the ellipse-JAX series. Step 1 (1_workspace_visualization.md) added the visualization integration test. This prompt adds the likelihood-function integration tests, still on the numpy path. They lock in the reference numbers we'll later assert against once the JAX path lands in 7_analysis_ellipse_jax.md.

The pattern to mirror is @autogalaxy_workspace_test/scripts/jax_likelihood_functions/imaging/lp.py, which (a) auto-simulates a small dataset on first run, (b) builds a model + analysis, (c) computes a baseline log_likelihood on the numpy path, and (d) compares the JIT path with np.testing.assert_allclose(rtol=1e-4). For this prompt, only build steps (a)-(c) — leave a # TODO(7_analysis_ellipse_jax.md) placeholder where the JIT comparison will go.

Please:

1. Add @autogalaxy_workspace_test/scripts/jax_likelihood_functions/ellipse/__init__.py.

2. Add @autogalaxy_workspace_test/scripts/jax_likelihood_functions/ellipse/simulator.py modelled on @autogalaxy_workspace_test/scripts/jax_likelihood_functions/imaging/simulator.py — small grid (e.g. shape_native=(50, 50), pixel_scales=0.2), single Sersic galaxy, written into dataset/ellipse/jax_test/.

3. Add @autogalaxy_workspace_test/scripts/jax_likelihood_functions/ellipse/fit.py:

   • Load (or auto-simulate via simulator.py) the dataset.
   • Build an af.Collection(ellipses=af.Collection(ellipse_0=af.Model(ag.Ellipse, major_axis=...)))-style model (consult @autogalaxy_workspace/scripts/ellipse/modeling.py for the exact composition shape).
   • Build analysis = ag.AnalysisEllipse(dataset=dataset) (today this defaults to use_jax=False).
   • Compute fit_np = analysis.fit_list_from(instance=model.instance_from_prior_medians()) and print every component — log_likelihood, chi_squared, noise_normalization, figure_of_merit — to capture the reference numbers.
   • Leave a # TODO(7_analysis_ellipse_jax.md): jax.jit(analysis.fit_from) round-trip placeholder block at the end.

4. Add @autogalaxy_workspace_test/scripts/jax_likelihood_functions/ellipse/multipoles.py: same as fit.py but with multipole_list=[ag.EllipseMultipole(m=4, multipole_comps=(0.05, 0.0))] per ellipse. This locks in the multipole code path which has its own JAX-incompatible while loops in EllipseMultipole.get_shape_angle (handled in prompt 5).

5. Use the workspace docstring style throughout (see prompt 1 for reference).

6. Test bar: both scripts run cleanly under bash run_all_scripts.sh and the printed reference numbers are stable to ~1e-6 across runs. No JAX imports yet — import jax should not appear in any of these files.

[Jammy2211]

Smoke Test Results — 2026-05-10

Workspace	Passed	Failed	Skipped	Total
autofit_workspace	7	0	0	7
autogalaxy_workspace	6	0	0	6
autolens_workspace	7	0	0	7
autolens_workspace_test	11	0	1	12
euclid_strong_lens_modeling_pipeline	6	0	0	6
HowToLens	6	0	0	6

All 41 smoke tests passed across 6 workspaces (1 skipped: database/scrape/general in autolens_workspace_test — NEEDS_FIX 2026-04-27, pre-existing).

[Jammy2211]

Shipped — 2026-05-10

PRs

• feat: add numpy ellipse jax_likelihood reference scripts #42 (merged)

Summary

Added scripts/jax_likelihood_functions/ellipse/ with simulator.py, fit.py, and multipoles.py — all numpy path only. Captures the reference numbers (log_likelihood, chi_squared, noise_normalization, figure_of_merit) for a fixed-parameter single-ellipse fit and the multipole-perturbed variant. Step 2 of 7 in z_features/ellipse_fitting_jax.md.

Reference numbers locked in

Component	fit.py	multipoles.py
log_likelihood	-382.22660947	-725.51150750
chi_squared	764.45321894	1451.02301500
noise_normalization	-27.02295060	-26.84755137
figure_of_merit	-368.71513417	-712.08773182

Session Notes

• The multipoles=[[multipole]] list-of-lists composition shape was the one that worked. AnalysisEllipse.fit_list_from (PyAutoGalaxy/autogalaxy/ellipse/model/analysis.py:121) does instance.multipoles[i] — for a single ellipse with one multipole that means the outer collection is per-ellipse and the inner list is the multipole list. Documented inline so prompt 7's author knows the shape.
• The simulator is dataset-independent of imaging/jax_test/ on purpose: we don't want a future tweak to imaging/simulator.py (e.g. changing the noise seed for an unrelated reason) to shift the ellipse reference numbers and force a downstream rewrite of prompt 7's tolerances.
• Smoke suite ran 41/41 green (1 pre-existing skip in autolens_workspace_test/database/scrape/general). Run delegated to Sonnet per CLAUDE.md model-split rule for workspace_test developer scripts.