test: CI JAX coverage for PointMass/SMBH profiles

scripts/CLAUDE.md

@@ -153,6 +153,7 @@ likelihoods of the same base name.
 | Script | Model type |
 |---|---|
 | `imaging/jax_likelihood/lp.py` | Light parametric (Sersic, Exponential) |
+| `imaging/jax_likelihood/smbh.py` | Central `SMBH` point mass with FREE (traced) `mass` — regression cover for PyAutoGalaxy#553; non-SMBH components pinned to simulator truth because at prior medians the positive-only solver zeroes the source and the literal goes blind to source-plane mass |
 | `imaging/jax_likelihood/mge.py` | Multi-Gaussian expansion |
 | `imaging/jax_likelihood/delaunay.py` | Delaunay pixelization |
 | `imaging/jax_likelihood/rectangular.py` | Rectangular pixelization |
@@ -207,6 +208,11 @@ targets the methods that are called internally by `LensCalc` and `Tracer`.
 **Mass profiles**: `mp.Isothermal`, `mp.PowerLaw`, `mp.NFW`, `mp.ExternalShear`, `mp.ExternalPotential`
 → `deflections_yx_2d_from`, `convergence_2d_from`
 
+**Point-mass profiles**: `mp.PointMass`, `mp.SMBH` (regression cover for PyAutoGalaxy#553)
+→ `deflections_yx_2d_from`, `potential_2d_from`, plus a raw-zeros `convergence_2d_from` check
+(their convergence is undecorated and returns a raw zeros array by design). The traced-mass half
+of #553 needs a free model parameter and lives in `imaging/jax_likelihood/smbh.py`.
+
 Each method is tested on both `Grid2DIrregular` and `Grid2D.uniform`.
 All three steps of the JAX pattern are applied.  NFW uses `rtol=1e-4` (looser) due
 to its analytic JAX implementation.

scripts/imaging/jax_likelihood/smbh.py

@@ -0,0 +1,268 @@
+"""
+Func Grad: SMBH Point Mass
+==========================
+
+This script tests if JAX can successfully compute the log likelihood of an `Imaging` dataset with a
+model whose lens galaxy contains a central supermassive black hole, modelled as an `SMBH` point-mass
+profile with its `mass` a free parameter.
+
+ __SMBH Fitting__
+
+A central black hole in the lens galaxy is modelled with the `SMBH` profile (a `PointMass` whose
+Einstein radius is derived from a physical mass and the lens/source redshifts). Two regressions
+previously broke this under JAX (PyAutoGalaxy#553):
+
+- `PointMass.deflections_yx_2d_from` produced an `ArrayIrregular` wrapper that `jnp.multiply`
+  rejects on the irregular PSF-evaluation grids every imaging fit uses.
+
+- `SMBH.__init__` converted mass to Einstein radius with `np.sqrt`, which raises
+  `TracerArrayConversionError` when `mass` is a free (traced) parameter.
+
+The second bug is only reachable when `mass` is free, because the model instance is then built
+inside the jit trace — which is why this script keeps `mass` free rather than fixed. The
+profile-level half of this coverage lives in `misc/profiles_jit.py`.
+
+__Env__
+
+Test-harness configuration (PyAutoHands docs/env_profile_redesign.md §10).
+JAX likelihood functions test JIT compilation; need JAX enabled and full-
+size datasets.
+
+ENV: jax full_datasets
+"""
+
+# %matplotlib inline
+# from pyprojroot import here
+# workspace_path = str(here())
+# %cd $workspace_path
+# print(f"Working Directory has been set to `{workspace_path}`")
+
+import numpy as np
+import jax.numpy as jnp
+import jax
+from os import path
+
+import autofit as af
+import autolens as al
+
+"""
+__Dataset__
+
+Load and plot the galaxy dataset via .fits files.
+"""
+dataset_path = path.join("dataset", "imaging", "jax_test")
+
+"""
+__Dataset Auto-Simulation__
+
+If the dataset does not already exist on your system, it will be created by running the corresponding
+simulator script. This ensures that all example scripts can be run without manually simulating data first.
+"""
+if al.util.dataset.should_simulate(dataset_path):
+    import subprocess
+    import sys
+
+    subprocess.run(
+        [sys.executable, "scripts/imaging/simulator/simple.py"],
+        check=True,
+    )
+
+dataset = al.Imaging.from_fits(
+    data_path=path.join(dataset_path, "data.fits"),
+    psf_path=path.join(dataset_path, "psf.fits"),
+    noise_map_path=path.join(dataset_path, "noise_map.fits"),
+    pixel_scales=0.2,
+)
+
+
+"""
+__Mask__
+
+The model-fit requires a 2D mask defining the regions of the image we fit the model to the data, which we define
+and use to set up the `Imaging` object that the model fits.
+"""
+mask_radius = 3.0
+
+mask = al.Mask2D.circular(
+    shape_native=dataset.shape_native,
+    pixel_scales=dataset.pixel_scales,
+    radius=mask_radius,
+)
+
+dataset = dataset.apply_mask(mask=mask)
+
+dataset = dataset.apply_over_sampling(over_sample_size_lp=1)
+
+positions = al.Grid2DIrregular(
+    al.from_json(file_path=path.join(dataset_path, "positions.json"))
+)
+
+"""
+__Model__
+
+We compose our model using `Model` objects, which represent the galaxies we fit to our data. In this
+example we fit a model where:
+
+ - The lens galaxy's light (`Sersic` bulge + `Exponential` disk), `Isothermal` mass and
+   `ExternalShear` are FIXED at the values of `simulator/simple.py` [0 parameters].
+ - The lens galaxy hosts a central `SMBH` point mass whose `centre` and `mass` are FREE
+   [3 parameters] — a free traced mass is the only configuration that exercises the
+   `SMBH.__init__` mass-to-Einstein-radius conversion inside the jit trace.
+ - The source galaxy's bulge is a linear parametric `Sersic` fixed at the simulated
+   values [0 parameters].
+
+The non-SMBH components are pinned to the simulation truth deliberately: at this workspace's
+default prior medians (e.g. `effective_radius=15.0`, `einstein_radius=4.0`) the positive-only
+linear solver zeroes the source's intensity, making the likelihood bit-identical for ANY
+source-plane mass structure — a literal generated there would not pin the SMBH's deflections
+at all. Anchored at truth the source is retained and the literal is sensitive to the SMBH.
+"""
+# Lens:
+
+bulge = af.Model(al.lp_linear.Sersic)
+bulge.centre = (0.0, 0.0)
+bulge.ell_comps = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+bulge.effective_radius = 0.6
+bulge.sersic_index = 3.0
+
+disk = af.Model(al.lp_linear.Exponential)
+disk.centre = (0.0, 0.0)
+disk.ell_comps = al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0)
+disk.effective_radius = 1.6
+
+mass = af.Model(al.mp.Isothermal)
+mass.centre = (0.0, 0.0)
+mass.ell_comps = al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0)
+mass.einstein_radius = 1.6
+
+shear = af.Model(al.mp.ExternalShear)
+shear.gamma_1 = 0.001
+shear.gamma_2 = 0.001
+
+smbh = af.Model(al.mp.SMBH)
+smbh.mass = af.LogUniformPrior(lower_limit=1e8, upper_limit=1e11)
+smbh.redshift_object = 0.5
+smbh.redshift_source = 1.0
+
+lens = af.Model(
+    al.Galaxy,
+    redshift=0.5,
+    bulge=bulge,
+    disk=disk,
+    mass=mass,
+    smbh=smbh,
+    shear=shear,
+)
+
+# Source:
+
+source_bulge = af.Model(al.lp_linear.Sersic)
+source_bulge.centre = (0.1, 0.1)
+source_bulge.ell_comps = al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0)
+source_bulge.effective_radius = 1.0
+source_bulge.sersic_index = 1.0
+
+source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)
+
+# Overall Lens Model:
+
+model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+"""
+The `info` attribute shows the model in a readable format.
+"""
+print(model.info)
+
+"""
+__Analysis__
+
+The `AnalysisImaging` object defines the `log_likelihood_function` which will be used to determine if JAX
+can compute its likelihood and batch it with `vmap`.
+"""
+analysis = al.AnalysisImaging(
+    dataset=dataset,
+    positions_likelihood_list=[al.PositionsLH(threshold=0.4, positions=positions)],
+)
+
+
+"""
+The analysis and `log_likelihood_function` are internally wrapped into a `Fitness` class in **PyAutoFit**, which pairs
+the model with likelihood.
+
+This is the function on which JAX gradients are computed, so we create this class here.
+"""
+from autofit.non_linear.fitness import Fitness
+import time
+
+batch_size = 50
+
+fitness = Fitness(
+    model=model,
+    analysis=analysis,
+    fom_is_log_likelihood=True,
+    resample_figure_of_merit=-1.0e99,
+)
+
+param_vector = jnp.array(model.physical_values_from_prior_medians)
+
+parameters = np.zeros((batch_size, model.total_free_parameters))
+
+for i in range(batch_size):
+    parameters[i, :] = model.physical_values_from_prior_medians
+
+parameters = jnp.array(parameters)
+
+start = time.time()
+print()
+print(fitness._vmap(parameters))
+print("JAX Time To VMAP + JIT Function", time.time() - start)
+
+start = time.time()
+print()
+result = fitness._vmap(parameters)
+print(result)
+print("JAX Time Taken using VMAP:", time.time() - start)
+print("JAX Time Taken per Likelihood:", (time.time() - start) / batch_size)
+
+np.testing.assert_allclose(
+    np.array(result),
+    1194.84699035,
+    rtol=1e-4,
+    err_msg="smbh: JAX vmap likelihood mismatch",
+)
+
+
+"""
+__Path A: jit-wrap ``analysis.fit_from``__
+
+Wrap ``analysis.fit_from`` in ``jax.jit`` and assert the returned ``FitImaging``
+has a ``jax.Array`` ``log_likelihood`` that matches the NumPy-path scalar.
+"""
+
+
+instance = model.instance_from_prior_medians()
+
+analysis_np = al.AnalysisImaging(
+    dataset=dataset,
+    positions_likelihood_list=[al.PositionsLH(threshold=0.4, positions=positions)],
+    use_jax=False,
+)
+fit_np = analysis_np.fit_from(instance=instance)
+print("NumPy fit.log_likelihood:", float(fit_np.log_likelihood))
+
+analysis_jit = al.AnalysisImaging(
+    dataset=dataset,
+    positions_likelihood_list=[al.PositionsLH(threshold=0.4, positions=positions)],
+    use_jax=True,
+)
+fit_jit_fn = jax.jit(analysis_jit.fit_from)
+fit = fit_jit_fn(instance)
+
+print("JIT fit.log_likelihood:", fit.log_likelihood)
+assert isinstance(
+    fit.log_likelihood, jnp.ndarray
+), f"expected jax.Array, got {type(fit.log_likelihood)}"
+np.testing.assert_allclose(
+    float(fit.log_likelihood), float(fit_np.log_likelihood), rtol=1e-4
+)
+print("PASS: jit(fit_from) round-trip matches NumPy scalar.")

scripts/misc/profiles_jit.py

@@ -46,6 +46,8 @@
   - ag.mp.NFW             → deflections_yx_2d_from, convergence_2d_from
   - ag.mp.ExternalShear   → deflections_yx_2d_from, convergence_2d_from
   - ag.mp.ExternalPotential → deflections_yx_2d_from, convergence_2d_from
+  - ag.mp.PointMass       → deflections_yx_2d_from, potential_2d_from (+ raw-zeros convergence check)
+  - ag.mp.SMBH            → deflections_yx_2d_from, potential_2d_from (+ raw-zeros convergence check)
 """
 
 import jax
@@ -660,4 +662,75 @@ def check_profile_method(
 
 print("  mp.ExternalPotential OK")
 
+"""
+ag.mp.PointMass and ag.mp.SMBH
+
+Point-mass profiles regressed under JAX (PyAutoGalaxy#553): deflections returned an
+`ArrayIrregular` wrapper into `jnp.multiply` on irregular grids, and `SMBH.__init__`
+called `np.sqrt` on a traced mass. This block pins the profile-level half of that
+coverage (the traced-mass half needs a free model parameter and lives in
+`imaging/jax_likelihood/smbh.py`).
+
+`convergence_2d_from` on these profiles is undecorated and returns a raw zeros
+array (the convergence is a Dirac delta, carried entirely by deflections and
+potential), so it gets a raw-array check instead of the autoarray-type helper.
+"""
+for label_prefix, point_profile in [
+    ("mp.PointMass", ag.mp.PointMass(centre=(0.0, 0.0), einstein_radius=0.1)),
+    (
+        "mp.SMBH",
+        ag.mp.SMBH(
+            centre=(0.0, 0.0), mass=1e10, redshift_object=0.5, redshift_source=1.0
+        ),
+    ),
+]:
+    for method_name, np_irr, jax_irr, np_uni, jax_uni in [
+        (
+            "deflections_yx_2d_from",
+            aa.VectorYX2DIrregular,
+            aa.VectorYX2DIrregular,
+            aa.VectorYX2D,
+            aa.VectorYX2D,
+        ),
+        (
+            "potential_2d_from",
+            aa.ArrayIrregular,
+            aa.ArrayIrregular,
+            aa.Array2D,
+            aa.Array2D,
+        ),
+    ]:
+        check_profile_method(
+            label=f"{label_prefix}.{method_name} (irregular)",
+            profile=point_profile,
+            method_name=method_name,
+            grid=grid_irr,
+            np_type=np_irr,
+            jax_type=jax_irr,
+        )
+        check_profile_method(
+            label=f"{label_prefix}.{method_name} (uniform)",
+            profile=point_profile,
+            method_name=method_name,
+            grid=grid_uni,
+            np_type=np_uni,
+            jax_type=jax_uni,
+        )
+
+    conv_np = point_profile.convergence_2d_from(grid=grid_irr)
+    assert isinstance(conv_np, np.ndarray), f"{label_prefix} convergence (numpy)"
+    assert np.all(np.array(conv_np) == 0.0), f"{label_prefix} convergence not zero"
+
+    conv_jit = jax.jit(
+        lambda p=point_profile: p.convergence_2d_from(grid=grid_irr, xp=jnp)
+    )()
+    assert isinstance(conv_jit, jax.Array), f"{label_prefix} convergence (jax jit)"
+    npt.assert_allclose(
+        np.array(conv_jit),
+        np.array(conv_np),
+        err_msg=f"{label_prefix}: convergence numpy vs jax (jit) mismatch",
+    )
+
+    print(f"  {label_prefix} OK")
+
 print("\nAll profiles_jit.py checks passed.")

smoke_tests.txt

@@ -5,6 +5,7 @@ misc/database/scrape/general.py
 imaging/jax_likelihood/rectangular.py
 imaging/jax_likelihood/mge.py
 imaging/jax_likelihood/lp.py
+imaging/jax_likelihood/smbh.py
 interferometer/jax_likelihood/rectangular.py
 interferometer/jax_likelihood/mge.py
 point_source/jax_likelihood/point.py