Learning an Effective Evolution Equation for Particle-Mesh Simulations Across Cosmologies

Particle-mesh simulations trade small-scale accuracy for speed compared to traditional, computationally expensive N-body codes in cosmological simulations. In this work, we show how a learning-based model could be used to learn an effective evolution equation for the particles, by correcting the errors of the particle-mesh potential incurred on small scales during simulations. We find that our learnt correction yields evolution equations that generalize well to new, unseen initial conditions and cosmologies. We further demonstrate that the resulting corrected maps can be used in a simulation-based inference framework to yield an unbiased inference of cosmological parameters. The model, a network implemented in Fourier space, is exclusively trained on the particle positions and velocities. This work is of particular importance in the context where, in the coming decade, cosmology will be transformed by unprecedented volumes of survey data from multi-billion-dollar instruments, and extracting all the information from these datasets will require fast and accurate cosmological simulators which are not yet available.