feat: end to end `jax.jit` (and ideally `vmap`) for a multi plane substructure simulator

[mwiet]

What I'm trying to do

I'm building a forward modelling pipeline on top of PyAutoLens for dark matter substructure. The forward model is a fairly standard galaxy scale strong lens with a couple of populations of perturbers layered on top:

• one macro lens (power law plus external shear) at z_lens = 0.5,
• a handful of lens plane subhaloes (typically 5 to 50, NFWTruncatedSph for CDM and WDM, a mix of cored cNFWSph / cNFWMCRLudlowSph and core collapsed NFWTruncatedSph for SIDM),
• a multi plane line of sight (LOS) population of order 1000 halos spread across 8 planes between observer and source, built via al.Tracer and the existing LOSSampler helpers,
• a Gaussian PSF (al.Kernel2D style), Poisson noise via al.SimulatorImaging, an over sampled al.Grid2D.uniform.

The hot loop is conceptually one function, theta -> noisy image, and the calling code wants to evaluate it of order 10^6 times with different theta. On the CPU this is the wall clock bottleneck of the whole project by a comfortable margin. A jit compiled, ideally vmap batched and GPU runnable, version of that path would change the practical scale of what we can do by a couple of orders of magnitude.

The point of this issue is to ask whether the simulator side specifically can be added to the JAX roadmap, and to flag the bits that this kind of substructure model actually needs so the priority ordering can be informed by a concrete use case.

What already exists (no need to redo any of this)

Skimming the issue trackers, quite a lot of the groundwork is either done or already in flight. Listing what I found so this issue can focus on the gaps.

PyAutoArray:

• Feature/point redesign #334: use_jax=True flag on Simulators and PointSolver.
• feature/fits interface #338: Array2D.native jit traceable for the JAX simulator path.
• feature/andrew implementation #331: audit of the JAX xp interface for user facing introduction.
• Bug with ShapeSolver and Circle #320: AbstractMeshGeometry picklable, xp module replaced by _use_jax bool plus property.
• use the new jax coordinate array triangles #317: KNN barycentric interpolator as a JAX native Delaunay alternative.
• Feature/multi image visualization #307: perf work on the JAX NNLS solver under vmap.
• feature/triangle analysis closest #278: NaN gradients from jaxnnls backward pass.
• remove preloads #306: 2D regular grid interpolator with JAX path, wires into DatasetInterp.
• feature/vectorised triangles #286: propagate xp through Grid2DIrregular derived constructors (called out as the thing that unblocks the source plane JIT).
• feature/positions_lh_mass_centre #326: TransformerNUFFT migration.

PyAutoGalaxy:

• jax: register FitImaging, Tracer, DatasetModel as pytrees on use_jax path #445: mass profile MGE/CSE fallback, JAX port, testing, docs.
• test: mask padding likelihood sanity check #436: self.centre JAX traceable under vmap.
• Rename PlotterInterface -> Plotter; fix _to_lines; update docs #403: replaced the Ludlow16 colossus pure_callback with a JAX native implementation.
• docs(point): add module-level docstrings to point package #397: the prior research issue that scoped Rename PlotterInterface -> Plotter; fix _to_lines; update docs #403.
• feat: re-key linear_light_profile_intensity_dict to survive jax.jit round-trip #448: MGE/CSE fallback for zero returning mass profiles.
• Add CSV I/O for PointDataset (cluster-friendly alternative to JSON) #446: port the CSE module to support JAX via the xp parameter.
• perf: Lazy imports for scipy/matplotlib to reduce import time #426: EllipseMultipoleScaled JAX traceable (move derivation out of __init__).
• Update ReadTheDocs to Python 3.12 #417: print_vram_use fails on lmp.Sersic lenses under a JAX tracer.
• Visualization cleanup #411: wire AnalysisEllipse for JAX via fit_from and pytree registration.

PyAutoLens:

• fix: subhalo redshift as free parameter raises TracerBoolConversionError under JAX #498: subhalo redshift as free parameter TracerBoolConversionError under JAX.
• refactor: AnalysisPoint auto-configures PointSolver when use_jax=True #466: AnalysisPoint auto configures PointSolver when use_jax=True.
• perf: slim LOS tests and remove JAX solver test #420: LOS test slimming and JAX solver test removal.
• feat: pytree-register FitImaging for jax.jit-returnable fits (Path A) #444: pytree register FitImaging for jax.jit returnable fits (Path A).
• feat: pytree PoC for group-scale MGE FitImaging (mge_group_pytree) #452: pytree PoC for the group scale MGE FitImaging.
• feat: jax.jit(analysis.fit_from) for FitInterferometer + MGE source #454: jax.jit(analysis.fit_from) for FitInterferometer plus MGE source.
• feat: jax.jit(analysis.fit_from) for FitPointDataset (point-source) #457: jax.jit(analysis.fit_from) for FitPointDataset.
• feat: jit(fit_from) round-trip for double-source-plane + rectangular #460,feat: jit(fit_from) round-trip for Delaunay source + MGE lens #461: jit(fit_from) round trip for double source plane plus rectangular, and Delaunay source plus MGE lens.

So the big picture is: the inference side (analysis.fit_from) is where most of the open jax.jit work lives at the moment, plus the last few xp propagation fixes around Grid2DIrregular and the CSE port. The bit that doesn't seem to be explicitly tracked yet is the substructure forward simulator path: al.SimulatorImaging.via_tracer_from(tracer=...) for a multi plane Tracer carrying a long, possibly variable length, list of dark matter perturbers.

What I'm asking for

A roadmap covering end to end jax.jit of the substructure forward simulator. Concretely, the following pieces would need to be JAX traceable and jit friendly so that this kind of multi plane substructure simulation can run inside jax.jit.

1. Dark matter mass profiles under convergence_2d_from and deflections_yx_2d_from

The relevant profiles are:

• autogalaxy.profiles.mass.dark.nfw_truncated.NFWTruncatedSph,
• autogalaxy.profiles.mass.dark.nfw_truncated_mcr.NFWTruncatedMCRLudlowSph,
• autogalaxy.profiles.mass.dark.cnfw.cNFWSph,
• autogalaxy.profiles.mass.dark.cnfw_mcr.cNFWMCRLudlowSph (cored SIDM, lens plane).

The closed #403 /#397 already give us a JAX native Ludlow16, so the MCR branch is presumably partway there. What's missing for the substructure path is vmap friendliness across a batch of N halos with different mass_at_200, concentration, centre, f_c, plus the same on the deflection field. Bluntly: vmap(profile.deflections_yx_2d_from) over the halo axis ought to give us the contribution of every halo with a single GPU launch, then a sum gives the plane deflection. Currently this is a Python loop over halos inside Tracer.

A precondition for the cNFW bits of this is the deflection sign flip and zero convergence bug in autogalaxy.profiles.mass.dark.cnfw that I've reported separately: PyAutoLabs/PyAutoGalaxy#451. That one should land first, otherwise jit'ing the broken branch just makes the broken behaviour faster.

2. Multi plane Tracer traceable under jit with O(1000) galaxies

A typical realisation builds a Tracer(galaxies=[macro, *subhaloes, *los, source]) with often 1000+ entries. Two things make this jit hostile today:

• The galaxy list length is theta dependent (different draws of the SHMF give different N), and jit specialises on shape, so each new N would trigger a recompile.
• The multi plane recursion inside Tracer uses Python iteration over per plane galaxy lists.

The usual JAX idioms (pad to a max N, mask the unused slots, jax.lax.scan over planes) would address both. I'd happily prototype this against a stripped down branch if it would help focus the design.

3. LOSSampler and friends

The sampler itself doesn't need to be jit'd, it's called once per realisation and dominates approximately none of the wall clock. But the output of the sampler currently feeds galaxies straight into the Tracer, so whatever decision is made on padded vs variable length galaxy lists in (2) drives the API here too. The closed #420 (LOS test slimming) suggests this area has been getting some love already.

4. Convolver / PSF convolution under JAX

The forward model uses a Gaussian PSF (al.Kernel2D) plus an over sampled al.Grid2D.uniform. JAX's own jax.scipy.signal.fftconvolve would be the obvious backend on the JAX path, but the wiring isn't there today as far as I can tell (there's no convolver specific issue I could find). Even a small use_jax aware shim around Convolver.convolved_image_from would cover the relevant use case.

5. Poisson noise with a JAX PRNGKey

al.SimulatorImaging(..., add_poisson_noise=True, noise_seed=...) currently takes a numpy seed. To make theta -> noisy image referentially transparent, jit friendly, and vmap able over a batch of keys, the simulator wants to thread a PRNGKey through instead. jax.random.poisson already exists, so this is fundamentally a plumbing change rather than a numerical one.

6. Macro lens (mp.PowerLaw plus mp.ExternalShear) and source (lp.SersicCore)

These are simpler than the dark profiles and I'd expect them to be mostly there once the broader mass / light profile JAX work is wrapped up. Flagging here just so the priority list captures them too.

7. (Stretch) vmap over theta for batched evaluation

The headline win is vmap(jit(simulate))(thetas, keys) running of order 1024 lensed images per GPU launch. Everything above is a prerequisite, but it's worth naming the stretch goal explicitly because some of the design choices (especially padded vs ragged galaxy lists in 2) make a much bigger difference once vmap is on the table.

Suggested ordering

If it's useful, here's how I'd sequence the chunks above against the in flight work:

1. Land the cNFW deflection bug fix (separate issue) so the cored branch isn't a moving target.
2. Land feature/vectorised triangles #286 (Grid2DIrregular xp propagation) and remove preloads #306 (regular grid interpolator JAX path); both look like they unblock several downstream simulator pieces.
3. Open a new issue for the dark profile vmap friendliness (item 1 above), since that's the bit none of the existing JAX issues explicitly cover.
4. Decide on the padded vs variable length galaxy list strategy (item 2). This is the single design decision that drives most of the rest.
5. Convolver and Poisson noise on the JAX path (items 4 and 5).
6. End to end jit(simulate) smoke test on a representative substructure configuration as a regression target, then the vmap extension (item 7).

[Jammy2211]

Hi @mwiet — thanks for the detailed issue. We've been working through the JAX roadmap and a good chunk of what you listed is already landed. Here's a status update plus a plan for the remaining work.

What's already done

These items from your list need no further work:

• Dark profile convergence_2d_from / deflections_yx_2d_from under JAX — all four profiles (NFWTruncatedSph, cNFWSph, and their MCR Ludlow variants) are fully xp-threaded and work under jax.jit.
• cNFW deflection boundary bug — fixed in PyAutoGalaxy feat: jax.jit(analysis.fit_from) for FitInterferometer + MGE source #454 (the F_func boundary condition you flagged in feat: pytree PoC for rectangular-pixelization FitImaging #451).
• Convolver / PSF convolution — both FFT (jnp.fft.rfft2) and real-space (jax.scipy.signal.convolve) paths are JAX-ready.
• Poisson noise — routes through jax.random.poisson when xp=jnp (integer seed; PRNGKey support coming in step 3 below).
• Macro lens profiles — PowerLaw (uses jax.lax.scan series expansion), ExternalShear, SersicCore all JAX-traceable.
• Cosmology scaling factors — fully xp-threaded in autogalaxy/cosmology/model.py.
• Grid2DIrregular — core class is xp-aware; regular grid interpolator has a JAX path (PyAutoArray remove preloads #306).
• CSE module — ported to JAX (PyAutoGalaxy feat: CSV I/O for PointDataset for cluster-scale workflows #447).

Remaining work — 4-step plan

We're building this as a separate pure-function path alongside the existing OO Tracer (which continues to work for normal modeling via Python-loop unrolling). The new path targets your substructure use case where variable N and O(1000) halos make unrolling prohibitive.

1. vmap over N subhalos (in progress)

A generic vmapped_deflections_from classmethod on MassProfile that takes (grid, params_batch, mask) arrays and uses jax.vmap to compute all N halo deflections in a single GPU launch. Each spherical profile provides a small radial_deflection_from static method (just the physics formula). The generic wrapper handles centre subtraction, radial distance, cartesian reconstruction, masking, and summing — written once on the base class.

Currently implemented for NFWTruncatedSph and cNFWSph, with integration tests passing against the existing loop path.

2. jax.lax.scan over redshift planes

Replace the nested Python loops in tracer_util.traced_grid_2d_list_from with a jax.lax.scan that carries deflection state through planes. Precomputed (n_planes, n_planes) scaling-factor matrix from cosmology. Fixed-shape per-plane halo stacks with boolean masks.

3. End-to-end jax.jit(simulate)

Wire the scan-based ray-tracing through PSF convolution and Poisson noise. Add prng_key parameter support for noise (enabling vmap over seeds). Thread xp through over-sampling. Smoke test on a representative substructure configuration.

4. vmap over theta (stretch)

vmap(jit(simulate))(thetas, keys) for ~1024 images per GPU launch. LOSSampler output → padded array conversion helper. Memory estimation utility.

The final user-facing API will live in autolens_workspace/scripts/imaging/features/advanced/los_halos/ alongside the existing simulator.

Happy to coordinate if you'd like to prototype against any of this — the vmapped_deflections_from API from step 1 is already usable on the current branch.

[Jammy2211]

All four prompts are now implemented and merged (prompt 4 PR pending):

Step	What	PR
1	vmapped_deflections_from — batched single-plane deflections	PyAutoGalaxy#455 ✅
2	traced_grids_via_scan — jax.lax.scan over planes	PyAutoLens#543 ✅
3	simulate_substructure — end-to-end jit(theta → noisy_image)	PyAutoLens#544 ✅
4	batched_simulate_substructure — vmap over realizations	PyAutoLens#545 (open)

Workspace example showing the full API (LOSSampler → padded arrays → single + batched simulation):
autolens_workspace/scripts/imaging/features/advanced/los_halos/simulator_jax.py

The example produces 996 halos across 10 planes, a single jitted image, and an 8-image vmap batch — all from the existing LOSSampler output converted to padded arrays via galaxies_to_halo_arrays.

[Jammy2211]

All four prompts plus the parameterization follow-up are now merged:

Step	What	PRs
1	vmapped_deflections_from — batched single-plane deflections	PyAutoGalaxy#455
2	traced_grids_via_scan — jax.lax.scan over planes	PyAutoLens#543
3	simulate_substructure — end-to-end jit(theta → noisy_image)	PyAutoLens#544
4	batched_simulate_substructure — vmap over realizations	PyAutoLens#545
5	Parameterized lens mass, source light, lens light	PyAutoLens#545 (amended), autolens_workspace#209

Key update since last comment: the lens mass model (PowerLaw + ExternalShear), source light (SersicCore), and lens light are now dynamic JAX inputs — not baked-in constants. All profile parameters flow through as traced arrays, so they can be varied between realizations without recompilation. This addresses @mwiet's feedback that the lens mass model wasn't previously jittable.

Workspace example (merged): autolens_workspace/scripts/imaging/features/advanced/los_halos/simulator_jax.py — demonstrates the full parameterized API with LOSSampler, single + batched simulation.

API summary:

from autolens.lens.substructure_util import (
    simulate_substructure,          # jit: single image
    batched_simulate_substructure,  # vmap: batch of images
    precompute_scaling_matrix,      # cosmology scaling (outside jit)
    galaxies_to_halo_arrays,        # LOSSampler → padded arrays
    los_realizations_to_arrays,     # multiple realizations → batched arrays
)

# Define parameterized lens mass function
def lens_mass_fn(grid, params):
    power_law = al.mp.PowerLaw(
        centre=(params[0], params[1]),
        ell_comps=(params[2], params[3]),
        slope=params[4],
        einstein_radius=params[5]),
    )
    shear = al.mp.ExternalShear(gamma_1=params[6], gamma_2=params[7])
    galaxy = al.Galaxy(redshift=0.5, mass=power_law, shear=shear)
    g = aa.Grid2DIrregular(values=grid, xp=jnp)
    return galaxy.deflections_yx_2d_from(grid=g, xp=jnp).array

# All profile params are dynamic — no recompilation when values change
image = jax.jit(simulate_substructure)(
    grid, image_shape, halo_params, halo_mask, scaling_matrix,
    lens_mass_fn, lens_mass_params, lens_plane_mask, sheet_kappas,
    source_light_fn, source_light_params, psf_kernel,
    exposure_time, background_sky_level, prng_key, halo_profile_cls,
)