feat: batched_simulate_substructure via jax.vmap

[Jammy2211]

Summary

• Adds batched_simulate_substructure — wraps simulate_substructure with jax.vmap over halo params, masks, sheet kappas, and PRNGKeys
• Adds los_realizations_to_arrays — stacks multiple LOSSampler Galaxy lists into batched (B, n_planes, max_N, n_params) padded arrays
• Produces (B, H, W) image stacks — ~1024 noisy lensed images per GPU launch

API

from autolens.lens.substructure_util import (
    batched_simulate_substructure,
    los_realizations_to_arrays,
)

# Stack B realizations
hp_batch, hm_batch, sk_batch = los_realizations_to_arrays(
    realization_galaxies, plane_redshifts, max_n=200,
    profile_cls=ag.mp.NFWTruncatedSph,
)
keys = jax.random.split(jax.random.PRNGKey(0), batch_size)

# Batch simulate
images = jax.jit(batched_simulate_substructure)(
    grid, image_shape, hp_batch, hm_batch, scaling_matrix,
    macro_fn, macro_mask, sk_batch, source_fn,
    psf, exposure_time, sky, keys, ag.mp.NFWTruncatedSph,
)  # shape: (B, H, W)

Test plan

• 8-element batch matches sequential loop (atol=1e-5)
• jax.jit(batched_simulate_substructure) compiles and matches
• Different realizations produce different images
• Companion PR: autolens_workspace_test#129

Ref: PyAutoLens#542 (prompt 4 of 4) — completes the full feature series

🤖 Generated with Claude Code