Searching optimal scales for reconstructing cosmological initial conditions using convolutional neural networks

TL;DR

This work tackles reconstructing the cosmological initial density field from late-time nonlinear matter distributions using convolutional neural networks. It systematically investigates how the input sub-box scale $L_ ext{sub}$ affects reconstruction quality and introduces a dual-input 3D CNN to fuse information from two different scales, improving performance especially on small scales. Using Indra simulations, the study finds an optimal single-input scale near $L_ ext{sub}\sim152\,h^{-1}\mathrm{Mpc}$ across multiple metrics (loss, PDFs, KL divergence, transfer function, correlation), and demonstrates that the dual-input model significantly enhances small-scale fidelity (e.g., KL divergence dropping to $0.0027$ and higher $T(k)$ and $r(k)$ in key $k$-bins) while being more data-efficient. These results indicate that multi-scale context is crucial for accurate cosmological field reconstruction and lay groundwork for improved inference in large-scale structure surveys, with potential extensions to redshift space, halos/galaxies, and physics-informed hybrids.

Abstract

Reconstructing the initial density field of the Universe from the late-time matter distribution is a nontrivial task with implications for understanding structure formation in cosmology, offering insights into early Universe conditions. Convolutional neural networks (CNNs) have shown promise in tackling this problem by learning the complex mapping from nonlinear evolved fields back to initial conditions. Here we investigate the effect of varying input sub-box size in single-input CNNs. We find that intermediate scales ($L_\mathrm{sub} \sim 152\,h^{-1}\,\mathrm{Mpc}$) strike the best balance between capturing local detail and global context, yielding the lowest validation loss and most accurate recovery across multiple statistical metrics. We then propose a dual-input model that combines two sub-boxes of different sizes from the same simulation volume. This model significantly improves reconstruction performance, especially on small scales over the best single-input case, despite utilizing the same parent simulation box. This demonstrates the advantage of explicitly incorporating multi-scale context into the network. Our results highlight the importance of input scale and network design in reconstruction tasks. The dual-input approach represents a simple yet powerful enhancement that leverages fixed input information more efficiently, paving the way for more accurate cosmological inference from large-scale structure surveys.

Figures (9)

• Figure 1: Left : Loss functions defined in equation (\ref{['eq:loss_function']}). Solid and dotted lines represent the train and validation loss, respectively. Results are shown for three representative sub-box sizes: $L_\mathrm{sub}\sim76\,h^{-1}\,\mathrm{Mpc}$ (matching the setup of Mao2021; gray), $L_\mathrm{sub}\sim152\,h^{-1}\,\mathrm{Mpc}$ (intermediate size; red) and $L_\mathrm{sub}\sim380\,h^{-1}\,\mathrm{Mpc}$ (large size; blue). Right : Train and validation loss as a function of sub-box size $L_\mathrm{sub}$. White circles and error-bars denote the mean and standard deviation of the loss over $[4,5]\times10^5$ iterations. Black points indicate loss values computed from all grid points in a test simulation.
• Figure 2: Left : PDFs of the density fluctuation $\delta$. Solid, dotted and dash-dotted lines denote PDFs of the target $\delta_\mathrm{ini}(z=10)$, input $\delta_\mathrm{fin}(z=0)$ and output $\delta_\mathrm{rec}$ density distributions, respectively. Results are shown for three representative sub-box sizes: $L_\mathrm{sub}\sim76,\,152\,\mathrm{and}\,380\,h^{-1}\,\mathrm{Mpc}$. Right : KL divergence as a function of sub-box size $L_\mathrm{sub}$. KL divergence between the predicted and target initial density distributions, defined in equation (\ref{['eq:kl_divergence']}).
• Figure 3: Left : Transfer function between the power spectrum of the reconstructed and target initial density fields, defined in equation \ref{['eq:transfer']}. Results are shown for three representative sub-box sizes: $L_\mathrm{sub}\sim76,\,152\,\mathrm{and}\,380\,h^{-1}\,\mathrm{Mpc}$ from a single simulation. Right : Transfer function averaged over five wavenumber bins, as a function of sub-box size $L_\mathrm{sub}$.
• Figure 4: Left : Correlation coefficient between the power spectrum of the reconstructed and target initial density fields, defined in equation \ref{['eq:correlation']}. Results are shown for three representative sub-box sizes: $L_\mathrm{sub}\sim76,\,152\,\mathrm{and}\,380\,h^{-1}\,\mathrm{Mpc}$ from a single simulation. Right : Correlation coefficient averaged over five wavenumber bins, as a function of sub-box size $L_\mathrm{sub}$.
• Figure 5: Loss functions as defined in equation (\ref{['eq:loss_function']}). Results are shown for two representative sub-box sizes: $L_\mathrm{sub}\sim76\,h^{-1}\,\mathrm{Mpc}$ (gray) and $L_\mathrm{sub}\sim228\,h^{-1}\,\mathrm{Mpc}$ (green), as well as for dual-input model using $L_\mathrm{sub}\sim\{76,228\}\,h^{-1}\,\mathrm{Mpc}$ as inputs (red).