Mean-Field Simulation-Based Inference for Cosmological Initial Conditions

Abstract

Reconstructing cosmological initial conditions (ICs) from late-time observations is a difficult task, which relies on the use of computationally expensive simulators alongside sophisticated statistical methods to navigate multi-million dimensional parameter spaces. We present a simple method for Bayesian field reconstruction based on modeling the posterior distribution of the initial matter density field to be diagonal Gaussian in Fourier space, with its covariance and the mean estimator being the trainable parts of the algorithm. Training and sampling are extremely fast (training: $\sim 1 \, \mathrm{h}$ on a GPU, sampling: $\lesssim 3 \, \mathrm{s}$ for 1000 samples at resolution $128^3$), and our method supports industry-standard (non-differentiable) $N$-body simulators. We verify the fidelity of the obtained IC samples in terms of summary statistics.

Figures (3)

• Figure 1: Our method produces posterior IC samples constrained by a given late-time density field ( training:$\simeq 1.5 \, \mathrm{h}$ on a GPU, sampling:$< 3 \, \mathrm{s}$ for 1000 samples at resolution $128^3$). Top row: slices of the initial and final overdensity fields of the target simulation. Bottom row: two examples of the generated IC samples. All the shown slices are averaged over the depth of $100 \ \text{Mpc}/h$ in the third axis direction. Rightmost column: power spectrum, transfer function, and cross-correlation of the generated samples and the ground truth. Shaded regions correspond to $1\sigma$ and $2\sigma$ errors, and the yellow $C(k)$ line corresponds to the cross-correlation between the final and the initial density fields. Note that the MAP estimate, which is the only quantity recovered by most other approaches, loses significant power at small scales and is not fully able to describe the field's statistical properties.
• Figure 2: Top row: slices of the initial density field ${\boldsymbol{z}}_{\rm truth}$ of the target simulation and the corresponding MAP estimate $\mathbf{\hat{\boldsymbol{\mu}}}_{\boldsymbol{\theta}}({\boldsymbol{x}}_{\mathrm{obs}})$. Bottom row left: diagonal values $\bm D_{{\boldsymbol{\theta}}}$ of the trained posterior precision matrix $\bm Q_{\boldsymbol{\theta}} = \bm Q^P + \bm Q_{\boldsymbol{\theta}}^L$ as a function of the wavenumber $k$. Bottom row right: the standard deviation computed from 1000 samples. All the shown slices are averaged over the depth of $100 \ \text{Mpc}/h$ in the third axis direction.
• Figure 3: Coverage test. We plot the histograms of the quantity $\Delta(\bm k)$ defined in (\ref{['coverage']}) for four different spherical shells in $k$-space. The red lines represent the univariate normal distribution $\mathcal{N}(0, 1)$. The fact that all the histograms closely follow the $\mathcal{N}(0, 1)$ distribution validates our posterior coverage test for different ranges of $k$.