Skip to content

feat: add numpy ellipse jax_likelihood reference scripts#42

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/ellipse-jax-likelihood-tests
May 10, 2026
Merged

feat: add numpy ellipse jax_likelihood reference scripts#42
Jammy2211 merged 1 commit into
mainfrom
feature/ellipse-jax-likelihood-tests

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Add the numpy-baseline reference scripts under scripts/jax_likelihood_functions/ellipse/ that lock in the reference numbers for ellipse fitting before the JAX path lands. Step 2 of 7 in the ellipse_fitting_jax feature decomposition (PyAutoPrompt/z_features/ellipse_fitting_jax.md).

Each script auto-simulates the dataset on first run, builds a fixed-parameter ellipse model (plus multipole for multipoles.py), and prints every component of every FitEllipse produced by AnalysisEllipse.fit_list_from(...). Prompt 7 of the feature will JIT-wrap analysis.fit_from and assert against these numbers with np.testing.assert_allclose(rtol=1e-4).

Numpy-only — no import jax anywhere in the new scripts. Both fit.py and multipoles.py are deterministic across runs (verified by diff-comparison of two runs).

Reference numbers captured at the prior medians of the model:

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

Issue #41.

Scripts Changed

  • scripts/jax_likelihood_functions/ellipse/__init__.py — new empty package marker.
  • scripts/jax_likelihood_functions/ellipse/simulator.py — new; simulates a 50×50 single-Sersic-galaxy Imaging dataset (noise_seed=1) into dataset/ellipse/jax_test/. Independent of imaging/simulator.py so the ellipse reference numbers don't drift with future imaging-side retuning.
  • scripts/jax_likelihood_functions/ellipse/fit.py — new; single-ellipse numpy reference, prints the four FOM components per fit, ends with a # TODO(7_analysis_ellipse_jax.md) placeholder block.
  • scripts/jax_likelihood_functions/ellipse/multipoles.py — new; same shape but with EllipseMultipole(m=4, multipole_comps=(0.05, 0.0)) per ellipse. Uses the multipoles=[[multipole]] list-of-lists composition shape so AnalysisEllipse.fit_list_from's instance.multipoles[i] access works correctly.

Test Plan

  • Smoke tests pass for autogalaxy_workspace_test
  • Both fit.py and multipoles.py re-runs produce byte-identical reference numbers

🤖 Generated with Claude Code

Adds scripts/jax_likelihood_functions/ellipse/ with simulator.py,
fit.py, and multipoles.py — numpy path only. Captures the
log_likelihood / chi_squared / noise_normalization / figure_of_merit
reference numbers that prompt 7 of the ellipse_fitting_jax feature
will assert against once AnalysisEllipse gains the JAX path.

Issue #41.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label May 10, 2026
@Jammy2211 Jammy2211 merged commit 25debdc into main May 10, 2026
4 checks passed
@Jammy2211 Jammy2211 deleted the feature/ellipse-jax-likelihood-tests branch May 10, 2026 20:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant