JERALD: high-fidelity dark matter, stellar mass and neutral hydrogen maps from fast N-body simulations

We present a new code and approach, JERALD -- JAX Enhanced Resolution Approximate Lagrangian Dynamics -- , that improves on and extends the Lagrangian Deep Learning method of Dai & Seljak (2021), producing high-resolution dark matter, stellar mass and neutral hydrogen maps from lower-resolution approximate $N$-body simulations. The model is trained using the Sherwood-Relics simulation suite (for a fixed cosmology), specifically designed for the intergalactic medium and the neutral hydrogen distribution in the cosmic web. The output is tested in the redshift range from $z=5$ to $z=0$ and the generalization properties of the learned mapping is demonstrated. JERALD produces maps with dark matter, stellar and neutral hydrogen power spectra in excellent agreement with full-hydrodynamic simulations with $8\times$ higher resolution, at large and intermediate scales; in particular, JERALD's neutral hydrogen power spectra agree with their higher-resolution full-hydrodynamic counterparts within 90% up to $k\simeq1\,h$Mpc$^{-1}$ and within 70% up to $k\simeq10\,h$Mpc$^{-1}$. JERALD provides a fast, accurate and physically motivated approach that we plan to embed in a statistical inference pipeline, such as Simulation-Based Inference, to constrain dark matter properties from large- to intermediate-scale structure observables.