Simple lessons from complex learning: what a neural network model learns about cosmic structure formation

TL;DR

This work trains a convolutional neural network to predict the full nonlinear phase-space evolution of cosmological dark matter from initial conditions, effectively learning a Green's-function–like expansion that links initial Gaussian fields to late-time outputs. The model preserves large-scale Zel'dovich approximation behavior while nonlinearizing small-scale dynamics, and is validated against structured tests including spherical collapse, isolated plane waves, and coupled plane waves, where it generalizes beyond its Gaussian training data. The CNN achieves percent-level accuracy in nonlinear regimes around $k \sim 1 Mpc^{-1} h$ and outperforms the fast COLA method, demonstrating strong potential for accelerating precision cosmology while providing diagnostic insight into mode coupling and network limitations. The results highlight both the promise and the caveats of physics-informed neural emulators for cosmic structure formation, with practical implications for rapid parameter inference and survey analyses.

Abstract

We train a neural network model to predict the full phase space evolution of cosmological N-body simulations. Its success implies that the neural network model is accurately approximating the Green's function expansion that relates the initial conditions of the simulations to its outcome at later times in the deeply nonlinear regime. We test the accuracy of this approximation by assessing its performance on well understood simple cases that have either known exact solutions or well understood expansions. These scenarios include spherical configurations, isolated plane waves, and two interacting plane waves: initial conditions that are very different from the Gaussian random fields used for training. We find our model generalizes well to these well understood scenarios, demonstrating that the networks have inferred general physical principles and learned the nonlinear mode couplings from the complex, random Gaussian training data. These tests also provide a useful diagnostic for finding the model's strengths and weaknesses, and identifying strategies for model improvement. We also test the model on initial conditions that contain only transverse modes, a family of modes that differ not only in their phases but also in their evolution from the longitudinal growing modes used in the training set. When the network encounters these initial conditions that are orthogonal to the training set, the model fails completely. In addition to these simple configurations, we evaluate the model's predictions for the density, displacement, and momentum power spectra with standard initial conditions for N-body simulations. We compare these summary statistics against N-body results and an approximate, fast simulation method called COLA. Our model achieves percent level accuracy at nonlinear scales of $k\sim 1\ \mathrm{Mpc}^{-1}\, h$, representing a significant improvement over COLA.

Figures (8)

• Figure 1: The top row shows the matter density power spectra (left) from the CNN displacement model (blue, solid) and COLA (red, dashed), the Lagrangian displacement field (middle), and the momentum field (right). The momentum field depends on both particle displacements and velocities. The displacements and momenta have been rescaled to coincide with the matter power spectrum on large scales. The middle row shows the stochasticities with respect to the N-body simulations, and the bottom row shows the relative error in the transfer functions.
• Figure 2: Spherical evolution test displacement fields. Our neural network models take linear displacement fields as input and predict nonlinear outputs to emulate N-body simulations. We consider test cases with spherical symmetry that admit exact solutions. The left panel shows the linearly varying $y$ displacement of particles in a 2D slice through the center of a $50\,\mathrm{Mpc}\,h^{-1}$ sphere in the Lagrangian space. The inner particles are set up in a linear overdensity of 1.55 so they collapse into a smaller region at redshift $z=0$ without shell crossing. The right panel shows the nonlinear network prediction of the same field. The nonlinear displacements also vary linearly, and are much larger in magnitude to create a nonlinear density of $\approx20$.
• Figure 3: Spherical evolution density relation. Underdense spherical top hats expand and overdense ones collapse. Depending on their linear overdensities, the top hats evolve into different nonlinear densities at redshift 0, which can be solved exactly before shell crossing. We compare the neural network predictions to the exact solutions and find good agreement, which is further improved if we avoid the fuzzy edges by only using the inner part (within 0.8 times the radius) of the sphere to estimate the nonlinear relative density. The fuzzy edges can be seen in the right panel of Fig. \ref{['fig:sphere']}.
• Figure 4: Displacement divergence amplitudes from isolated longitudinal modes. Top panels show the N-body simulation results, middle panels show the neural network predictions, and bottom panels show the ratio of neural network to N-body mode amplitudes. The left two columns are for longitudinal modes with displacements only along the $x$ direction. The right two columns are for modes with equal displacement amplitudes along the $x$ and $y$ directions and double these displacements in the $z$ direction. The dashed lines in the top two rows show the expected amplitudes from theory. The circled data points correspond to the original input modes in the initial conditions. The first and third columns show the input mode and even harmonic modes, while the second and fourth columns show the odd harmonics. The higher harmonics are absent in the continuum limit and arise from the discreteness of the N-body systems, which the neural network learns from the N-body training data.
• Figure 5: Displacement divergence mode amplitudes from N-body simulations and neural network predictions for a fixed wave vector and varied initial amplitude. The dashed line shows the expected amplitude from linear theory.