Field-level Neural Network Emulator for Cosmological N-body Simulations

For each simulation in the testing data set, we estimate the Eulerian density field using a cloud-in-cell (CIC) particle mesh assignment on a grid with 1024³ voxels. At the voxel centered on position x , we sum the CIC particle weights to obtain the effective number of particles n( x ) and then determine the overdensity ${\delta }_{{\rm{m}}}({\boldsymbol{x}})=n({\boldsymbol{x}})/\bar{n}-1$ with respect to the mean number of particles per cell $\bar{n}$. We Fourier transform the density fields using FFTW3 (Frigo & Johnson 2005) and deconvolve the CIC window function, yielding the density modes δ_m( k ). We then compute the density power spectrum P_mm(k), given by

$$\begin{eqnarray}&&\langle {\delta }_{{\rm{m}}}({\boldsymbol{k}}){\delta }_{{\rm{m}}}({\boldsymbol{k}}^{\prime} )\rangle ={\left(2\pi \right)}^{3}{\delta }_{{\rm{D}}}^{3}({\boldsymbol{k}}+{\boldsymbol{k}}^{\prime} ){P}_{\mathrm{mm}}(k).\end{eqnarray} \tag{ 4 }$$

Here, ${\delta }_{{\rm{D}}}^{(3)}({\boldsymbol{k}})$ denotes the 3D Dirac delta function, which enforces the homogeneity of the density statistics, while the fact that the power spectrum depends only on the magnitude k ≡ ∣ k ∣ is required by isotropy. In practice, we compute the power spectrum as the mean squared mode amplitude in bins the width of the Nyquist wavenumber of the Fourier mesh.

Let δ_m,pred be the predicted density field constructed from either the emulator, COLA, or the fiducial model, and let δ_m,true be the density field constructed from the N-body simulations. We characterize the errors in the model predictions using the cross-correlation coefficient:

$$\begin{eqnarray}&&r(k)=\displaystyle \frac{\langle {\delta }_{{\rm{m}},\mathrm{pred}}({\boldsymbol{k}}){\delta }_{{\rm{m}},\mathrm{true}}({\boldsymbol{k}}^{\prime} )\rangle }{\sqrt{\langle {\delta }_{{\rm{m}},\mathrm{pred}}({\boldsymbol{k}}){\delta }_{{\rm{m}},\mathrm{pred}}({\boldsymbol{k}}^{\prime} )\rangle \langle {\delta }_{{\rm{m}},\mathrm{true}}({\boldsymbol{k}}){\delta }_{{\rm{m}},\mathrm{true}}({\boldsymbol{k}}^{\prime} )\rangle }}.\end{eqnarray} \tag{ 5 }$$

The stochasticity is defined as 1 − r(k)², which quantifies the excess fraction of correlation in the prediction that cannot be accounted for in the target simulation data. We also characterize the errors using the fractional difference between the predicted and true transfer functions:

$$\begin{eqnarray}&&\displaystyle \frac{{T}_{\mathrm{pred}}(k)}{{T}_{\mathrm{true}}(k)}-1=\sqrt{\displaystyle \frac{{P}_{\mathrm{mm},\mathrm{pred}}(k)}{{P}_{\mathrm{mm},\mathrm{true}}(k)}}-1.\end{eqnarray} \tag{ 6 }$$

The results for these density power spectrum errors are shown in the top two rows of Figure 1. The color of each curve indicates the value of σ₈ from the corresponding simulation. We use σ₈ rather than the style parameter Ω_m to label the curves because the small-scale errors are determined by the abundances of high-density structures; σ₈ parameterizes the amplitude of small-scale power, so a higher σ₈ generally indicates a greater abundance of collapsed objects (i.e., dark matter halos).

Zoom In Zoom Out Reset image size

The SNN is considerably more accurate than COLA at scales out to k ∼ 1 Mpc⁻¹ h. The emulator achieves percent-level accuracy in the transfer functions for all cosmologies. The emulator also achieves percent-level stochasticities for all cosmologies except those with extremely high σ₈. As we will continue to see in the results presented later, it is the virialized motion inside of collapsed regions that limits the emulator's accuracy, and this explains the elevated errors in the matter power spectrum for high σ₈ simulations. Due to the fast orbits and fairly chaotic behavior of particles inside a virialized halo, if we output the simulation snapshots at slightly earlier or later times, the particle positions and velocities can change significantly. Since the CNNs map the ZA inputs to nonlinear fields at the particle level, it is very difficult for the CNNs to learn general rules about the fates of these particles. For COLA, the accuracy is limited by the combination of LPT and large time steps. Even with these limitations, in the worst cases, the emulator achieves stochasticities of less than 2% at k = 1 Mpc⁻¹ h, which is comparable with COLA's accuracy. We also see that the FNN errors have a much stronger cosmology dependence, so the SNN is using its style parameter to effectively interpolate between different cosmological models and make more accurate predictions.

The emulator performs worst on cosmologies with a combination of high σ₈ and low Ω_m. The large-scale power in these cosmologies has an extremely high amplitude. This results in a large divergence in the linear displacement field on large scales, so particles generally flow out further from their initial positions. In these extreme cosmologies, the nonlinear displacements are more nonlocal, sourced by a wider region of the density field. This poses a different limitation on the SNN, besides small-scale resolution, since it has a limited field of view for the environment surrounding each particle. In principle, this limitation can be alleviated by adding more convolutional layers, widening the model's field of view. However, in practice, this would increase the evaluation time and require more memory for the large number of additional parameters. We initially implemented the neural network with only three layers, and in this case, the sensitivity of the stochasticities to low values of Ω_m was even more dramatic, and the worst cases were worse than COLA. Increasing to four layers improved these worst cases without too severe a cost in memory and computational time.

From Equation (1), the N-body displacement field is defined with respect to the initial Lagrangian grid, so these can be Fourier transformed directly. The displacement power spectrum, P_{Ψ Ψ} (k), is then given by

$$\begin{eqnarray}&&\langle {\boldsymbol{\Psi }}({\boldsymbol{k}})\cdot {\boldsymbol{\Psi }}({\boldsymbol{k}}^{\prime} )\rangle ={\left(2\pi \right)}^{3}{\delta }_{{\rm{D}}}^{(3)}({\boldsymbol{k}}+{\boldsymbol{k}}^{\prime} ){P}_{{\boldsymbol{\Psi }}{\boldsymbol{\Psi }}}(k),\end{eqnarray} \tag{ 7 }$$

where Ψ( k ) are the Fourier modes of the displacement field, and the dot indicates taking the Cartesian inner product between the displacement mode vectors. The displacement stochasticity and transfer function errors are defined analogously to those of the density field in Equations (5) and (6). These errors are plotted in the middle two rows of Figure 1.

Again, we find that the SNN achieves smaller errors than COLA and has fewer cosmology-dependent errors than the FNN. Similar to the density power spectrum results, the worst cases for the emulator are the high σ₈ cosmology with low Ω_m. The SNN displacement stochasticities are less than 3% at k = 1 Mpc⁻¹ h, and the transfer function errors are negligible (<0.1%) down to this scale.

Note that the small bump in the FNN errors at k ∼ 0.04 Mpc⁻¹ h⁻¹ corresponds to the baryonic acoustic oscillations (BAOs) in the power spectrum. The shape of these oscillations depends on the ratio Ω_b/Ω_m. Since the FNN only encountered one example of this ratio, it is unable preserve the shape of the BAOs in cosmologies with parameters that are far from the fiducial values of its training data. A similar bump in the FNN errors can also be seen in the results for the density power spectrum.

The previous two power spectra depended on the displacement field alone. To compare the accuracy of the velocity part of our model, we construct the momentum power spectrum. We estimate the momentum field, p ( x ), by distributing the particles to a 1024³ mesh using the CIC assignment scheme. Note that this is an Eulerian momentum field, not the Lagrangian momentum field in Equation (3), which was defined with respect to the initial particle grid. We Fourier transform the momentum field and deconvolve the CIC window function to obtain the momentum modes p ( k ). We then compute the momentum power spectrum P _{p p} (k),

$$\begin{eqnarray}&&\langle {\boldsymbol{p}}({\boldsymbol{k}})\cdot {\boldsymbol{p}}({\boldsymbol{k}}^{\prime} )\rangle ={\left(2\pi \right)}^{3}{\delta }_{{\rm{D}}}^{(3)}({\boldsymbol{k}}+{\boldsymbol{k}}^{\prime} ){P}_{{\boldsymbol{p}}{\boldsymbol{p}}}(k),\end{eqnarray} \tag{ 8 }$$

using the same method that was used for the density power spectrum. The errors are plotted in the bottom two rows of Figure 1.

The SNN momentum transfer functions are much more accurate on small scales than those from COLA and show much less cosmology dependence than the FNN transfer function errors. The same is true for the stochasticities for all but the most extreme high-σ₈ low-Ω_m cosmology, where the stochasticities are slightly larger, although comparable with COLA. The emulator has stochasticities of less than 2% at k = 1 Mpc⁻¹ h for most cosmologies. The three worst cases have stochasticities of 4%–5% at this scale. The SNN transfer functions at k = 1 Mpc⁻¹ h are biased slightly too high by a few percent but remain considerably more accurate than COLA.

In general, the SNN momentum power spectrum errors are worse than those from the density or displacement power spectra. This is due to the fact that constructing the momentum field relies on the particle displacements, which are taken from the output of the displacement SNN. Thus, the errors from the displacement field propagate to the momentum field.

In addition to the two-point statistics, we also considered the three-point statistics of the Eulerian matter density field. The matter bispectrum B(k₁, k₂, k₃) is defined as

$$\begin{eqnarray}\begin{array}{rcl}\langle {\delta }_{{\rm{m}}}({{\boldsymbol{k}}}_{1}){\delta }_{{\rm{m}}}({{\boldsymbol{k}}}_{2}){\delta }_{{\rm{m}}}({{\boldsymbol{k}}}_{3})\rangle & = & {\left(2\pi \right)}^{3}{\delta }_{{\rm{D}}}^{(3)}({{\boldsymbol{k}}}_{1}+{{\boldsymbol{k}}}_{2}+{{\boldsymbol{k}}}_{3})\\ & & \times \ B({k}_{1},{k}_{2},{k}_{3}).\end{array}\end{eqnarray} \tag{ 9 }$$

Since the Dirac delta function requires the three wavevectors to form a closed triangle, the bispectrum can also be expressed as a function of the magnitude of two wavevectors and the angle between them, B(k₁, k₂, θ). We use both notations interchangeably. It is also convenient to define the reduced bispectrum,

$$\begin{eqnarray}&&Q({k}_{1},{k}_{2},{k}_{3})\equiv \displaystyle \frac{B({k}_{1},{k}_{2},{k}_{3})}{{P}_{1}{P}_{2}+{P}_{2}{P}_{3}+{P}_{3}{P}_{1}},\end{eqnarray} \tag{ 10 }$$

where P_i = P_mm(k_i ) is the density power spectrum evaluated at k_i . We use the Pylians3 ¹³ library to compute the bispectrum.

We characterize the bispectrum errors by taking the fractional difference with respect to the N-body bispectrum. These are shown in Figure 2 for two different configurations. The two rows of plots show the bispectra in the equilateral configuration as a function of the magnitude of the three wavevectors. The bottom two rows show the bispectra fixing k₁ = 0.1 and k₂ = 1 Mpc⁻¹ h as a function of the angle between them. The top row of each set of plots shows the reduced bispectrum errors, while the bottom row shows the error in the bispectrum. The color of each bispectrum curve corresponds to its value of Ω_m.

Zoom In Zoom Out Reset image size

Unlike the power spectrum errors, the SNN bispectrum errors do not exhibit a strong dependence on Ω_m or σ₈. The SNN bispectrum errors are significantly better than those from COLA or the FNN. This demonstrates that the model is inferring general physical principles about nonlinear gravitational clustering and its cosmology dependence, enabling it to accurately model structure formation on small scales. The SNN errors are less than 5% at k = 0.5 Mpc⁻¹ h and do not increase significantly until after k = 1 Mpc⁻¹ h. The errors in both the bispectra and reduced bispectra are comparable for the SNN because its power spectrum errors are small on these scales.

The matter density field is not directly measured in large-scale structure surveys. Instead, biased tracers of the underlying density field, such as galaxies, are used to infer the statistics and initial conditions of the density field. Under gravitational clustering, the dark matter forms high-density, virialized structures called halos, which are the environments in which galaxies form and reside.

We identify halos in the N-body, CNN, and COLA outputs using the halo-finder code Rockstar (Behroozi et al. 2013). Rockstar identifies halos using a friend-of-friends algorithm and then determines the halo properties by spherical overdensity calculations out to each halo's virial radius. Particles that are not bound to a halo's center of mass are omitted, so the six-dimensional phase space of the N-body particle data is utilized by this algorithm. Rockstar provides us with an excellent framework for further assessing the combined full phase-space predictions of our emulator.

One important observable related to dark matter halos is their density profiles. These density profiles are the result of extremely nonlinear processes of gravitational collapse, mergers, and accretion. To construct halo density profiles from our data, we select a sample of 500 halos with an average mass of 5 × 10¹⁴ M_⊙ h⁻¹. These were taken from a single simulation in the testing data set with cosmological parameters closest to the fiducial ΛCDM cosmology. The simulation has Ω_m = 0.3223 and σ₈ = 0.8311, while the fiducial cosmology has Ω_m = 0.3175 and σ₈ = 0.834. We construct the density field from the positions of the N-body particles using a CIC mesh assignment on a grid of 0.5 Mpc h⁻¹ in cell length. We average the densities of each halo in the z-direction over a 3 Mpc h⁻¹ slice centered on each halo's center of mass. The results are shown in the top two rows of the left panel of Figure 3.

Zoom In Zoom Out Reset image size

The SNN achieves much more accurate halo interiors than COLA and also shows improvements over COLA in the environments surrounding halos. The SNN even shows small improvements over the FNN, even though these halos form in a cosmology close to the FNN's training data. This demonstrates that even small changes in Ω_m can result in a loss of accuracy for the FNN in the deeply nonlinear regime.

Another important observable related to halos is their number density as a function of mass, n_h (M), which is known as the halo mass function. We construct the halo mass function for each testing cosmology and compute its error as the fractional difference with respect to the N-body results. These errors are shown in the top row in the right panel of Figure 3.

The SNN shows a dramatic improvement over COLA, and again we find a reduced cosmology dependence in the errors compared with the FNN. The abundances of low-mass halos (M < 5 × 10¹³ M_⊙ h⁻¹) from the SNN are biased high by about 10%–20% but scattered around the N-body mass functions at higher mass. To be fair to COLA, its purpose is not to accurately recover the small-scale details of halos, and its systematically low halo abundances are well known. Abundance matching methods have been developed to identify dense structures in COLA output that correspond to halos from the full N-body evolution (Izard et al. 2016).

In large-scale structure surveys, the distance to a galaxy is inferred through its observed redshift. This redshift is not due to cosmic expansion alone but gets an additional Doppler contribution from the galaxy's peculiar velocity. Even with the correct cosmological model, the observed positions of galaxies are radially displaced compared to their true positions. The shapes of galaxies are also distorted along the line of sight. These effects are known as RSDs. Matter inside of halos tends to have large, virialized velocities, which results in a smearing of the halo density along the line of sight, known as the "Fingers of God" phenomenon. On larger scales, where matter falls toward the local minimum of the gravitational potential, spherical profiles are squashed along the line of sight. It is crucial to accurately describe the late-time, nonlinear density field in redshift space in order to compare with realistic survey data.

We use the Pylians3 library to apply RSD to the N-body particle data and then construct the density profiles for halos in redshift space. The stacked redshift space profiles of the same 500 halos as in the previous subsection are shown in the bottom two rows of the left panel of Figure 3, along with their residuals with respect to the N-body simulations. All of the models exhibit the larger-scale flattening effect due to coherent infall, which appears as the two bulges on either side of the halo. However, only the SNN accurately reproduces the scale of the Fingers of God smearing along the y-axis compared with the N-body simulations.

As a final test of our model's accuracy in the nonlinear regime, we consider the multipole moments of the matter power spectrum in redshift space. Since RSDs contribute only radial distortions, they select the radial direction as special and induce anisotropy in the power spectrum. This anisotropy can be characterized by expanding the redshift space power in multipole moments:

$$\begin{eqnarray}&&{P}_{\mathrm{mm}}({\boldsymbol{k}})=\displaystyle \sum _{{\ell }}(2{\ell }+1){P}_{{\ell }}(k){{ \mathcal L }}_{l}\left(\cos (\theta )\right).\end{eqnarray} \tag{ 11 }$$

Here, ${{ \mathcal L }}_{{\ell }}$ are the Legendre polynomials, and θ is the angle between the wavevector k and the radial line-of-sight direction. We used the Pylians library to determine the monopole (ℓ = 0) and quadrupole (ℓ = 2) moments of the redshift space density power spectrum. The fractional errors with respect to the N-body results are shown in the bottom two rows of the right panel of Figure 3.

The SNN reproduces the monopole and quadrupole power spectra more accurately than COLA on small scales. Once again, we see that the SNN errors have less cosmology dependence than the FNN errors, indicating that the full phase-space evolution is modeled better across a wide range of Ω_m values due to its inclusion as a style parameter. Note that since RSDs tend to move small-scale power to larger scales, the errors are larger at the 1 Mpc⁻¹ h scale compared with the power spectrum errors in real space.

Searching for observational evidence of PNG is one of the highest-priority objectives of contemporary large-scale structure surveys. If detected, PNG would provide crucial evidence for inflation and possibly rule out a wide range of early universe models (Creminelli & Zaldarriaga 2004; Meerburg et al. 2019). The simulations that our emulator was trained on have purely Gaussian initial conditions. Non-Gaussianity can be added to the initial conditions by first drawing the modes of a random Gaussian scalar potential Φ_g( k ) and then constructing the non-Gaussian potential (Scoccimarro et al. 2012):

$$\begin{eqnarray}&&\begin{array}{l}{{\rm{\Phi }}}_{\mathrm{ng}}({\boldsymbol{k}})={{\rm{\Phi }}}_{{\rm{g}}}({\boldsymbol{k}})+{f}_{\mathrm{NL}}\int \displaystyle \frac{{{\rm{d}}}^{3}k^{\prime} }{{\left(2\pi \right)}^{3}}K({\boldsymbol{k}}^{\prime} ,{\boldsymbol{k}}-{\boldsymbol{k}}^{\prime} )\\ \,\times \,{{\rm{\Phi }}}_{{\rm{g}}}({\boldsymbol{k}}^{\prime} ){{\rm{\Phi }}}_{{\rm{g}}}({\boldsymbol{k}}-{\boldsymbol{k}}^{\prime} ).\end{array}\end{eqnarray} \tag{ 12 }$$

Here, f_NL sets the amplitude of the primordial bispectrum, and K( k ₁, k ₂) is an integration kernel specifying the shape of the primordial bispectrum. There are several typical choices for this kernel, usually referred to as bispectrum templates (Komatsu et al. 2003; Babich et al. 2004; Senatore et al. 2010). These include the local (LC), orthogonal (OR), and equilateral (EQ) templates. Since PNG enters only through the initial conditions of a simulation, our model should immediately generalize to these cosmologies without the need for additional parameters or training.

To test this, we use the new QUIJOTE-PNG (Coulton et al. 2023) N-body simulations. These augment the previous QUIJOTE data set with simulations that include the PNG templates mentioned above in their initial conditions. We use a set of seven simulations, one with Gaussian (G) initial conditions and two from each of the bispectrum template types listed earlier, setting f_NL = ±100. All other cosmological parameters are the same as in the fiducial cosmology, and the simulation configurations are identical to those of the training data. We compute the ZA displacements and velocities at redshift zero for all of these simulations and predict the nonlinear configurations using our emulator. We then compare these with the full N-body simulations. The results are shown in Figure 4.

Zoom In Zoom Out Reset image size

The top left panel shows the standard deviation of the displacement errors for the Gaussian and non-Gaussian simulations. There is no significant or systematic increase in error in the simulations with PNG; some even have smaller errors than the Gaussian case. The bottom left panel shows the standard deviation of the velocity errors, again without any systematic increase from the presence of PNG in the initial conditions. The right panels show the probability density functions for the errors in the displacement (top) and velocity (bottom) components for all seven simulations. We do not include a legend to label the different template types, since their error distributions are indistinguishable. This indicates that the emulator generalizes to these PNG simulations as expected.

Zoom In Zoom Out Reset image size

These results are unsurprising, since the model has been exposed to an enormous diversity of density environments during training. The presence of PNG alters the statistics of these configurations, but the physics of nonlinear gravitational clustering in these environments remains the same. It is this underlying physics of late-time structure formation that the emulator infers from its training data. Further tests are still required to determine how well the emulator preserves the observational features of PNG in various summary statistics, including the power spectrum and bispectrum, as well as at the field level. We leave this to future work.

Current cosmological N-body simulations model structure formation in the universe by evolving a large number (10⁶– 10¹²) of massive dark matter particles that interact with each other only via gravity. Since our universe begins with an almost uniform and Gaussian initial condition, the simulations start with particles only slightly perturbed from a uniform configuration, here a grid. For each particle, this initial small displacement Ψ from its grid location q can be computed accurately with the LPT. The particle position is then x = q + Ψ( q ). The leading-order LPT (ZA) predicts linear particle motion Ψ_ZA( q , z) = D(z)/D(z₀)Ψ_ZA( q , z₀), where the growth function D is a function of redshift z. As a result, the ZA velocity is linearly related to the displacement v _ZA( q , z) = aH(z)f(z)Ψ_ZA, where a = 1/(1 + z) is the expansion scale factor, $H={\rm{d}}\,\mathrm{log}\,a/{\rm{d}}\,t$ is the Hubble expansion rate, and $f={\rm{d}}\,\mathrm{log}\,D/{\rm{d}}\,\mathrm{log}\,a$ is known as the growth rate.

Structures collapse under gravity and cluster hierarchically in the universe. On small scales, particle motion becomes nonperturbative, invalidating ZA and higher-order LPT. This process becomes so nonlinear that it can only be accurately modeled by expensive N-body simulations via time integration of thousands of time steps. On the other hand, on large scales, the universe remains relatively uniform, and perturbation theories still apply. This motivates us to build a model to emulate the small-scale structure formation in the N-body simulations while retaining the ZA predictions on large scales. The CNNs integrate well into this framework, since they excel at local textures but lack a global view beyond their finite receptive fields. Therefore, our goal is to train CNN models to nonlinearize the linear ZA inputs to make predictions comparable in accuracy to N-body simulations.

For benchmarks, we compare the neural network predictions to alternative fast simulation models. Fast approximate simulations (Tassev et al. 2013; Feng et al. 2016) save computation time by only integrating tens of time steps and thus are less accurate than the full N-body simulations. Here, we compare the neural network model to the popular method COLA (Tassev et al. 2013), as implemented in L-PICOLA (Howlett et al. 2015). COLA solves the particle motions relative to the second-order LPT trajectory. We run L-PICOLA with the same settings as the full N-body simulations but for only 10 time steps.

As explained above, we want the neural network to keep the ZA predictions (Ψ_ZA and v _ZA at the target redshift) on large scales, where perturbation theories are valid and accurate. This accounts for the nonlocal gravitational interactions beyond the local views of CNNs. The CNN models can then take the linear inputs and form the appropriate small-scale textures. This is most easily implemented by a global residual operation, in which the CNN predicts the difference between the N-body and ZA motions.

The design with global residual connection has other advantages too. One is that the Galilean symmetry is approximately preserved by effective data augmentations. This is because even though the ZA and N-body particle motions may be affected by large-scale shifts in displacements or velocities, e.g., by adding a global constant vector, their difference is not. The other advantage is that the large-scale dependence on cosmological parameters is automatically accounted for by the ZA inputs, so the neural network emulators only need to model the small-scale cosmology modulation, as explained below. Therefore, we train two separate networks for displacement and velocity, respectively, whose inputs are linearly related.

By construction, CNNs should preserve the translational symmetry if they are fully convolutional. However, this is broken in most cases by the nonperiodic paddings commonly adopted in computer vision tasks. Our models fully preserve the translational symmetry using the periodic boundary condition of the N-body simulation volume. In addition, they approximately preserve the rotational symmetry of the cubic geometry by data augmentations that rotate and reflect the input and output fields as in He et al. (2019).

Before applying CNN models to our problem, we need to prepare the data in an image-like format, which is straightforward thanks to the uniform initial condition of the universe. Both the ZA inputs and N-body target are functions of the Lagrangian or initial positions, q , which form a uniform grid. Therefore, the displacements form a 3D image, with three channels being the Cartesian components of the particle displacements. Likewise, we format the velocity fields in a similar way.

For both displacement and velocity, we adopt a simple U-Net/V-Net-type architecture (Ronneberger et al. 2015; Milletari et al. 2016) similar to that in Alves de Oliveira et al. (2020). A diagram depicting the network architecture is shown in Figure 5. It works on four levels of resolution connected in a V shape first by three downsampling layers and then by three upsampling layers by stride-2 2³ convolutions and stride-1/2 2³ transposed convolutions, respectively. Blocks of two 3³ convolutions connect the input, resampling, and output layers. Similar to V-Net, a residual connection (ResNet; He et al. 2016), here a 1³ convolution instead of identity, is added over each block. We use a batch normalization layer after every convolution except the first one and the last two and a leaky ReLU activation (with negative slope 0.01) after each batch normalization, as well as the first and second-to-last convolution layers. As in the original ResNet architecture, the second or last activation in each residual block applies after the addition. And as in U-Net, the inputs to each of the downsampling layers are concatenated to the outputs of the upsampling layers of the same resolution at the top three resolution levels. All layers have 64 channels, except the input and the output (3) and those after concatenations (128). Finally, an important difference from the original U-Net architecture is that we add the input displacement or velocity fields directly to the output. Therefore, the network is effectively learning the corrections from linear to nonlinear motions, as we have mentioned above.

With the Lagrangian description, the simplest loss functions for training one can formulate aim to minimize the residuals in the displacements. Alves de Oliveira et al. (2020) and Li et al. (2021) showed that combining that with a loss of Eulerian quantities can greatly improve the performance of the trained model in the Eulerian space. For the displacement model, we can compute the particle positions and thus their (Eulerian) density distribution, described by the overdensity $\delta ({\boldsymbol{x}})\equiv n({\boldsymbol{x}})/\bar{n}-1$, where n( x ) is the particle number in voxel x , and $\bar{n}$ is its mean value. Specifically, we use the CIC, or trilinear scheme, to assign particles to the voxels. The total loss for our displacement network is $\mathrm{log}{L}_{\delta }+\lambda \mathrm{log}{L}_{{\boldsymbol{\Psi }}}$, with L_δ and L_Ψ being the mean squared error (MSE) losses on n( x ) and Ψ( q ), respectively. Combining the two losses with a logarithm allows us to ignore their absolute magnitudes and trade between the relative changes. The parameter λ is a weight on this trade-off, and we find that setting λ = 3 for the displacement network achieves faster training. For the velocity network, we use a loss of $\mathrm{log}{L}_{{\boldsymbol{v}}({\bf{q}})}+\mathrm{log}{L}_{{\boldsymbol{v}}({\boldsymbol{x}})}+\mathrm{log}{L}_{{{\boldsymbol{v}}}^{2}({\boldsymbol{x}})}$. The first term is the MSE loss of the particle velocities (in Lagrangian space, thus the q argument). The second and last terms are defined by

$$\begin{eqnarray}&&{L}_{{{\boldsymbol{v}}}^{n}({\boldsymbol{x}})}=\int m{\parallel {\rm{\Delta }}\underset{i=1}{\overset{n}{\displaystyle \bigotimes }}{\boldsymbol{v}}\parallel }_{2}^{2}\,{\rm{d}}{\boldsymbol{x}},\end{eqnarray} \tag{ A1 }$$

where m is the particle mass, ⨂ forms tensor products of the particle velocity, and Δ takes the residual from predicted to N-body results. For n = 1, this is simply the loss on the momentum fields, and for n = 2, the loss minimizes the residual of the mass-weighted velocity dispersions. We find the n = 2 loss crucial for correctly predicting the RSD, especially the Fingers of God effect, as shown in Figure 3.

The most important feature of our models is the ability to incorporate the dependence on cosmological parameters. The late-time structure evolution is sensitive to the expansion history of the universe, i.e., the expansion rate H₀ = 100 h km s⁻¹ Mpc⁻¹ and the matter density parameter Ω_m. The former has been accounted for by using proper units that include h; thus, we only need to include Ω_m as an input to our networks. We achieve this following StyleGAN2 (Karras et al. 2019), in which style parameters are incorporated by weight modulations and demodulations on the convolution kernels. In the modulation phase, the style parameters (here Ω_m) are first affine transformed to the same dimension as the number of input convolution channels and then multiplied by the kernel of each channel. In the demodulation phase, the convolution kernels are renormalized by their fan-ins, so that propagating data do not explode or vanish. We refer readers to Equations (1)–(3) in Karras et al. (2019) for more details.

Limited by the size of the GPU memory, an entire simulation field (512³) cannot be fed to the network at once and has to be cropped into smaller chunks of size 128³. To preserve the translational equivariance, we keep no paddings in the convolutions, resulting in smaller network outputs than their inputs, and we compensate for this by periodically padding 48 voxels per side to the input fields. We train with a batch size of 16 distributed on 16 Nvidia V100 GPUs and use the Adam optimizer (Kingma & Ba 2014) with learning rate 0.0004 and hyperparameters β₁ = 0.9 and β₂ = 0.999. We reduce the learning to 10% when the loss does not improve for three epochs. We trained both models until the loss converged, which took 80 epochs for the displacement model and 60 epochs for the velocity model. The training ran for 2 × 10⁴ GPU hr, roughly equally for each model.

Finally, it is worth mentioning that the total data size amounts to 12 TB. This poses technical challenges because training the networks requires repeatedly loading the whole data set to the GPU nodes. If the data sampling and loading are not carefully designed, the GPUs can easily outrun the cluster IO and be starved for data. We solve this and implement efficient training in map2map.