feat: scan-based multi-plane ray-tracing for substructure

[Jammy2211]

Summary

• Adds autolens/lens/substructure_util.py with traced_grids_via_scan — a jax.lax.scan-based multi-plane ray-tracing function that replaces the Python-loop approach for O(1000) halos
• Includes precompute_scaling_matrix to build the (n_planes, n_planes) cosmological scaling factor matrix outside jit
• Includes galaxies_to_halo_arrays to convert LOSSampler Galaxy lists to padded (n_planes, max_N, n_params) arrays
• The scan handles NFWTruncatedSph halos (via vmapped_deflections_from from PyAutoGalaxy#455), a macro lens (PowerLaw+ExternalShear captured in closure), and negative kappa sheets per plane

API

from autolens.lens.substructure_util import (
    traced_grids_via_scan,
    precompute_scaling_matrix,
    galaxies_to_halo_arrays,
)

# Precompute (outside jit)
scaling_matrix = precompute_scaling_matrix(plane_redshifts, cosmology)
halo_params, halo_mask, sheet_kappas = galaxies_to_halo_arrays(
    galaxies, plane_redshifts, max_n=200, profile_cls=ag.mp.NFWTruncatedSph,
)

# Trace (inside jit)
traced_grids = traced_grids_via_scan(
    grid, halo_params, halo_mask, scaling_matrix,
    macro_deflections_fn, macro_plane_mask, sheet_kappas,
    halo_profile_cls=ag.mp.NFWTruncatedSph,
)
source_grid = traced_grids[-1]  # (M, 2) at source redshift

Test plan

• 4-plane parity test (3 LOS planes + source, 15 halos + macro + sheets) — all plane grids match existing Tracer loop path to atol=1e-5
• JIT compilation succeeds and matches
• Different parameters produce different results (no stale compilation)
• Companion PR: autolens_workspace_test#126
• Depends on: PyAutoGalaxy NaN-safe masking fix (pushed to main)

Ref: PyAutoLens#542 (prompt 2 of 4)

🤖 Generated with Claude Code