Learning to predict the cosmological structure formation

Significance

To understand the evolution of the Universe requires a concerted effort of accurate observation of the sky and fast prediction of structures in the Universe. N-body simulation is an effective approach to predicting structure formation of the Universe, though computationally expensive. Here, we build a deep neural network to predict structure formation of the Universe. It outperforms the traditional fast-analytical approximation and accurately extrapolates far beyond its training data. Our study proves that deep learning is an accurate alternative to the traditional way of generating approximate cosmological simulations. Our study shows that one can use deep learning to generate complex 3D simulations in cosmology. This suggests that deep learning can provide a powerful alternative to traditional numerical simulations in cosmology.

Abstract

Matter evolved under the influence of gravity from minuscule density fluctuations. Nonperturbative structure formed hierarchically over all scales and developed non-Gaussian features in the Universe, known as the cosmic web. To fully understand the structure formation of the Universe is one of the holy grails of modern astrophysics. Astrophysicists survey large volumes of the Universe and use a large ensemble of computer simulations to compare with the observed data to extract the full information of our own Universe. However, to evolve billions of particles over billions of years, even with the simplest physics, is a daunting task. We build a deep neural network, the Deep Density Displacement Model (${\mathrm{D}}^3\mathrm{M}$), which learns from a set of prerun numerical simulations, to predict the nonlinear large-scale structure of the Universe with the Zel’dovich Approximation (ZA), an analytical approximation based on perturbation theory, as the input. Our extensive analysis demonstrates that ${\mathrm{D}}^3\mathrm{M}$ outperforms the second-order perturbation theory (2LPT), the commonly used fast-approximate simulation method, in predicting cosmic structure in the nonlinear regime. We also show that ${\mathrm{D}}^3\mathrm{M}$ is able to accurately extrapolate far beyond its training data and predict structure formation for significantly different cosmological parameters. Our study proves that deep learning is a practical and accurate alternative to approximate 3D simulations of the gravitational structure formation of the Universe.

Astrophysicists require a large amount of simulations to extract the information from observations (1–8). At its core, modeling structure formation of the Universe is a computationally challenging task; it involves evolving billions of particles with the correct physical model over a large volume over billions of years (9–11). To simplify this task, we either simulate a large volume with simpler physics or a smaller volume with more complex physics. To produce the cosmic web (12) in large volume, we select gravity, the most important component of the theory, to simulate at large scales. A gravity-only $N$-body simulation is the most popular and effective numerical method to predict the full 6D phase-space distribution of a large number of massive particles whose position and velocity evolve over time in the Universe (13). Nonetheless, $N$-body simulations are relatively computationally expensive, thus making the comparison of the $N$-body–simulated large-scale structure (of different underlying cosmological parameters) with the observed Universe a challenging task. We propose to use a deep model that predicts the structure formation as an alternative to $N$-body simulations.

Deep learning (14) is a fast-growing branch of machine learning, where recent advances have led to models that reach and sometimes exceed human performance across diverse areas, from analysis and synthesis of images (15–17), sound (18, 19), text (20, 21), and videos (22, 23) to complex control and planning tasks as they appear in robotics and game play (24–26). This new paradigm is also significantly impacting a variety of domains in the sciences, from biology (27, 28) to chemistry (29, 30) and physics (31, 32). In particular, in astronomy and cosmology, a growing number of recent studies are using deep learning for a variety of tasks, ranging from analysis of cosmic microwave background (33–35), large-scale structure (36, 37), and gravitational lensing effects (38, 39) to classification of different light sources (40–42).

The ability of these models to learn complex functions has motivated many to use them to understand the physics of interacting objects, leveraging image, video, and relational data (43–53). However, modeling the dynamics of billions of particles in N-body simulations poses a distinct challenge.

In this work, we show that a variation on the architecture of a well-known deep-learning model (54) can efficiently transform the first-order approximations of the displacement field and approximate the exact solutions, thereby producing accurate estimates of the large-scale structure. Our key objective is to prove that this approach is an accurate and computationally efficient alternative to expensive cosmological simulations, and, to this end, we provide an extensive analysis of the results in the following section.

The outcome of a typical N-body simulation depends on both the initial conditions and on cosmological parameters which affect the evolution equations. A striking discovery is that the Deep Density Displacement Model (${\mathrm{D}}^3\mathrm{M}$), trained by using a single set of cosmological parameters, generalizes to new sets of significantly different parameters, minimizing the need for training data on a diverse range of cosmological parameters.

Setup

We build a deep neural network, ${\mathrm{D}}^3\mathrm{M}$, with similar input and output of an $N$-body simulation. The input to our ${\mathrm{D}}^3\mathrm{M}$ is the displacement field from the Zel’dovich Approximation (ZA) (55). A displacement vector is the difference of a particle position at target redshift $z=0$—i.e., the present time—and its Lagrangian position on a uniform grid. ZA evolves the particles on linear trajectories along their initial displacements. It is accurate when the displacement is small; therefore, ZA is frequently used to construct the initial conditions of $N$-body simulations (56). As for the ground truth, the target displacement field is produced by using FastPM (57), a recent approximate N-body–simulation scheme that is based on a particle-mesh (PM) solver. FastPM quickly approaches a full N-body simulation with high accuracy and provides a viable alternative to direct N-body simulations for the purpose of our study.

A significantly faster approximation of N-body simulations is produced by second-order Lagrangian perturbation theory (2LPT), which bends each particle’s trajectory with a quadratic correction (58). The 2LPT is used in many cosmological analyses to generate a large number of cosmological simulations for comparison of the astronomical dataset against the physical model (59, 60) or to compute the covariance of the dataset (61–63). We regard 2LPT as an effective way to efficiently generate a relatively accurate description of the large-scale structure, and therefore we select 2LPT as the reference model for comparison with ${\mathrm{D}}^3\mathrm{M}$.

We generate 10,000 pairs of ZAs as input and accurate FastPM approximations as the target. We use simulations of $32^3\hspace{0.17em}N$-body particles in a volume of $128\hspace{0.17em}{h}^{-1}\mathrm{M}\mathrm{p}\mathrm{c}$ (600 million light years, where $h=0.7$ is the Hubble parameter). The particles have a mean separation of $4\hspace{0.17em}{h}^{-1}\mathrm{M}\mathrm{p}\mathrm{c}$ per dimension.

An important choice in our approach is training with a displacement field rather than a density field. Displacement field $\mathbf{\Psi }$ and density field $\rho$ are two ways of describing the same distribution of particles. And an equivalent way to describe a density field is the overdensity field, defined as $\delta =\rho /\overline{\rho }-1$, with $\overline{\rho }$ denoting the mean density. The displacement field and overdensity field are related by Eq. 1.

$$\begin{array}{cc}\hfill \mathbf{x}& =\mathbf{\Psi }\left(\mathbf{q}\right)+\mathbf{q}\hfill \\ \hfill \delta \left(\mathbf{x}\right)& =\int d^{3}q{\delta }_D\left(\mathbf{x}-\mathbf{q}-\mathbf{\Psi }\left(\mathbf{q}\right)\right)-1.\hfill \end{array}$$ [1]

When the displacement field is small and has zero curl, the choice of overdensity vs. displacement field for the output of the model is irrelevant, as there is a bijective map between these two representations, described by the equation:

$$\mathbf{\Psi }=\int \frac{d^{3}k}{{\left(2\pi \right)}^3}{e}^{i\mathbf{k}\cdot \mathbf{q}}\frac{i\mathbf{k}}{k^2}\delta \left(\mathbf{k}\right).$$ [2]

However, as the displacements grow into the nonlinear regime of structure formation, different displacement fields can produce identical density fields (e.g., ref. 64). Therefore, providing the model with the target displacement field during the training eliminates the ambiguity associated with the density field. Our inability to produce comparable results when using the density field as our input and target attests that relevant information resides in the displacement field (SI Appendix, Fig. S1).

Results and Analysis

Fig. 1 shows the displacement vector field as predicted by ${\mathrm{D}}^3\mathrm{M}$ (Left) and the associated point-cloud representation of the structure formation (Right). It is possible to identify structures such as clusters, filaments, and voids in this point-cloud representation. We proceed to compare the accuracy of ${\mathrm{D}}^3\mathrm{M}$ and 2LPT compared with ground truth.

Fig. 1.

The displacement vector field (Left) and the resulting density field (Right) produced by ${\mathrm{D}}^3\mathrm{M}$. The vectors in Left are uniformly scaled down for better visualization.

Point-Wise Comparison.

Let $\mathbf{\Psi }\in {\mathbb{R}}^{d\times d\times d\times 3}$ denote the displacement field, where $d$ is the number of spatial-resolution elements in each dimension ($d=32$). A natural measure of error is the relative error $|\widehat{\mathbf{\Psi }}-{\mathbf{\Psi }}_t|/\left|{\mathbf{\Psi }}_t\right|$, where ${\mathbf{\Psi }}_t$ is the true displacement field (FastPM), and $\widehat{\mathbf{\Psi }}$ is the prediction from 2LPT or ${\mathrm{D}}^3\mathrm{M}$. Fig. 2 compares this error for different approximations in a 2D slice of a single simulation. We observe that ${\mathrm{D}}^3\mathrm{M}$ predictions are very close to the ground truth, with a maximum relative error of 1.10 over all 1,000 simulations. For 2LPT, this number is significantly higher at 4.23. In average, the result of ${\mathrm{D}}^3\mathrm{M}$ comes with a 2.8% relative error, while for 2LPT, it equals 9.3%.

Fig. 2.

The columns show 2D slices of full-particle distribution (Upper) and displacement vector (Lower) by various models: FastPM, the target ground truth, a recent approximate N-body simulation scheme that is based on a PM solver (A); ZA, a simple linear model that evolves particle along the initial velocity vector (B); 2LPT, a commonly used analytical approximation (C); and deep-learning model (${\mathrm{D}}^3\mathrm{M}$) as presented in this work (D). While FastPM (A) served as our ground truth, B–D include color for the points or vectors. The color indicates the relative difference $\left(q_{model}-q_{target}\right)/q_{target}$ between the target location (A) or displacement vector and predicted distributions by various methods (B–D). The error bar shows that denser regions have a higher error for all methods, which suggests that it is harder to predict highly nonlinear region correctly for all models: ${\mathrm{D}}^3\mathrm{M}$, 2LPT, and ZA. Our model ${\mathrm{D}}^3\mathrm{M}$ has the smallest differences between predictions and ground truth among the above models B–D.

Two-Point Correlation Comparison.

As suggested by Fig. 2, the denser regions seem to have a higher error for all methods—that is, more nonlinearity in structure formation creates larger errors for both ${\mathrm{D}}^3\mathrm{M}$ and 2LPT. The dependence of error on scale is computed with two- and three-point correlation analysis.

Cosmologists often use compressed summary statistics of the density field in their studies. The most widely used of these statistics are the two-point correlation function (2PCF) $\xi \left(r\right)$ and its Fourier transform, the power spectrum $P_{\delta \delta }\left(k\right)$:

$$\begin{array}{cc}\hfill \xi \left(\left|\mathit{r}\right|\right)& =⟨{\delta }_A\left(\mathit{r}\prime \right){\delta }_B\left(\mathit{r}\prime +\mathit{r}\right)⟩,\hfill \\ \hfill P_{\delta \delta }\left(\left|\mathit{k}\right|\right)& =\int {\mathrm{d}}^3\mathit{r}\hspace{0.17em}\xi \left(r\right)e^{i\mathit{k}\cdot \mathit{r}},\hfill \end{array}$$ [3]

where the ensemble average $⟨⟩$ is taken over all possible realizations of the Universe. Our Universe is observed to be both homogeneous and isotropic on large scales—i.e., without any special location or direction. This allows one to drop the dependencies on $\mathit{r}\prime$ and on the direction of $\mathit{r}$, leaving only the amplitude $\left|\mathit{r}\right|$ in the final definition of $\xi \left(r\right)$. In the second equation, $P_{\delta \delta }\left(k\right)$ is simply the Fourier transform of $\xi \left(r\right)$ and captures the dispersion of the plane-wave amplitudes at different scales in the Fourier space. $\mathit{k}$ is the 3D wavevector of the plane wave, and its amplitude $k$ (the wavenumber) is related to the wavelength $\lambda$ by $k=2\pi /\lambda$. Due to isotropy of the Universe, we drop the vector form of $\mathit{r}$ and $\mathit{k}$.

Because FastPM, 2LPT, and ${\mathrm{D}}^3\mathrm{M}$ take the displacement field as input and output, we also study the two-point statistics for the displacement field. The displacement power spectrum is defined as:

$$P_{\mathbf{\Psi }\mathbf{\Psi }}\left(k\right)=⟨{\mathrm{\Psi }}_x\left(k\right){\mathrm{\Psi }}_x^{*}\left(k\right)⟩+⟨{\mathrm{\Psi }}_y\left(k\right){\mathrm{\Psi }}_y^{*}\left(k\right)⟩+⟨{\mathrm{\Psi }}_z\left(k\right){\mathrm{\Psi }}_z^{*}\left(k\right)⟩.$$ [4]

We focus on the Fourier-space representation of the two-point correlation. Because the matter and the displacement power spectrum take the same form, in what follows, we drop the subscript for matter and displacement field and use $P\left(k\right)$ to stand for both matter and displacement power spectrum. We use the transfer function $T\left(k\right)$ and the correlation coefficient $r\left(k\right)$ as metrics to quantify the model performance against the ground truth (FastPM) in the two-point correlation. We define the transfer function $T\left(k\right)$ as the square root of the ratio of two power spectra,

$$T\left(k\right)=\sqrt{\frac{P_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}}\left(k\right)}{P_{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}\left(k\right)}},$$ [5]

where $P_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}}\left(k\right)$ is the density or displacement power spectrum as predicted by 2LPT or ${\mathrm{D}}^3\mathrm{M}$, and, analogously, $P_{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}\left(k\right)$ is the ground truth predicted by FastPM. The correlation coefficient r(k) is a form of normalized cross-power spectrum,

$$r\left(k\right)=\frac{P_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}\times \mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}\left(k\right)}{\sqrt{P_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}}\left(k\right)P_{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}\left(k\right)}},$$ [6]

where $P_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}\times \mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}\left(k\right)$ is the cross-power spectrum between 2LPT or ${\mathrm{D}}^3\mathrm{M}$ predictions and the ground-truth (FastPM) simulation result. The transfer function captures the discrepancy between amplitudes, while the correlation coefficient can indicate the discrepancy between phases as functions of scales. For a perfectly accurate prediction, $T\left(k\right)$ and $r\left(k\right)$ are both 1. In particular, $1-r^2$ describes stochasticity, the fraction of the variance in the prediction that cannot be explained by the true model.

Fig. 3A shows the average power spectrum, transfer function $T\left(k\right)$, and stochasticity $1-r^2\left(k\right)$ of the displacement field and the density field over 1,000 simulations. The transfer function of density from 2LPT predictions is 2% smaller than that of FastPM on large scales ($k\approx 0.05\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$). This is expected since 2LPT performs accurately on very large scales ($k<0.01\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$). The displacement transfer function of 2LPT increases >1 at $k\approx 0.35\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$ and then drops sharply. The increase of the 2LPT displacement transfer function is because 2LPT overestimates the displacement power at small scales (e.g., ref. 65). There is a sharp drop of power near the voxel scale because smoothing over voxel scales in our predictions automatically erases power at scales smaller than the voxel size.

Fig. 3.

(A) Displacement and density power-spectrum of FastPM (orange), 2LPT (blue), and c (green) (Top); transfer function—i.e., the square root of the ratio of the predicted power-spectrum to the ground truth (Middle); and 1 – $r^2$, where $r$ is the correlation coefficient between the predicted fields and the true fields (Bottom). Results are the averaged values of 1,000 test simulations. The transfer function and correlation coefficient of the ${\mathrm{D}}^3\mathrm{M}$ predictions are nearly perfect from large to intermediate scales and outperform our benchmark 2LPT significantly. (B) The ratios of the multipole coefficients (${\zeta }_l\left(r_1,r_2\right)$) (to the target) of the two 3PCFs for several triangle configurations. The results are averaged over 10 test simulations. The error bars (padded regions) are the SDs derived from 10 test simulations. The ratio shows that the 3PCF of ${\mathrm{D}}^3\mathrm{M}$ is closer than 2LPT to our target FastPM with lower variance.

Now, we turn to the ${\mathrm{D}}^3\mathrm{M}$ predictions: Both the density and displacement transfer functions of the ${\mathrm{D}}^3\mathrm{M}$ differ from 1 by a mere 0.4% at scale $k\lesssim 0.4\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$, and this discrepancy only increases to 2% and 4% for density field and displacement field, respectively, as $k$ increases to the Nyquist frequency at ?$0.7\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$. The stochasticity hovers at ∼$10^{-3}$ and $10^{-2}$ for most scales. In other words, for both the density and displacement fields, the correlation coefficient between the ${\mathrm{D}}^3\mathrm{M}$ predictions and FastPM simulations, all the way down to small scales of $k=0.7\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$, is >90%. The transfer function and correlation coefficient of the ${\mathrm{D}}^3\mathrm{M}$ predictions show that it can reproduce the structure formation of the Universe from large to seminonlinear scales. ${\mathrm{D}}^3\mathrm{M}$ significantly outperforms our benchmark model 2LPT in the two-point function analysis. ${\mathrm{D}}^3\mathrm{M}$ only starts to deviate from the ground truth at fairly small scales. This is not surprising, as the deeply nonlinear evolution at these scales is more difficult to simulate accurately and appears to be intractable by current analytical theories (66).

Three-Point Correlation Comparison.

The three-point correlation function (3PCF) expresses the correlation of the field of interest among three locations in the configuration space, which is equivalently defined as bispectrum in Fourier space. Here, we concentrate on the 3PCF for computational convenience:

$$\zeta \left(r_1,r_2,\theta \right)=⟨\delta \left(\mathbf{x}\right)\delta \left(\mathbf{x}+{\mathbf{r}}_{\mathbf{1}}\right)\delta \left(\mathbf{x}+{\mathbf{r}}_{\mathbf{2}}\right)⟩,$$ [7]

where $r_1$ = $\left|{\mathbf{r}}_{\mathbf{1}}\right|$ and $r_2$ = $\left|{\mathbf{r}}_{\mathbf{2}}\right|$. Translation invariance guarantees that $\zeta$ is independent of $\mathit{x}$. Rotational symmetry further eliminates all direction dependence except dependence on $\theta$, the angle between ${\mathit{r}}_{\mathbf{1}}$ and ${\mathit{r}}_{\mathbf{2}}$. The multipole moments of $\zeta \left(r_1,r_2,\theta \right)$, ${\zeta }_{\ell }\left(r_1,r_2\right)=\left(2\ell +1\right)\int d\theta P_{\ell }\left(\cos \theta \right)\zeta \left(r_1,r_2,\theta \right)$, where $P_{\ell }\left(\cos \theta \right)$ is the Legendre polynomial of degree $\ell$, can be efficiently estimated with pair counting (67). While the input (computed by ZA) do not contain significant correlations beyond the second order (power spectrum level), we expect ${\mathrm{D}}^3\mathrm{M}$ to generate densities with a 3PCF that mimics that of ground truth.

We compare the 3PCF calculated from FastPM, 2LPT, and ${\mathrm{D}}^3\mathrm{M}$ by analyzing the 3PCF through its multipole moments ${\zeta }_{\ell }\left(r_1,r_2\right)$. Fig. 3B shows the ratio of the binned multipole coefficients of the two 3PCFs for several triangle configurations, $\overline{{\xi }_{\ell }}{\left(r_1,r_2\right)}_{pred}/\overline{{\xi }_{\ell }}{\left(r_1,r_2\right)}_{true}$, where $\overline{{\xi }_{\ell }}{\left(r_1,r_2\right)}_{pred}$ can be the 3PCF for ${\mathrm{D}}^3$M or 2LPT and $\overline{{\xi }_{\ell }}{\left(r_1,r_2\right)}_{true}$ is the 3PCF for FastPM. We used 10 radial bins with $\mathrm{\Delta }r=5\hspace{0.17em}{h}^{-1}\mathrm{M}\mathrm{p}\mathrm{c}$. The results are averaged over 10 test simulations, and the error bars are the SD. The ratio shows that the 3PCF of ${\mathrm{D}}^3$M is closer to FastPM than 2LPT, with smaller error bars. To further quantify our comparison, we calculate the relative 3PCF residual defined by

$$\begin{array}{ccc}\hfill & \hfill & 3\text{PCF relative residual}\hfill \\ \hfill & \hfill & =\frac{1}{9\times N_r}\sum _{\ell =0}^8\sum _{r_1,r_2}\frac{|{\zeta }_{\ell }{\left(r_1,r_2\right)}_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}}-{\zeta }_{\ell }{\left(r_1,r_2\right)}_{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}|}{\left|{\zeta }_{\ell }{\left(r_1,r_2\right)}_{\mathrm{t}\mathrm{r}\mathrm{u}\mathrm{e}}\right|},\hfill \end{array}$$ [8]

where $N_r$ is the number of ($r_1,r_2$) bins. The mean relative 3PCF residual of the ${\mathrm{D}}^3\mathrm{M}$ and 2LPT predictions compared with FastPM are $0.79\%$ and $7.82\%$, respectively. The ${\mathrm{D}}^3\mathrm{M}$ accuracy on 3PCF is also an order of magnitude better than 2LPT, which indicates that the ${\mathrm{D}}^3\mathrm{M}$ is far better at capturing the non-Gaussian structure formation.

Generalizing to New Cosmological Parameters

So far, we train our model using a “single” choice of cosmological parameters $A_{\mathrm{s}}=2.142\times 10^{-9}$ (hereafter $A_0=2.142\times 10^{-9}$) and ${\mathrm{\Omega }}_m=0.3089$ (68). $A_{\mathrm{s}}$ is the primordial amplitude of the scalar perturbation from cosmic inflation, and ${\mathrm{\Omega }}_m$ is the fraction of the total energy density that is matter at the present time, and we will call it matter density parameter for short. The true exact value of these parameters is unknown, and different choices of these parameters change the large-scale structure of the Universe; Fig. 4.

Fig. 4.

We show the differences of particle distributions and displacement fields when we change the cosmological parameters $A_s$ and ${\mathrm{\Omega }}_m$. (A) The error bar shows the difference of particle distribution (Upper) and displacement fields (Lower) between $A_s=A_0$ and the two extremes for A_s = 0.2 A₀ (Center) and $A_s=1.8A_0$ (Right). (B) A similar comparison showing the difference of the particle distributions (Upper) and displacement fields (Lower) for smaller and larger values of ${\mathrm{\Omega }}_m\in \left\{0.1,0.5\right\}$ with regard to ${\mathrm{\Omega }}_m=0.3089$, which was used for training. While the difference for the smaller value of $A_s$ (${\mathrm{\Omega }}_m$) is larger, the displacement for the larger $A_s$ (${\mathrm{\Omega }}_m$) is more nonlinear. This nonlinearity is due to concentration of mass and makes the prediction more difficult.

Here, we report an interesting observation: The ${\mathrm{D}}^3\mathrm{M}$ trained on a single set of parameters in conjunction with ZA (which depends on $A_{\mathrm{s}}$ and ${\mathrm{\Omega }}_m$) as input can predict the structure formation for widely different choices of $A_{\mathrm{s}}$ and ${\mathrm{\Omega }}_m$. From a computational point of view, this suggests a possibility of producing simulations for a diverse range of parameters, with minimal training data.

Varying Primordial Amplitude of Scalar Perturbations $A_{\mathrm{s}}$.

After training the ${\mathrm{D}}^3\mathrm{M}$ using $A_{\mathrm{s}}=A_0$, we change $A_{\mathrm{s}}$ in the input of our test set by nearly one order of magnitude: $A_s=1.8A_0$ and $A_{\mathrm{s}}=0.2A_0$. Again, we use 1,000 simulations for analysis of each test case. The average relative displacement error of ${\mathrm{D}}^3\mathrm{M}$ remains <$4\%$ per voxel (compared with $<3\%$ when train and test data have the same parameters). This is still well below the error for 2LPT, which has relative errors of $15.5\%$ and $6.3\%$ for larger and smaller values of $A_s$, respectively.

Fig. 5A shows the transfer function and correlation coefficient for both ${\mathrm{D}}^3\mathrm{M}$ and 2LPT. The ${\mathrm{D}}^3\mathrm{M}$ performs much better than 2LPT for $A_s=1.8A_0$. For small $A_s=0.2A_0$, 2LPT does a better job than ${\mathrm{D}}^3\mathrm{M}$ predicting the density transfer function and correlation coefficient at the largest scales; otherwise, ${\mathrm{D}}^3\mathrm{M}$ predictions are more accurate than 2LPT at scales larger than $k=0.08\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$. We observe a similar trend with 3PCF analysis: The 3PCFs of ${\mathrm{D}}^3\mathrm{M}$ predictions are notably better than 2LPT ones for larger $A_s$, compared with smaller $A_s$, where it is only slightly better. These results confirm our expectation that increasing $A_s$ increases the nonlinearity of the structure-formation process. While 2LPT can predict fairly well in linear regimes, compared with ${\mathrm{D}}^3\mathrm{M}$, its performance deteriorates with increased nonlinearity. It is interesting to note that the ${\mathrm{D}}^3\mathrm{M}$ prediction maintains its advantage, despite being trained on data from more linear regimes.

Fig. 5.

Similar plots as in Fig. 3A, except that we test the two-point statistics when we vary the cosmological parameters without changing the training set (which has different cosmological parameters) or the trained model. We show predictions from ${\mathrm{D}}^3\mathrm{M}$ and 2LPT when tested on different $A_{\mathrm{s}}$ (A) and ${\mathrm{\Omega }}_m$ (B). We show the transfer function—i.e., the square root of the ratio of the predicted power spectrum to the ground truth (Upper)—and 1 – $r^2$, where $r$ is the correlation coefficient between the predicted fields and the true fields (Lower). The ${\mathrm{D}}^3\mathrm{M}$ prediction outperforms 2LPT prediction at all scales except in the largest scales, as the perturbation theory works well in linear regime (large scales).

Varying Matter Density Parameter ${\mathbf{\Omega }}_{\mathbf{m}}$.

We repeat the same experiments, this time changing ${\mathrm{\Omega }}_m$ to 0.5 and 0.1, while the model is trained on ${\mathrm{\Omega }}_m=0.3089$, which is quite far from both of the test sets. For ${\mathrm{\Omega }}_m=0.5$, the relative residual displacement errors of the ${\mathrm{D}}^3\mathrm{M}$ and 2LPT averaged over 1,000 simulations are 3.8% and 15.2%, and for ${\mathrm{\Omega }}_m=0.1$, they are 2.5% and 4.3%. Fig. 5 C and D show the two-point statistics for density field predicted by using different values of ${\mathrm{\Omega }}_m$. For ${\mathrm{\Omega }}_m=0.5$, the results show that the ${\mathrm{D}}^3\mathrm{M}$ outperforms 2LPT at all scales, while for smaller ${\mathrm{\Omega }}_m=0.1$, ${\mathrm{D}}^3\mathrm{M}$ outperforms 2LPT on smaller scales ($k>0.1\hspace{0.17em}h{\mathrm{M}\mathrm{p}\mathrm{c}}^{-1}$). As for the 3PCF of simulations with different values of ${\mathrm{\Omega }}_m$, the mean relative 3PCF residual of the ${\mathrm{D}}^3\mathrm{M}$ for ${\mathrm{\Omega }}_m=0.5$ and ${\mathrm{\Omega }}_m=0.1$ are 1.7% and 1.2%, respectively, and for 2LPT, they are 7.6% and 1.7%, respectively. The ${\mathrm{D}}^3\mathrm{M}$ prediction performs better at ${\mathrm{\Omega }}_m=0.5$ than ${\mathrm{\Omega }}_m=0.1$. This is again because the Universe is much more nonlinear at ${\mathrm{\Omega }}_m=0.5$ than ${\mathrm{\Omega }}_m=0.1$. The ${\mathrm{D}}^3\mathrm{M}$ learns more nonlinearity than is encoded in the formalism of 2LPT.

Conclusions

To summarize, our deep model ${\mathrm{D}}^3\mathrm{M}$ can accurately predict the large-scale structure of the Universe as represented by FastPM simulations, at all scales, as seen in Table 1. Furthermore, ${\mathrm{D}}^3\mathrm{M}$ learns to predict cosmic structure in the nonlinear regime more accurately than our benchmark model 2LPT. Finally, our model generalizes well to test simulations with cosmological parameters ($A_{\mathrm{s}}$ and ${\mathrm{\Omega }}_m$) significantly different from the training set. This suggests that our deep-learning model can potentially be deployed for a range of simulations beyond the parameter space covered by the training data (Table 1). Our results demonstrate that the ${\mathrm{D}}^3\mathrm{M}$ successfully learns the nonlinear mapping from first-order perturbation theory to FastPM simulation beyond what higher-order perturbation theories currently achieve.

Table 1.

A summary of our analysis

Data	Point-wise	$T\left(k\right)$k = 0.11	$r\left(k\right)k=0.11$	$T\left(k\right)k=0.50$	$r\left(k\right)k=0.50$	3PCF
Test phase
2LPT density	N/A	0.96	1.00	0.74	0.94	0.0782
${\mathrm{D}}^3\mathrm{M}$ density	N/A	1.00	1.00	0.99	1.00	0.0079
2LPT displacement	0.093	0.96	1.00	1.04	0.90	N/A
${\mathrm{D}}^3\mathrm{M}$ displacement	0.028	1.00	1.00	0.99	1.00	N/A
${\mathit{A}}_{\mathit{s}}=\mathbf{1.8}{\mathit{A}}_{\mathbf{0}}$
2LPT density	N/A	0.93	1.00	0.49	0.78	0.243
${\mathrm{D}}^3\mathrm{M}$ density	N/A	1.00	1.00	0.98	1.00	0.039
2LPT displacement	0.155	0.97	1.00	1.07	0.73	N/A
${\mathrm{D}}^3\mathrm{M}$ displacement	0.039	1.00	1.00	0.97	0.99	N/A
$A_{\mathrm{s}}=0.2A_0$
2LPT density	N/A	0.99	1.00	0.98	0.99	0.024
${\mathrm{D}}^3\mathrm{M}$ density	N/A	1.00	1.00	1.03	1.00	0.022
2LPT displacement	0.063	0.99	1.00	0.95	0.98	N/A
${\mathrm{D}}^3\mathrm{M}$ displacement	0.036	1.00	1.00	1.01	1.00	N/A
${\mathbf{\Omega }}_{\mathbf{m}}=\mathbf{0.5}$
2LPT density	N/A	0.94	1.00	0.58	0.87	0.076
${\mathrm{D}}^3\mathrm{M}$ density	N/A	1.00	1.00	1.00	1.00	0.017
2LPT displacement	0.152	0.97	1.00	1.10	0.80	N/A
${\mathrm{D}}^3\mathrm{M}$ displacement	0.038	1.00	1.00	0.98	0.99	N/A
${\mathbf{\Omega }}_{\mathbf{m}}=\mathbf{0.1}$
2LPT density	N/A	0.97	1.00	0.96	0.99	0.017
${\mathrm{D}}^3\mathrm{M}$ density	N/A	0.99	1.00	1.04	1.00	0.012
2LPT displacement	0.043	0.97	1.00	0.97	0.98	N/A
${\mathrm{D}}^3\mathrm{M}$ displacement	0.025	0.99	1.00	1.02	1.00	N/A

The unit of k is hMpc⁻¹. N/A, not applicable.

Looking forward, we expect that replacing FastPM with exact N-body simulations would improve the performance of our method. As the complexity of our ${\mathrm{D}}^3\mathrm{M}$ model is linear in the number of voxels, we expect to be able to further improve our results if we replace the FastPM simulations with higher-resolution simulations. Our work suggests that deep learning is a practical and accurate alternative to the traditional way of generating approximate simulations of the structure formation of the Universe.

Materials and Methods

Dataset.

The full simulation data consists of 10,000 simulations of boxes with ZA and FastPM as input–output pairs, with an effective volume of 20 (Gpc/h)³ ($190\times 10^9{\mathrm{l}\mathrm{y}}^3$), comparable to the volume of a large spectroscopic sky survey like Dark Energy Spectroscopic Instrument or EUCLID. We split the full simulation dataset into 80%, 10%, and 10% for training, validation, and test, respectively. We also generated 1,000 simulations for 2LPT for each set of tested cosmological parameters.

Model and Training.

The ${\mathrm{D}}^3\mathrm{M}$ adopts the U-Net architecture (54) with 15 convolution or deconvolution layers and ∼$8.4\times 10^6$ trainable parameters. Our ${\mathrm{D}}^3\mathrm{M}$ generalizes the standard U-Net architecture to work with 3D data (69–71). The details of the architecture are described in the following sections, and a schematic figure of the architecture is shown in SI Appendix, Fig. S2. In the training phase, we use the Adam Optimizer (72) with a learning rate of 0.0001, and first- and second-moment exponential decay rates equal to 0.9 and 0.999, respectively. We use the mean-squared error as the loss function (Loss Function) and $L2$ regularization with regularization coefficient 0.0001.

Details of the ${\mathrm{D}}^3\mathrm{M}$ Architecture.

The contracting path follows the typical architecture of a convolution network. It consists of two blocks, each of which consists of two successive convolutions of stride 1 and a down-sampling convolution with stride 2. The convolution layers use 3$\times 3\times$3 filters with a periodic padding of size 1 (Padding and Periodic Boundary) on both sides of each dimension. Notice that at each of the two down-sampling steps, we double the number of feature channels. At the bottom of the ${\mathrm{D}}^3\mathrm{M}$, another two successive convolutions with stride 1 and the same periodic padding as above are applied. The expansive path of our ${\mathrm{D}}^3\mathrm{M}$ is an inverted version of the contracting path of the network. (It includes two repeated applications of the expansion block, each of which consists of one up-sampling–transposed convolution with stride 1/2 and two successive convolutions of stride 1. The transposed convolution and the convolution are constructed with 3$\times 3\times$3 filters.)

We take special care in the padding and cropping procedure to preserve the shifting and rotation symmetry in the up-sampling layer in expansive path. Before the transposed convolution, we apply a periodic padding of length 1 on the right, down, and back sides of the box [padding = (0,1,0,1,0,1) in pytorch], and after the transposed convolution, we discard one column on the left, up, and front sides of the box and two columns on the right, down, and back sides [crop = (1,2,1,2,1,2)].

A special feature of the ${\mathrm{D}}^3\mathrm{M}$ is the concatenation procedure, where the up-sampling layer halves the feature channels and then concatenates them with the corresponding feature channels on the contracting path, doubling the number of feature channels.

The expansive building block then follows a 1$\times 1\times$1 convolution without padding, which converts the 64 features to the final 3D displacement field. All convolutions in the network except the last one are followed by a rectified linear unit activation and batch normalization.

Padding and Periodic Boundary.

It is common to use constant or reflective padding in deep models for image processing. However, these approaches are not suitable for our setting. The physical model we are learning is constructed on a spatial volume with a periodic boundary condition. This is sometimes also referred to as a torus geometry, where the boundaries of the simulation box are topologically connected—that is, $x_{i+L}=x_i$, where $i$ is the index of the spatial location, and $L$ is the periodicity (size of box). Constant or reflective padding strategies break the connection between the physically nearby points separated across the box, which not only loses information but also introduces noise during the convolution, further aggravated with an increased number of layers.

We find that the periodic padding strategy significantly improves the performance and expedites the convergence of our model, comparing to the same network using a constant padding strategy. This is not surprising, as one expects that it is easier to train a model that can explain the data than to train a model that does not.

Loss Function.

We train the ${\mathrm{D}}^3\mathrm{M}$ to minimize the mean square error on particle displacements

$$\mathcal{L}=\frac{1}{N}\sum _i{\left({\widehat{\mathbf{\Psi }}}_i-{\mathbf{\Psi }}_{t,i}\right)}^2,$$ [9]

where $i$ labels the particles and N is the total number of particles. This loss function is proportional to the integrated squared error, and by using a Fourier transform and Parseval’s theorem, it can be rewritten as

$$\begin{array}{ll}\hfill & \int {\left(\widehat{\mathbf{\Psi }}-{\mathbf{\Psi }}_t\right)}^2{\mathrm{d}}^3\mathit{q}=\int |\widehat{\mathbf{\Psi }}-{\mathbf{\Psi }}_t{|}^2{\mathrm{d}}^3\mathit{k}=\hfill \\ \hfill & \int {\mathrm{d}}^3\mathit{k}\left(|{\mathbf{\Psi }}_t{|}^2{\left(1-T\right)}^2+2\left|\widehat{\mathbf{\Psi }}\right|\left|{\mathbf{\Psi }}_t\right|\left(1-r\right)\right),\hfill \end{array}$$ [10]

where $\mathit{q}$ is the Lagrangian space position, and $\mathit{k}$ its corresponding wavevector. $T$ is the transfer function defined in Eq. 5, and $r$ is the correlation coefficient defined in Eq. 6, which characterize the similarity between the predicted and true fields, in amplitude and phase, respectively. Eq. 10 shows that our simple loss function jointly captures both of these measures: As $T$ and $r$ approach 1, the loss function approaches 0.

Data Availability.

The source code of our implementation is available at https://github.com/siyucosmo/ML-Recon. The code to generate the training data is also available at https://github.com/rainwoodman/fastpm.