NECOLA: Toward a Universal Field-level Cosmological Emulator

We made use of the Quijote full N-body simulations (Villaescusa-Navarro et al. 2020) to both train and test the model. The simulations used in this work follow the evolution of 512³ cold dark matter (CDM) particles (plus 512³ neutrino particles in the case of massive neutrino cosmologies) from z = 127 down to z = 0 in a periodic volume of ${(1000\,{h}^{-1}\mathrm{Mpc})}^{3}$. We train the network using a set of 100 simulations from the fiducial cosmology, where the values of the cosmological parameters are fixed to Ω_m = 0.3175, Ω_b = 0.049, h = 0.6711, n_s = 0.9624, σ₈ = 0.834, w = −1, and M_ν = 0.0 eV. These simulations are only different in the value of the initial random seed.

We test the accuracy of our network on simulations with very different cosmologies to the one used in the training. For this, we made use of 100 out of 2000 simulations of the latin hypercube contained in the Quijote simulations, where the values of the cosmological parameters span the range Ω_m ∈ [0.1, 0.5], Ω_b ∈ [0.03, 0.07], h ∈ [0.5, 0.9], n_s ∈ [0.8, 1.2], and σ₈ ∈ [0.6, 1.0]. In these simulations, not only are the set of values of the cosmological parameters different, but the initial random seed varies as well. Furthermore, we test the accuracy of our network on models with massive neutrinos and on models where the dark energy equation of state is w ≠ − 1, making use of Quijote simulations labeled ${M}_{\nu }^{+}$, ${M}_{\nu }^{++}$, ${M}_{\nu }^{+++}$, w⁺, and w⁻¹. On average, each of the Quijote simulations used in this work required ∼500 Central Processing Unit (CPU) hours to run. We refer the reader to Villaescusa-Navarro et al. (2020) for further details on the Quijote simulations.

The fast and approximate simulations we use in this work are run with the COmoving Lagrangian Acceleration (COLA; Tassev et al. 2013) method, which combines second-order Lagrangian perturbation theory (2LPT; Bernardeau et al. 2002) on large scales with N-body methods on small scales. In particular, we use the MG-PICOLA (Wright et al. 2017) package. For each Quijote full N-body simulation, we run a COLA simulation by matching (1) the number of particles, (2) the set of values of the cosmological parameters, and (3) the value of the initial random seed, which gives rise to an identical initial Gaussian field for both Quijote and COLA. These simulations require fewer time steps than the full N-body simulations and are therefore much more computationally efficient. Each COLA simulation is run with 30 time steps equally spaced in log from z = 9 down to z = 0. On average, these simulations only take 3 CPU hours to run.

Let us write the displacement vector of a particle as d = x _f − x _i , where x _f and x _i are the final (z = 0) and initial (Lagrangian) position of the particle. Our goal is to train a neural network to correct the positions of the particles generated by COLA, to match them with those from a full N-body simulation, i.e.,

$$\begin{eqnarray}&&{{\boldsymbol{x}}}_{f,\mathrm{Nbody}}=g({{\boldsymbol{x}}}_{f,\mathrm{COLA}}),\end{eqnarray} \tag{ 1 }$$

where g is an unknown function. Note that the right-hand side of Equation (1) should not be taken as the position of the particular particle considered, but also of all its neighboring particles. To preserve translational equivariance, we use displacement vectors instead of absolute particle positions. Thus, the input to the network is d _COLA, rather than x _f,COLA. The network is trained to learn d _Nbody − d _COLA = x _f,Nbody − x _f,COLA.

We follow Alves de Oliveira et al. (2020) and use a V-Net-based model (Milletari et al. 2016) that consists of two downsampling and three upsampling layers connected in a "V" shape. Blocks of two 3³ convolutions connect the input, the resampling, and the output layers; 1³ convolutions are added over each of these convolution blocks to realize a residual connection. We add batch normalization after every convolution except after the first one and the last two, and leaky ReLU activation with a negative slope of 0.01 after every batch normalization, as well as after the first and the second-to-last convolutions. The last activation in each residual block acts after the summation, following Milletari et al. (2016). As in U-Net/V-Net, at all resolution levels (with the exception of the bottleneck levels), the inputs to the downsampling layers are concatenated to the outputs of the upsampling layers. All layers have a channel size of 64, except for the input and the output that have three channels (the displacement vector along each Cartesian coordinate), as well as those after the concatenations (128-channeled). Finally, the input ( d _COLA) is directly added to the output, so that the network can learn the corrections to match the target ( d _Nbody − d _COLA). Stride-2 2³ convolutions and stride-1/2 2³ transposed convolutions are used in downsampling and upsampling layers, respectively. A diagram of the network architecture is shown in Figure 1.

Zoom In Zoom Out Reset image size

Following Alves de Oliveira et al. (2020), we minimize a loss function given by $L={\mathrm{log}}_{e}({{\rm{L}}}_{\delta }{{\rm{L}}}_{{\rm{\Delta }}}^{\lambda })$, where L_δ is the mean squared error (MSE) loss on n( x ) (the particle number in voxel x ) and L_Δ is the MSE on the displacement vector d . With this loss function, we are able to train the model to make accurate predictions in both Lagrangian and Eulerian spaces. By combining the two losses with logarithm rather than summation, we can account for their absolute magnitudes and trade between their relative values; λ here serves as a weight on this trade-off of relative losses and λ = 1 works pretty well in our case.

The input cannot be fed into the network all at once due to the big size of the data (3 × 512³), and we thus divide it into smaller chunks first. We crop the data into subcubes of size 3 × 128³, corresponding to a simulation box of length 250 h⁻¹Mpc. In order to preserve the physical translational equivariance, no padding has been used in the 3³ convolutions, which results in an output that is smaller than the input in spatial size. This limitation is compensated by padding the input cubes periodically with 20 voxels on each side so that the effective spatial size of the input becomes 3 × 168³. Furthermore, data augmentation is implemented to enforce the equivariance of displacement fields under rotational and parity transformations. We use the Adam optimizer (Kingma & Ba 2017) with a learning rate of 0.0001, β₁ = 0.9, and β₂ = 0.999, and reduce the learning rate by half when the loss does not improve for 3 epochs. The model is trained on 70 realizations for 100 epochs and the remaining realizations are used for validation (20) and final testing (10). From now on, we will refer to this model as NECOLA, from Neural Enhanced COLA, in order to avoid any confusion with the model by Alves de Oliveira et al. (2020), which uses Zel'dovich simulations as input and a different value of λ. Note that the model architecture of NECOLA is the same as that of Alves de Oliveira et al. (2020).

Before quantifying the accuracy of the network using summary statistics, we perform a visual inspection of its output. In Figure 2, we show the distribution of matter at z = 0 from the full N-body simulation (top row), the COLA simulation (middle row), and NECOLA (bottom row).

Zoom In Zoom Out Reset image size

While looking at large scales, the agreement between the three methods is really good, but when we look at small scales, some differences are visible. In the case of COLA, the output is more diffuse and halos do not exhibit a high concentration in their centers, in contrast to the corresponding N-body simulation. On the other hand, NECOLA produces much sharper results, clearly defining the positions and boundaries of dark matter halos.

The power spectrum is defined as the Fourier transform of the two-point correlation function, which measures the excess probability of finding a pair of random galaxies (or points) at a given separation compared to the one from a random distribution. The power spectrum is one of the most important summary statistics used in cosmology because for Gaussian density fields (like the one our universe resembles on large, linear scales), it fully characterizes the statistical properties of the field.

In the top-left panel of Figure 3, we show with a solid black line the average power spectra from 10 Quijote simulations of the test set. The dotted blue line shows the average power spectrum from the corresponding COLA simulations, while the green dotted–dashed line outputs the average power spectrum of Zel'dovich-evolved simulations. The solid yellow and dashed red lines show the average power spectrum from mod(ZA) and NECOLA, respectively. As can be seen, the worst model is the one that only employs the Zel'dovich approximation, followed by the COLA simulation.

Zoom In Zoom Out Reset image size

In order to better visualize the differences between the output of the N-body simulation and the networks, we plot in the middle-left panel of Figure 3, the transfer function, defined as

$$\begin{eqnarray}&&T(k)=\sqrt{\displaystyle \frac{{P}_{\mathrm{pred}}(k)}{{P}_{\mathrm{target}}(k)}},\end{eqnarray} \tag{ 2 }$$

where P_pred(k) and P_target(k) are the average matter power spectra of the predictions and the target density fields, respectively. Values close to one indicate a better agreement between the prediction and the target. As can be seen, both networks achieve a subpercent accuracy on the power spectrum down to k = 1 hMpc⁻¹, though the results obtained from NECOLA are slightly more accurate. We note that in the case of the Quijote simulations, it does not make sense to look into much smaller scales than k ∼ 1 hMpc⁻¹, as those are not numerically converged in the simulations due to mass resolution (Villaescusa-Navarro et al. 2020).

We note that there exists state-of-the-art power spectrum emulators such as COSMIC EMU (Heitmann et al. 2009, 2010; Lawrence et al. 2010), FRANKEN EMU (Heitmann et al. 2014), and MIRA TITAN (Heitmann et al. 2016; Lawrence et al. 2017) that are computationally faster and more accurate in estimating the power spectrum but are not used in our comparisons as the primary objective of our work is to provide a field-level emulator itself and not a power spectrum emulator.

In Fourier space, every mode can be written as δ( k ) = Ae^{i θ} , where A and θ are the mode amplitude and phase, respectively. When using the power spectrum, we are effectively comparing how well the amplitude of the modes from the network and the simulation agree. However, that statistic neglects the correlations in mode phases, which are very important in the nonlinear regime. To quantify the correlations between the mode phases, we use the cross-correlation coefficient, r, defined as

$$\begin{eqnarray}&&r(k)=\displaystyle \frac{{P}_{\mathrm{pred}\times \mathrm{target}}(k)}{\sqrt{{P}_{\mathrm{pred}}(k){P}_{\mathrm{target}}(k)}},\end{eqnarray} \tag{ 3 }$$

where the numerator is the cross-power spectrum between the predictions and the target, and the denominator contains the autopower spectrum of the prediction and the target. Values of r close to one indicate a very good correlation in mode phases. In the bottom-left panel of Figure 3, we show the cross-correlation coefficient averaged over the testing set for the different cases considered. We find that NECOLA achieves the highest accuracy, being within 1% down to k = 1 hMpc⁻¹.

The third statistic that we consider to quantify the agreement between the full simulations and the network predictions is the bispectrum, defined as

$$\begin{eqnarray}&&\langle {\delta }_{{{\boldsymbol{k}}}_{1}}{\delta }_{{{\boldsymbol{k}}}_{2}}{\delta }_{{{\boldsymbol{k}}}_{3}}\rangle \equiv {\delta }_{D}({{\boldsymbol{k}}}_{123})B({{\boldsymbol{k}}}_{1},{{\boldsymbol{k}}}_{2},{{\boldsymbol{k}}}_{3}),\end{eqnarray} \tag{ 4 }$$

where δ( k ) is the overdensity in the Fourier space and k ₁₂₃ ≡ k ₁ + k ₂ + k ₃.

Differently to the power spectrum, the bispectrum quantifies the correlation between triplets of modes in closed triangles. For Gaussian density fields, this quantity is zero, and therefore, its amplitude and shape capture information about the non-Gaussianities in a given field. In the top-right panel of Figure 3, we show the bispectrum for k₁ = 0.15 hMpc⁻¹ and k₂ = 0.25 hMpc⁻¹ as a function of the angle between k₁ and k₂, θ. On this scale, we cannot see large differences, besides the fact that the Zel'dovich approximation underestimates the amplitude of the bispectrum, as expected. In the bottom-right panel of Figure 3, we show the ratio between the different bispectra to the bispectrum of the N-body simulation. We find that both neural networks give very accurate results, although NECOLA is slightly more accurate. The Appendix also shows a comparison of the results from various models on a single realization in the testing set.

The above values of k₁ and k₂ are chosen in order to probe the nonlinear scales of the universe at which the non-Gaussian signatures in the mass distribution (induced by nonlinear gravitational instability) are imprinted. The model has been evaluated at other values of k₁ and k₂ as well, and performs equally well.