test: CI JAX coverage for PointMass/SMBH profiles - #252
Merged
Conversation
Follow-up to PyAutoGalaxy#553/#554. profiles_jit.py gains PointMass and SMBH (deflections + potential three-step pattern, raw-zeros convergence check); new imaging/jax_likelihood/smbh.py exercises a FREE traced SMBH mass through the full vmap+jit likelihood pipeline and joins smoke_tests.txt. Model is truth-anchored because at prior medians the positive-only solver zeroes the source, making the literal blind to source-plane mass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Aug 6, 2026
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
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
Follow-up to PyAutoLabs/PyAutoGalaxy#553 (fixed in PyAutoLabs/PyAutoGalaxy#554):
al.mp.PointMass/al.mp.SMBHbroke every JAX fit and no CI exercised the JAX path for point-mass profiles. This adds coverage for both bug shapes:scripts/misc/profiles_jit.py—mp.PointMassandmp.SMBHadded to the profile-level JIT tests:deflections_yx_2d_from+potential_2d_fromthrough the standard three-step pattern onGrid2DIrregularandGrid2D.uniform, plus a raw-zerosconvergence_2d_fromcheck (undecorated by design). Catches the ArrayIrregular-deflections bug shape.scripts/imaging/jax_likelihood/smbh.py(new) — full-pipeline vmap+jit likelihood with the lens hosting anSMBHwhosecentre/massare FREE. A free traced mass is the only configuration that exercises theSMBH.__init__mass→Einstein-radius conversion inside the jit trace (profiles_jit.pyinstantiates with floats and can never see it). Added tosmoke_tests.txtso the PR gate runs it (25s measured under the smoke env, vs the 300s cap).scripts/CLAUDE.md— coverage docs updated.Design note (why the model is truth-anchored)
At this workspace's default prior medians (
effective_radius=15.0,einstein_radius=4.0) the positive-only linear solver zeroes the source's solved intensity, making the likelihood bit-identical for any source-plane mass structure — verified: Isothermal θ_E 4.0 vs 2.0 and SMBH mass 1e9 vs 1e11 all give LL −352316.6109446431. A vmap literal generated there would not pin the SMBH at all. The new script therefore pins the non-SMBH components tosimulator/simple.pytruth values, where the source is retained and the literal responds to the SMBH (LL 1187.6 without SMBH → 1194.8 at median mass → −42.6 at 1e11). Note this zeroed-source property also means the existing jax_likelihood literals at prior medians constrain source-plane deflections only weakly — possibly worth a separate look.Scripts Changed
scripts/misc/profiles_jit.py(extended; full run passes locally, all profiles OK)scripts/imaging/jax_likelihood/smbh.py(new; passes plain and under the exact smoke envPYAUTO_TEST_MODE=2 PYAUTO_FAST_PLOTS=1withENV: jax full_datasetsreleased — verified viabuild_env_for_script)smoke_tests.txt(+1 entry)scripts/CLAUDE.md(docs)Test Plan
imaging/jax_likelihood/lp.pyreproduces its pinned literal in this environment before generating the new onesmbh.pyvmap literal 1194.84699035 generated and re-run to a full pass (jit round-trip matches NumPy scalar to 13 significant figures)smbh.pypasses under the exact CI smoke env vars (25.9s)profiles_jit.pyfull run passes (mp.PointMass OK,mp.SMBH OK)Heart gate: RED at ship time with the same single unrelated reason (
release validation FAILED (stage integrate), in-flight release arc) explicitly human-acknowledged for the PyAutoGalaxy#554 ship earlier today; this ack covers that reason set only.Generated by the PyAutoLabs agent workflow.