From bb60eccd0fa7c6e2f8dfbcbe16e928229896eab3 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 6 Aug 2026 14:33:03 -0400 Subject: [PATCH] test: CI JAX coverage for PointMass/SMBH profiles 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 --- scripts/CLAUDE.md | 6 + scripts/imaging/jax_likelihood/smbh.py | 268 +++++++++++++++++++++++++ scripts/misc/profiles_jit.py | 73 +++++++ smoke_tests.txt | 1 + 4 files changed, 348 insertions(+) create mode 100644 scripts/imaging/jax_likelihood/smbh.py diff --git a/scripts/CLAUDE.md b/scripts/CLAUDE.md index 453af319..1bf749d1 100644 --- a/scripts/CLAUDE.md +++ b/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. diff --git a/scripts/imaging/jax_likelihood/smbh.py b/scripts/imaging/jax_likelihood/smbh.py new file mode 100644 index 00000000..c4803668 --- /dev/null +++ b/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.") diff --git a/scripts/misc/profiles_jit.py b/scripts/misc/profiles_jit.py index 9c887073..90ae3699 100644 --- a/scripts/misc/profiles_jit.py +++ b/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.") diff --git a/smoke_tests.txt b/smoke_tests.txt index 30a9df49..4075308e 100644 --- a/smoke_tests.txt +++ b/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