Probabilistic cosmological inference on HI tomographic data

TL;DR

This paper presents a deep learning approach to extract cosmological parameters from 21-cm hydrogen tomography data using an encoder and likelihood-free inference, achieving high precision and robustness across different datasets.

Contribution

It introduces a novel pipeline combining a 3D convolutional encoder with a Masked Autoregressive Flow for probabilistic inference of cosmology from 21-cm data.

Findings

Encoder learns separable latent representations.

Density estimator constrains cosmology with R^2 ≥ 0.91.

Model remains effective on out-of-distribution data with R^2 ≥ 0.80.

Abstract

We explore the possibility of retrieving cosmological information from 21-cm tomographic data at intermediate redshift. The first step in our approach consists of training an encoder, composed of several three dimensional convolutional layers, to cast the neutral hydrogen 3D data into a lower dimension latent space. Once pre-trained, the featurizer is able to generate 3D grid representations which, in turn, will be mapped onto cosmology ($\Omega_{\rm m}$, $\sigma_{8}$) via likelihood-free inference. For the latter, which is framed as a density estimation problem, we consider a Bayesian approximation method which exploits the capacity of Masked Autoregressive Flow to estimate the posterior. It is found that the representations learned by the deep encoder are separable in latent space. Results show that the neural density estimator, trained on the latent codes, is able to constrain cosmology with a precision of $R^2 \ge 0.91$ on all parameters and that most of the ground truth of the instances in the test set fall within $1\sigma$ uncertainty. It is established that the posterior uncertainty from the density estimator is reasonably calibrated. We also investigate the robustness of the feature extractor by using it to compress out-of-distribution dataset, that is either from a different simulation or from the same simulation but at different redshift. We find that, while trained on the latent codes corresponding to different types of out-of-distribution dataset, the probabilistic model is still reasonably capable of constraining cosmology, with $R^2 \ge 0.80$ in general. This highlights both the predictive power of the density estimator considered in this work and the meaningfulness of the latent codes retrieved by the encoder. We believe that the approach prescribed in this proof of concept will be of great use when analyzing 21-cm data from various surveys in the near future.