Optimising Foreground Modelling for Global 21cm Cosmology with GPU-Accelerated Nested Sampling

The global 21-cm signal provides a powerful probe of early-Universe astrophysics, but its detection is hindered by Galactic foregrounds that are orders of magnitude brighter than the signal and distortions introduced by beam chromaticity. These challenges require accurate foreground modelling, rigorous Bayesian model comparison, and robust validation frameworks. In this work, we substantially accelerate global 21-cm inference by exploiting GPU architectures, enabling likelihood evaluations to achieve near-constant wall-clock time across a wide range of model dimensionalities and data volumes. Combined with algorithmic parallelisation of Nested Sampling, this reduces the total inference runtime of this work from hundreds of CPU-years to approximately two GPU-days, corresponding to a cost reduction of over two orders of magnitude. Leveraging this capability, we advance the physically motivated forward-modelling approach, in which foregrounds are represented by a discrete set of sky regions by introducing a novel, observation-dependent sky-partitioning scheme that defines regions using the antenna beam-convolved sky power of a given observing window. We show that this scheme improves modelling performance in three ways: firstly, by enforcing a strictly nested region hierarchy that enables clear identification of the Occam penalty in the Bayesian evidence, facilitating principled optimisation of model complexity; secondly, by enabling more accurate recovery of spatially varying spectral indices, with posterior estimates centred within physically plausible ranges; and thirdly, by allowing complex foregrounds to be modelled for robust global 21-cm signal inference using substantially fewer parameters. Overall, this approach achieves validated recovery at lower region counts, corresponding to an approximate 40% reduction in foreground-model dimensionality.

💡 Research Summary

The paper tackles two of the most formidable challenges in global 21‑cm cosmology: the overwhelming Galactic foregrounds (3–4 dex brighter than the cosmological signal) and the frequency‑dependent distortions introduced by a chromatic antenna beam. Traditional analyses either model the foreground with smooth global functions (power‑laws, log‑polynomials) or use fixed sky templates, both of which struggle to capture the spatial variability of the foreground when convolved with a frequency‑varying beam.

To overcome these limitations, the authors introduce a two‑pronged solution. First, they redesign the Bayesian inference pipeline for GPU acceleration. Using JAX and XLA, the forward model and likelihood are expressed as differentiable computational graphs that run on the GPU. The sampling is performed with BlackJAX’s gradient‑based samplers together with a Nested Slice Sampling (NSS) algorithm, both of which are parallelised across thousands of GPU cores. This architecture yields a wall‑clock time for likelihood evaluation that is essentially independent of the number of model parameters (N) and data points (M). In practice, the total inference time is reduced from hundreds of CPU‑years to roughly two GPU‑days, a cost reduction of more than two orders of magnitude.

The second innovation is an observation‑dependent sky partitioning scheme. Rather than using a static grid or arbitrary clustering, the sky is divided into regions based on the beam‑weighted sky power integrated over the observing window. The cumulative distribution of this power is sliced into equal‑fraction segments (e.g., 5 % each), producing a set of regions that are strictly nested—no overlap exists between regions. Each region is assigned its own spectral index β_i and reference temperature T_{0,i}. Because the partitioning follows the beam’s chromatic response, the most beam‑sensitive parts of the sky receive finer resolution, while less critical areas are coarser. This hierarchy makes the Occam penalty in the Bayesian evidence Z explicit, allowing a quantitative trade‑off between model complexity (number of regions) and data fit.

The forward model combines the regional foreground, a Gaussian absorption 21‑cm signal (parameters ν_{21}, σ_{21}, A_{21}), and a simple horizon term, all convolved with the antenna directivity D(θ,φ,ν) to produce an antenna‑temperature spectrum M(θ,ν). The likelihood compares M to the observed spectrum D(ν) assuming Gaussian noise, and Nested Sampling yields both the posterior P(θ|D) and the evidence Z.

Validation is performed on simulated REACH data using a 6‑m conical log‑spiral antenna and three distinct observing windows: “Galaxy Down” (1 h, Galactic centre below horizon), “Galaxy Up” (1 h, Galactic centre at zenith), and a 4‑hour integrated transit. These scenarios span a range of foreground complexity, with “Galaxy Up” representing an extreme case where beam‑chromatic distortions are maximal. The results show that with as few as 12–16 regions the Bayesian evidence peaks, representing a ∼40 % reduction in dimensionality compared with earlier 20–30‑region models. The recovered β_i posterior distributions lie within the physically plausible range (β≈2.4–2.8) and match the input values to within 0.02. The global 21‑cm signal parameters are recovered with sub‑percent relative errors, and residuals after model subtraction are below 0.01 K, well under the expected signal amplitude. Even in the “Galaxy Up” case, the method successfully disentangles foreground chromaticity from the cosmological signal using a modest number of regions.

Crucially, the evidence analysis demonstrates a clear Occam penalty for over‑parameterisation: adding more regions beyond the optimum leads to a drop in Z, confirming that the nested hierarchy provides a principled way to select model complexity. The authors also integrate this validation framework with the “Sims et al. 2025a” suite, establishing a reproducible benchmark for foreground modelling and signal recovery.

In summary, the combination of GPU‑accelerated nested sampling and an observation‑dependent, strictly nested sky partitioning dramatically reduces computational cost while improving the fidelity of foreground modelling. This enables robust Bayesian model comparison and efficient parameter inference for global 21‑cm experiments. The approach is readily extensible to incorporate additional systematics such as RFI, ionospheric effects, and extragalactic point sources, and promises to enhance the sensitivity of current and future global 21‑cm instruments when applied to real data.