Hybrid Physical-Neural Simulator for Fast Cosmological Hydrodynamics

TL;DR

This work introduces a differentiable forward model for cosmological hydrodynamics by hybridizing a differentiable particle-mesh gravity solver with a neural network that learns an effective gas pressure. Trained in a solver-in-the-loop setup on a single fully hydrodynamical Camels Simba simulation, the model evolves both dark matter and gas in comoving coordinates, with gravity computed on a mesh and gas dynamics augmented by a learned pressure field. Results show improved field-level agreement and comparable or better two-point statistics relative to baselines like Enthalpy Gradient Descent, while achieving strong data efficiency. This approach enables field-level inference of cosmological parameters and initial conditions, with potential to fit directly to observational data such as Sunyaev–Zeldovich signals and weak lensing, albeit with limitations on history dependence which could be addressed via latent-variable neural ODE extensions.

Abstract

Cosmological field-level inference requires differentiable forward models that solve the challenging dynamics of gas and dark matter under hydrodynamics and gravity. We propose a hybrid approach where gravitational forces are computed using a differentiable particle-mesh solver, while the hydrodynamics are parametrized by a neural network that maps local quantities to an effective pressure field. We demonstrate that our method improves upon alternative approaches, such as an Enthalpy Gradient Descent baseline, both at the field and summary-statistic level. The approach is furthermore highly data efficient, with a single reference simulation of cosmological structure formation being sufficient to constrain the neural pressure model. This opens the door for future applications where the model is fit directly to observational data, rather than a training set of simulations.

Figures (5)

• Figure 1: Projected gas density $\rho_\mathrm{gas}$ (logarithmic scale) comparing four methods: reference hydrodynamical Camels simulation, our hybrid simulator with learned gas pressure force, gravity-only JaxPM, and its EGD post-processing. Our pressure model suppresses small-scale structure formation, better matching the hydrodynamical reference than purely gravitational evolution while retaining more detail than EGD.
• Figure 2: Two-point statistics of gas densities $\rho_\mathrm{gas}$ from \ref{['fig:final_fields']}. Solid lines and shaded bands indicate means and standard deviations over random initial conditions (cosmic variance), respectively. Left and middle: Power spectra demonstrate that both our method and EGD suppress small-scale power relative to gravity-only evolution (JaxPM), achieving better agreement with the Camels reference. Right: Our hybrid simulator shows consistently higher cross-correlation with the reference than EGD. These quantitative results confirm the visual intuition from \ref{['fig:final_fields']}.
• Figure D.1: Projections of the four local features $\mathbf{h}(\mathbf{x})$ used to predict the effective pressure $P_\varphi$: gas density, scalar force (see \ref{['eq:fscalar']}), velocity divergence, and velocity dispersion (left to right). For visualization, each field $h$ is separately scaled by $\mathrm{arcsinh}(h/\sigma_h)$.
• Figure D.2: Like \ref{['fig:final_fields']}, but for the dark matter density $\rho_\mathrm{dm}$. At this resolution, differences between methods are minimal, demonstrating that our neural pressure model does not adversely impact dark matter evolution. The large-scale structure resembles that in \ref{['fig:final_fields']} because gas traces the underlying dark matter distribution up to scales where gas feedback effects become significant.
• Figure D.3: Evolution of gas density $\rho_\mathrm{gas}$ from our initial conditions ($a=0.14$) to the present day ($a = 1$). For this example, the penalized scale factors in the loss function \ref{['eq:loss']} are $a_s \in \{0.30, 0.44, 0.65, 1.0\}$. As a post-processing method, EGD does not model a self-consistent evolution and is therefore not shown here.