Jacob Tutt

End-to-end uncertainty propagation in scientific inference often requires repeated evaluation of coupled simulators spanning multiple physical domains, often at prohibitive computational cost. Global 21-cm cosmology is an extreme example: a millikelvin cosmological signal must be separated from radio foregrounds and instrumental effects four orders of magnitude larger, making inference acutely sensitive to model misspecification. Reliable recovery is therefore contingent on an exceptionally accurate model that jointly propagates uncertainty through the complete observational chain. We present a differentiable, GPU-accelerated hybrid forward model combining neural emulation of the cosmological signal, a physics-informed multi-resolution parameterisation of the galactic foregrounds, and a linear surrogate model of the antenna. This unifies simulators spanning astrophysics and computational electromagnetics while reducing evaluation time by a factor of 7.2×10⁷. Using GPU-accelerated Nested Slice Sampling, we benchmark the framework across 100 unseen signal realisations and find that 94 injected profiles lie entirely within their corresponding 95% credible bands. Together, the simulator and reference Nested Slice Sampling results provide a demanding benchmark of posterior calibration, mode coverage, and evidence accuracy across next-generation stochastic samplers, simulation-based inference, and variational inference frameworks.

Recovering the 21-cm power spectrum from interferometric visibilities requires a faint cosmological signal to be distinguished from foreground emission many orders of magnitude brighter, after the chromatic instrumental response has mapped both into highly degenerate structures. Physically motivated Bayesian forward models provide a principled framework for joint inference and rigorous uncertainty quantification, but can require 10⁴–10⁶ parameters, rendering posterior exploration and Bayesian-evidence evaluation computationally prohibitive. We first extend BayesEoR beyond its existing tangent-plane foreground treatment by parameterising wide-field emission directly in the spherical image domain. We then derive a Karhunen–Loève basis that preserves the noise-relevant prior-predictive foreground uncertainty by construction, reducing the dense system required for each marginal-likelihood evaluation. The retained rank therefore directly quantifies the noise-relevant dimensionality of foreground uncertainty in visibility space. Extending the model across the final 10° in zenith angle to the horizon increases this dimensionality by a factor of 4.63, demonstrating the disproportionate contribution of near-horizon emission and the need to propagate its uncertainty in cosmological inference. For full-field analysis, the method reduces the joint coefficient space from 306,163 to just 9,551, enabling unbiased power-spectrum recovery and Bayesian-evidence evaluation in just 0.32 GPU-hours, giving an estimated computational saving of  830 relative to an unreduced GPU-resident implementation. Throughout, we highlight the importance of Bayesian validation, allowing robust cosmological recovery to be distinguished from apparently plausible but spurious detections caused by model misspecification.

Solving large dense linear systems and eigenvalue problems is a core requirement in many areas of scientific computing, but scaling these operations beyond a single GPU remains challenging within modern programming frameworks. While highly optimized multi-GPU solver libraries exist, they are typically difficult to integrate into composable, just-in-time compiled Python workflows. JAXMg provides distributed dense linear algebra for JAX, enabling linear solves and decompositions for matrices that exceed single-GPU memory limits. By interfacing JAX with NVIDIA’s cuSOLVERMp through an XLA Foreign Function Interface, JAXMg exposes distributed GPU routines as JIT-compatible JAX primitives. This design allows scalable linear algebra to be embedded directly within JAX programs, preserving composability with JAX transformations and enabling multi-GPU and multi-node execution in end-to-end scientific workflows.