fix: make PointMass and SMBH JAX-compatible

[Jammy2211]

Overview

al.mp.PointMass and al.mp.SMBH fail every JAX-mode imaging fit (reported by a user on 2026.8.4.1; reproduced on current main). PointMass raises TypeError: multiply requires ndarray or scalar arguments, got ArrayIrregular inside deflections_yx_2d_from; SMBH raises jax.errors.TracerArrayConversionError from np.sqrt on a traced mass in __init__. Root causes are diagnosed and the fixes were validated end-to-end (jitted AnalysisImaging.log_likelihood_function with subclassed fixed profiles returns a finite likelihood; stock profiles reproduce both errors as controls; fixed deflections match stock numpy deflections to ~1e-18).

Plan

• Rewrite PointMass.deflections_yx_2d_from with raw xp ops instead of the legacy radial_grid_from → _cartesian_grid_via_radial_from helpers, which return an ArrayIrregular wrapper under xp=jnp on irregular PSF-evaluation grids.
• Simplify PointMass.convergence_2d_from to xp.zeros (its central-pixel np.argmin code is dead — the assignment is commented out) so it is tracer-safe.
• Make SMBH.__init__ tracer-safe: (mass_angular / np.pi) ** 0.5 instead of np.sqrt(...).
• Validate with the full test_autogalaxy/ suite (NumPy parity) and jitted-likelihood control scripts (JAX path).

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoGalaxy	main	clean

Suggested branch: feature/pointmass-smbh-jax
Worktree root: ~/Code/PyAutoLabs-wt/pointmass-smbh-jax/ (created by /start_library)
Work Classification: Library

Implementation Steps

1. autogalaxy/profiles/mass/point/point.py — deflections_yx_2d_from: keep @aa.decorators.to_vector_yx + @aa.decorators.transform; replace the body with raw xp ops:
   y = xp.asarray(grid.array[:, 0]), x = xp.asarray(grid.array[:, 1]),
   alpha = self.einstein_radius**2 / (y**2 + x**2 + 1e-20),
   return xp.stack((alpha * y, alpha * x), axis=-1).
   This mirrors how IsothermalSph.deflections_yx_2d_from (total/isothermal.py) already passes raw arrays and works under JAX.
2. autogalaxy/profiles/mass/point/point.py — convergence_2d_from: return xp.zeros(grid.shape[0]); delete the dead np.argmin central-pixel computation (its assignment has been commented out for a long time). Keep the docstring explaining the Dirac-delta convention.
3. autogalaxy/profiles/mass/point/smbh.py — __init__: einstein_radius = (mass_angular / np.pi) ** 0.5 (Python pow is tracer-safe; the astropy critical_surface_density_between_redshifts_from call is fine because redshifts stay floats).
4. Run test_autogalaxy/ (unit tests are NumPy-only by repo policy — no JAX in unit tests).
5. JAX validation: jitted AnalysisImaging.log_likelihood_function with stock PointMass (einstein_radius free) and stock SMBH (mass free) in the model must return finite likelihoods; both previously failed with the two user-reported errors. Numerical parity of new deflections vs the old numpy implementation.
6. At ship time, decide via workspace-impact analysis whether to add PointMass/SMBH to an autogalaxy_workspace_test JAX parity script so CI covers this regression class.

Key Files

• autogalaxy/profiles/mass/point/point.py — PointMass.deflections_yx_2d_from, convergence_2d_from
• autogalaxy/profiles/mass/point/smbh.py — SMBH.__init__
• autogalaxy/profiles/mass/total/isothermal.py — reference pattern (no change)

Notes

• PowerLawSph(slope=3.0) is not a workaround for users: the (3 - slope) normalisation is NaN at exactly 3.
• PointMass.potential_2d_from already uses raw xp ops — no change needed.

Original Prompt

Click to expand starting prompt

Prompt file: PyAutoMind/bug/autogalaxy/pointmass_smbh_jax_incompatibility.md

User support report (2026.8.4.1, JAX GPU fit on cosma): modelling a central black hole in a lens galaxy — al.mp.PointMass fails with TypeError: multiply requires ndarray or scalar arguments, got ArrayIrregular (full Nautilus traceback through AnalysisImaging.log_likelihood_function → Tracer.blurred_image_2d_from → PointMass.deflections_yx_2d_from → _cartesian_grid_via_radial_from); al.mp.SMBH fails with jax.errors.TracerArrayConversionError.

Maintainer instruction: "ok do the work on the source code, the user can wait until we do a new release so I wont send them the code above"

[Jammy2211]

Library PR Created

PR: #554 (branch feature/pointmass-smbh-jax, commit b68df8b, label pending-release verified).

Validation: test_autogalaxy/ 1017 passed · PyAutoLens test_autolens/point/ 115 passed against this branch · jitted AnalysisImaging.log_likelihood_function with stock PointMass (θ_E free) and stock SMBH (mass free) returns finite likelihoods (both reproduced the user-reported errors on unfixed main as controls) · deflections parity vs analytic formula 8.7e-19 · jax.grad finite for all point-mass parameters.

Heart gate: verdict RED at ship time — red_reasons: ['release validation FAILED (stage integrate)'] plus non-red workspace-validation failures (autolens interferometer likelihood_function script/notebook, cloud#30938311069) and 1 tenant-firewall manifest mismatch. All pre-existing and unrelated to this branch (in-flight release-push arc). Human explicitly authorized PR-open despite RED (2026-08-06); merge/release remain separate human decisions.

Workspace impact: API changes none — internal bug fix, numpy path numerically identical. Workspace grep: only autolens_workspace/scripts/guides/profiles/mass.py, autolens_workspace_test/scripts/misc/mass/point.py, and prior-config YAMLs reference these profiles; no migration needed (option iii). Six-workspace smoke run in flight against this branch — results to follow.

[Jammy2211]

Shipped

PR #554 merged to main (2026-08-06, commit b68df8b). al.mp.PointMass and al.mp.SMBH now work in JAX-mode fits; the fix reaches users in the next release.

Notes:

• Merged on explicit human instruction ahead of PR CI completion (docs + 3.12/3.13 unittest legs were pending, not failing) and with the six-workspace smoke run cut short — the only smoke failures observed were parallel-run race artifacts (missing auto-simulated datasets from shared state), the known false-failure mode, none serially confirmed. Local validation on the branch: test_autogalaxy/ 1017 passed, PyAutoLens test_autolens/point/ 115 passed, jitted-likelihood + gradient controls green.
• Heart was RED at ship time for the unrelated in-flight release-validation arc; PR-open and merge were both explicitly human-authorized.