test: CI JAX coverage for PointMass/SMBH profiles

[Jammy2211]

Overview

Follow-up to PyAutoLabs/PyAutoGalaxy#553 (fixed by PyAutoLabs/PyAutoGalaxy#554): al.mp.PointMass and al.mp.SMBH broke every JAX-mode fit, and nothing in CI exercised the JAX path for point-mass profiles, so the regression class was invisible. Add coverage for both bug shapes.

Plan

• Add mp.PointMass and mp.SMBH to scripts/misc/profiles_jit.py (profile-level JIT: deflections_yx_2d_from + convergence_2d_from on Grid2DIrregular and Grid2D.uniform, three-step pattern) — catches the ArrayIrregular deflections bug shape.
• New scripts/imaging/jax_likelihood/smbh.py — full-pipeline vmap+jit likelihood with lens = Isothermal + SMBH where mass is a free parameter (a free traced mass is the only way to exercise the SMBH.__init__ bug shape; profiles_jit.py instantiates with floats). Mirrors imaging/jax_likelihood/lp.py; added to smoke_tests.txt so the PR CI gate runs it.
• Update scripts/CLAUDE.md coverage docs.

Detailed implementation plan

Affected Repositories

• autolens_workspace_test (primary)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_test	main	clean

Suggested branch: feature/pointmass-smbh-jax-ci
Worktree root: ~/Code/PyAutoLabs-wt/pointmass-smbh-jax-ci/
Work Classification: Workspace

Implementation Steps

1. scripts/misc/profiles_jit.py — extend the mass-profile list with mp.PointMass(centre=(0.01, 0.02), einstein_radius=0.1) and mp.SMBH(centre=(0.01, 0.02), mass=1e10); both grids, all three JAX steps.
2. scripts/imaging/jax_likelihood/smbh.py — mirror lp.py: auto-sim jax_test dataset, mask, model with lens Isothermal + SMBH (mass LogUniformPrior(1e8, 1e11), redshifts fixed 0.5/1.0) + Sersic source, Fitness + _vmap batch with hardcoded expected literal, Path A jax.jit(analysis.fit_from) round-trip vs numpy scalar. __Env__ section ENV: jax full_datasets.
3. Generate the vmap literal by running locally under the CI env tokens; run lp.py first as a control to confirm the env reproduces its pinned literal (-1.34797827e09).
4. Add imaging/jax_likelihood/smbh.py to smoke_tests.txt.
5. Update scripts/CLAUDE.md (jax_likelihood table + profiles_jit profile list).

Key Files

• scripts/misc/profiles_jit.py
• scripts/imaging/jax_likelihood/smbh.py (new)
• smoke_tests.txt
• scripts/CLAUDE.md

Original Prompt

Click to expand starting prompt

Prompt file: PyAutoMind/test/autolens_workspace_test/pointmass_smbh_jax_ci_coverage.md

Maintainer instruction: "ok yes do a CI exercise in autolens_workspace_test with JAX on these profiles as a quick follow up"

[Jammy2211]

Workspace PR Created

PR: #252 (branch feature/pointmass-smbh-jax-ci, commit bb60ecc, pending-release verified). The PR smoke gate now runs the new imaging/jax_likelihood/smbh.py on both matrix legs.

Validation: control lp.py reproduced its pinned literal before the new literal was generated · smbh.py passes plain and under the exact smoke env (25.9s vs 300s cap; __Env__ release proven via build_env_for_script) · profiles_jit.py full run passes with mp.PointMass OK / mp.SMBH OK · literal sensitivity verified (changes with SMBH mass).

Finding worth flagging: 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 imaging likelihood bit-identical for any source-plane mass structure (verified: Isothermal θ_E 4.0 vs 2.0 and SMBH 1e9 vs 1e11 all → LL −352316.6109446431, reconstruction [.., 0.0]). The new script pins its non-SMBH components to simulator/simple.py truth for this reason. Consequence: the existing jax_likelihood literals generated at prior medians constrain source-plane deflections only weakly — they pin the pipeline numerically but a mass-model regression would not move them. Possibly worth a separate prompt.

Heart gate: RED with the same single unrelated reason (release validation FAILED (stage integrate)) explicitly human-acknowledged for the PyAutoGalaxy#554 ship earlier today; PR-open proceeded under that same reason-set ack.

[Jammy2211]

Shipped

PR #252 merged to main (2026-08-06, commit bb60ecc). CI now exercises the JAX path for PointMass/SMBH: profiles_jit.py covers the profile level (full sweep / release profile) and imaging/jax_likelihood/smbh.py covers the full pipeline with a free traced mass in the PR smoke gate.

Merged on explicit human instruction on the strength of local validation ("if it runs ok locally then merge") because GitHub Actions was in a partial outage and never created runs for this PR (close/reopen re-trigger included). Local evidence: new script passes plain and under the exact CI smoke env (25.9s vs the 300s cap, env-token release proven via build_env_for_script); profiles_jit.py full run passes; the lp.py control reproduced its pinned literal first. The main-branch smoke run will provide the first CI execution once Actions recovers.