Neural Network Reconstruction of Non-Gaussian Initial Conditions from Dark Matter Halos

We develop a machine learning approach to reconstructing the cosmological initial conditions from late-time dark matter halo number density fields in redshift space, with the goal of improving sensitivity to cosmological parameters, and in particular primordial non-Gaussianity. Using an U-Net architecture, our model achieves a cross-correlation accuracy of 44% for scales out to $k = 0.4 \text{ h}/\text{Mpc}$ between reconstructed and true initial conditions of Quijote 1 Gpc$^3$ simulation boxes with an average halo number density of $\bar{n} = 4\times 10^{-4}$ (h/Mpc)$^{3}$ in the tracer field at $z=0$ . We demonstrate that our reconstruction is likely to be optimal for this setup and that it is highly effective at reducing redshift-space distortions. Using a Fisher analysis, we show that reconstruction improves cosmological parameter constraints derived from the power spectrum and bispectrum. By combining the power spectrum monopole, quadrupole, and bispectrum monopole up to $k_{\rm{max}} = 0.52 \text{ h}/\text{Mpc}$, our joint analysis of pre- and post-reconstructed fields from the Quijote simulation suite finds improved marginalized errors on all cosmological parameters. In particular, reconstruction improves constraints on $f_{\rm{NL}}$ by factors of 1.33, 1.88, and 1.57 for local, equilateral, and orthogonal shapes. Our findings demonstrate the effectiveness of reconstruction in decoupling modes, mitigating redshift-space distortions and maximizing information on cosmology. The results provide important insights into the amount of cosmological information that can be extracted from small scales, and can potentially be used to complement standard analysis of observational data, upon further development.