Skip to content

test: CI JAX coverage for PointMass/SMBH profiles - #252

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/pointmass-smbh-jax-ci
Aug 6, 2026
Merged

test: CI JAX coverage for PointMass/SMBH profiles#252
Jammy2211 merged 1 commit into
mainfrom
feature/pointmass-smbh-jax-ci

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Follow-up to PyAutoLabs/PyAutoGalaxy#553 (fixed in PyAutoLabs/PyAutoGalaxy#554): al.mp.PointMass / al.mp.SMBH broke 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.pymp.PointMass and mp.SMBH added to the profile-level JIT tests: deflections_yx_2d_from + potential_2d_from through the standard three-step pattern on Grid2DIrregular and Grid2D.uniform, plus a raw-zeros convergence_2d_from check (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 an SMBH whose centre/mass are FREE. A free traced mass is the only configuration that exercises the SMBH.__init__ mass→Einstein-radius conversion inside the jit trace (profiles_jit.py instantiates with floats and can never see it). Added to smoke_tests.txt so 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 to simulator/simple.py truth 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 env PYAUTO_TEST_MODE=2 PYAUTO_FAST_PLOTS=1 with ENV: jax full_datasets released — verified via build_env_for_script)
  • smoke_tests.txt (+1 entry)
  • scripts/CLAUDE.md (docs)

Test Plan

  • Control: imaging/jax_likelihood/lp.py reproduces its pinned literal in this environment before generating the new one
  • smbh.py vmap literal 1194.84699035 generated and re-run to a full pass (jit round-trip matches NumPy scalar to 13 significant figures)
  • smbh.py passes under the exact CI smoke env vars (25.9s)
  • profiles_jit.py full run passes (mp.PointMass OK, mp.SMBH OK)
  • Sensitivity verified: the literal changes when the SMBH mass changes (not a blind literal)

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.

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>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Aug 6, 2026
@Jammy2211 Jammy2211 closed this Aug 6, 2026
@Jammy2211 Jammy2211 reopened this Aug 6, 2026
@Jammy2211
Jammy2211 merged commit 9a631d3 into main Aug 6, 2026
@Jammy2211
Jammy2211 deleted the feature/pointmass-smbh-jax-ci branch August 6, 2026 18:45
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