test: EllipseMultipoleScaled JAX parity script#51
Merged
Conversation
Mirrors the existing multipoles.py (EllipseMultipole coverage) for the scaled variant: numpy reference + fitness._vmap batch + jax.jit(fit_from) round-trip with rtol=1e-4 assertion against the numpy total_log_likelihood. This closes the variant-coverage gap that allowed PyAutoGalaxy#426 to slip through the prompt-7 verification — going forward any regression in EllipseMultipoleScaled's JAX path gets caught by smoke. Caveat: with multipole.major_axis=1.0 the scaling factor is a no-op (k = k_scaled * 1.0), so the numerical reference values match the non-scaled multipoles.py. The script still exercises the EllipseMultipoleScaled.points_perturbed_from code path including the deferred k_scaled * major_axis derivation under vmap/jit — that's the code path that was broken before PR #427. Adding a non-trivial major_axis case is a future enhancement. Issue PyAutoGalaxy#426. Upstream library PR: PyAutoGalaxy#427. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
scripts/jax_likelihood_functions/ellipse/multipoles_scaled.py— the JAX parity test forEllipseMultipoleScaled, mirroring the existingmultipoles.py(EllipseMultipole). This closes the variant-coverage gap that allowed PyAutoGalaxy#426 to slip through prompt-7's verification:EllipseMultipoleScaledhad novmap-validation script, so its__init__-timenp.*call on JAX tracers stayed invisible to smoke until a user HPC job hit it.The new script runs numpy reference →
fitness._vmapbatch evaluation →jax.jit(analysis.fit_from)(instance)round-trip withnp.testing.assert_allclose(rtol=1e-4)against the numpytotal_log_likelihood. Same shape asmultipoles.pyso any future regression in either variant gets caught.Caveat for reviewers: with
multipole.major_axis=1.0the scaling factork = k_scaled * 1.0is mathematically a no-op, so the numerical reference values match the non-scaledmultipoles.py. The script still exercisesEllipseMultipoleScaled.points_perturbed_from(including the deferredk_scaled * major_axisderivation under vmap/jit) — that's the code path that was broken before PR #427. Strengthening with a non-trivialmajor_axiscase is a future enhancement.Scripts Changed
scripts/jax_likelihood_functions/ellipse/multipoles_scaled.py— new (~160 lines). Mirrorsmultipoles.pyline-by-line but builds the multipole asaf.Model(ag.EllipseMultipoleScaled)withscaled_multipole_compspriors andmajor_axis=1.0. Numpy reference, vmap block (50-batch throughFitness._vmap), and JIT round-trip block (jax.jit(analysis_jit.fit_from)(instance)+assert_allclose(rtol=1e-4)+isinstance(fit.log_likelihood, jnp.ndarray)).Upstream PR
PyAutoLabs/PyAutoGalaxy#427
Test Plan
python scripts/jax_likelihood_functions/ellipse/multipoles_scaled.pyprintsPASS: jit(fit_from) round-trip matches NumPy scalar.🤖 Generated with Claude Code