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}<sup>{-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.