A Differentiable, End-to-End Forward Model for 21 cm Cosmology: Estimating the Foreground, Instrument, and Signal Joint Posterior

Abstract: We present a differentiable, end-to-end Bayesian forward modeling framework for line intensity mapping cosmology experiments, with a specific focus on low-frequency radio telescopes targeting the redshifted 21 cm line from neutral hydrogen as a cosmological probe. Our framework is capable of posterior density estimation of the cosmological signal jointly with foreground and telescope parameters at the field level. Our key aim is to be able to optimize the model's high-dimensional, non-linear, and ill-conditioned parameter space, while also sampling from it to perform robust uncertainty quantification within a Bayesian framework. We show how a differentiable programming paradigm, accelerated by recent advances in machine learning software and hardware, can make this computationally-demanding, end-to-end Bayesian approach feasible. We demonstrate a proof-of-concept on a simplified signal recovery problem for the Hydrogen Epoch of Reionization Array experiment, highlighting the framework's ability to build confidence in early 21 cm signal detections even in the presence of poorly understood foregrounds and instrumental systematics. We use a Hessian-preconditioned Hamiltonian Monte Carlo algorithm to efficiently sample our parameter space with a dimensionality approaching $N\sim10^5$, which enables joint, end-to-end nuisance parameter marginalization over foreground and instrumental terms. Lastly, we introduce a new spherical harmonic formalism that is a complete and orthogonal basis on the cut sky relevant to drift-scan radio surveys, which we call the spherical stripe harmonic formalism, and it's associated three-dimensional basis, the spherical stripe Fourier-Bessel formalism.