TUNeS: Neural Emulation of Large-Scale Structure Across Redshifts

Abstract

In this work, we introduce TUNeS (Temporal UNet emulator for Structure formation), a neural network framework for accelerating N-body simulations by predicting the nonlinear evolution of the matter density field from an initial particle distribution. TUNeS employs a two-stage modeling strategy, combining particle-based inference with a density-field refinement on a regular grid, enabling accurate reconstruction of both large- and small-scale structures. The model is designed to operate across redshift, taking particle snapshots at arbitrary input redshifts and predicting density fields at arbitrary target redshifts. In this work, we evaluate its performance using simulations initialized at $z=100$, with predictions generated at multiple lower redshifts. Trained on only eight N-body simulations, TUNeS reproduces reference results with good agreement in both Gaussian and non-Gaussian statistics, including two-point correlations, one-point distributions, peak counts, and three-dimensional Minkowski functionals. In particular, at $k \simeq 1\,h\,\mathrm{Mpc}^{-1}$, the power spectrum error remains at the few-percent level. End-to-end inference from $256^3$ particles to a $256^3$ density grid can be completed in $\sim25\,\mathrm{second}$ on a single GPU. Thanks to its architectural design, the model naturally scales to larger particle numbers and larger volumes through particle batching and window-based refinement.

Figures (9)

• Figure 1: Overview of the pipeline for the proposed framework. The rightmost panels show two-dimensional mean density projections of slices with thickness $64\,h^{-1}\mathrm{Mpc}$ at different stages. For visualization, the particle field is randomly down-sampled to $10^6$ particles.
• Figure 2: Architecture of the Stage-1 network. The model takes as input the initial particle positions $\boldsymbol{x}_{\rm ini}$, velocities $\boldsymbol{v}_{\rm ini}$, and the initial and target redshifts $(z_{\rm ini}, z_{\rm fin})$. Particle positions are encoded using Fourier features and concatenated with velocities, followed by a linear projection to a fixed-width hidden representation. Temporal information is embedded from the logarithmic scale factors and injected into each residual block via feature-wise linear modulation. The residual block, highlighted by a dashed blue frame, is repeated six times and consists of a FiLM-modulated layer normalization, a per-particle multilayer perceptron, and a residual skip connection. Feature dimensions at each stage are annotated along the arrows in the figure, and all residual blocks operate at the same feature dimension. The network outputs a per-particle displacement field $\Delta\boldsymbol{x}$. Here $B$ denotes the batch size and $N$ the number of particles.
• Figure 3: Stage-2 network architecture for density-field refinement. Colored blocks indicate different computational modules as defined in the legend. Strided and transposed 3D convolutions are used to change spatial resolution, while gray arrows denote skip connections that transfer feature maps across scales. Redshift information is embedded from $(\ln a_{\rm ini}, \ln a_{\rm fin})$ and broadcast to all FiLM3D layers, enabling redshift-dependent modulation throughout the network. The final output is obtained via a $1\times1\times1$ convolution producing a single-channel density-field window. The spatial resolution of each level is indicated on the left, and the label on top of each block denotes its number of channels.
• Figure 4: Two-dimensional mean density field of a slice of thickness $64\,h^{-1}\,\mathrm{Mpc}$. The left panel shows the input initial condition at $z=100$. The middle panel shows the reference N-body result at $z=0$, and the right panel shows the corresponding model prediction.
• Figure 5: One-point probability distribution functions (PDFs) of the density field, smoothed with a Gaussian kernel of radius $2\,h^{-1}\,\mathrm{Mpc}$, measured from the model predictions and the reference N-body simulations. The PDFs are averaged over 10 independent test simulations. Colors represent different redshifts, as indicated by the color bar on the right. Upper panel: mean PDFs, where dashed lines represent the reference N-body results and solid lines denote the model predictions. Lower panel: relative errors with a $1\sigma$ error band. The $\pm 2\%$ reference levels are indicated by gray dashed lines, and the zero level by a black dashed line.