astro-ph.IMJul 29, 2026

Emulating Cosmic Structure Formation with a Lagrangian Neural Cellular Automaton

Authors: Cooper Jacobus, Beatriz Tucci, Oliver Philcox

Organizations: Department of Physics, Stanford University, Stanford, CA 94305, USA · Kavli Institute for Particle Astrophysics and Cosmology, Stanford · Leinweber Institute for Theoretical Physics, Stanford

Abstract

Field-level inference of cosmological initial conditions from galaxy surveys requires a forward model that is simultaneously accurate in the non-linear regime, computationally efficient, and fully differentiable. Traditional N-body simulations are accurate but computationally prohibitive for iterative inference, while approximate solvers like Lagrangian Perturbation Theory (LPT) fail to capture the knotty halo-forming dynamics of the cosmic web at late times. We introduce the \textit{Lagrangian Neural Cellular Automaton} (LNCA), a hybrid deep learning framework that can be applied to emulate structure formation as a local, iterative dynamical process on a comoving lattice. Unlike standard Eulerian Convolutional Neural Networks (CNNs) which map fixed density fields, the LNCA operates in the Lagrangian frame, advecting the computational graph itself to follow the flow of mass. By training the network to learn only the \textit{residual} displacement corrections to the Zeldovich approximation, we achieve high-fidelity emulation of the non-linear physics while guaranteeing accuracy at large scales. We further constrain our model to produce complete trajectories, not just final states, by adopting an equivariant cellular automaton architecture, which recurrently iterates on its internal states to yield a dynamic history. The resulting model is strictly local, translationally and rotationally equivariant, and naturally supports continuous time integration, making it a reliable differentiable forward model for reconstructing the initial conditions of the universe from lightcone data. Our trained model supports percent-level precision in the power and cross spectra well into the non-linear regime (k≲0.5 hMpc−1k \lesssim 0.5 \, h \text{Mpc}^{-1}), while requiring ∼104\sim10^4 times fewer learned parameters than comparable models which take the form of an interpretable internal dynamic rule set.

Explore similar work

CardsList
Jun 8, 2026astro-ph.CO

Learning the Universe: Posterior Reliability of Neural Generative Models in High-Dimensional Field-Level Inference of Cosmic Initial Conditions

Accurate posterior estimation is central to scientific inference, as uncertainties determine what can be reliably learned from observational data. While Markov chain Monte Carlo methods provide asymptotic convergence guarantees, they are computationally demanding in high-dimensional settings. Neural network-based generative models for entire discretized 3D fields enable fast amortized inference but often lack convergence guarantees and principled accuracy assessment. Using Hamiltonian Monte Carlo to obtain reference posterior samples, we conduct a controlled field-level evaluation of an implicit generative model (Stochastic Interpolants) and an explicit likelihood-based model (GLOW normalizing flows). This comparison, unavailable in typical applications, enables the detection of posterior geometry failures that standard metrics cannot capture. As a case study, we consider the cosmological inverse problem of inferring cosmic initial conditions from present-day large-scale structure. To match the precision of modern cosmological data, this problem increasingly relies on complex, non-linear, and non-differentiable simulators, which are incompatible with gradient-based inference frameworks. Generative models offer a route to address these challenges, provided their inferred posteriors are reliable. In this work, we show that matching posterior means, marginal distributions, or achieving high cross-correlation does not imply correct uncertainty structure, as revealed by posterior variance fields and sample-based evaluations. Through this work, we aim to raise awareness of the challenges of uncertainty estimation in high-dimensional field-level settings, highlighting the importance of careful design and validation of neural generative approaches for scientific applications.
Ludvig Doeser, Jens Jasche
CosmologyPosteriorGenerative ModelsHamiltonian Monte CarloHigh-DimensionalReliability+1
AskPDF
Jun 22, 2026astro-ph.CO

Field-level weak lensing cosmology with 60 simulations using multifidelity simulation-based inference

We perform a realistic KiDS-Legacy mock analysis with field-level neural compression and simulation-based inference using just 60 NN-body simulations. The weak lensing shear field encodes substantially more cosmological information than standard two-point summary statistics such as the power spectrum. Field-level inference can fully exploit this information, but physical realism at the field-level requires very high-fidelity simulations. This poses a major challenge for simulation-based inference (SBI): accurate empirical density modelling and deep-learning-based neural compression require tens of thousands of training samples, but achieving physical realism at the field level makes each simulation extremely costly. We demonstrate that multifidelity SBI can alleviate this tension by substantially reducing the number of high-fidelity simulations needed for accurate cosmological inference. We pre-train neural inference models on realistic KiDS-Legacy-like shear mocks using fast log-normal \texttt{GLASS} simulations and fine-tune them on a small set of high-fidelity NN-body simulations. We show that 6060 high-fidelity simulations are sufficient to obtain informative and well-calibrated cosmological posteriors, enabling at least an order-of-magnitude reduction in simulation cost for accurate field-level inference in a realistic setting.
Alex A. Saoulis, Kiyam Lin, Niall Jeffrey +5
CosmologyMulti-FidelityBayesian Inference ApproachesConvolutional Neural Network Compression
AskPDF
May 27, 2026astro-ph.CO

Dark Quest II: A Wide-Coverage Neural Network Emulator of the Nonlinear Matter Power Spectrum Across Extended Cosmologies

\textsc{DarkEmulator2} is a neural network emulator of the nonlinear matter power spectrum in a nine-dimensional w0waνoCDMw_0 w_a νo \mathrm{CDM} parameter space, developed as the emulator component of the \textsc{Dark Quest II} (DQ2) program. It is trained on simulations generated with the \textsc{Ginkaku} code, whose numerical implementation, accuracy tests, and post-processing pipeline are described in the companion paper. The design follows a unified strategy: in addition to the cosmological parameter vector, we supplement the neural network's inputs with three families of physically motivated auxiliary quantities -- the linear matter power spectrum, descriptors of the simulation resolution, and a low-dimensional summary of the initial Gaussian random field -- that are expected to improve generalization across the parameter space. Training a single network jointly across three simulation resolution tiers allows the emulator to exploit a small number of high-resolution simulations while retaining broad coverage from lower-resolution simulations. For a Lbox=1 \hiGpcL_{\mathrm{box}}=1\,\hiGpc box with N=30003N=3000^{3} particles, the emulator reproduces the simulated matter power spectrum to subpercent accuracy up to the particle Nyquist scale, kNy≃10 \hMpcik_{\mathrm{Ny}}\simeq 10\,\hMpci. The emulator remains accurate over the calibrated wavenumber range, while its highest-kk predictions depend on the simulation resolution and shot noise. We validate the emulator on independent test suites and, through a cross-comparison with several public emulators and widely used fitting formulas, characterize the inter-model consistency and the parameter-dependent trends in their residuals.
Satoshi Tanaka, Takahiro Nishimichi, Yosuke Kobayashi
Neural Network EmulatorCosmology
AskPDF