Differentiable Hydrodynamical Simulations and Solver-in-Loop for Cosmological Inference

Introduction to Differentiable Hydrodynamics

• 💡 Ben Horowitz's work focuses on the intersection of large scale structure, machine learning, and hydrodynamical simulations, particularly for understanding the cosmic web at small scales.
• 🎯 The motivation stems from the PFS galaxy evolution survey, aiming to track galaxy formation and evolution from high to low redshift, especially during the peak star formation period known as "cosmic noon."
• 🔑 The approach emphasizes field level inference, moving beyond traditional surrogate modeling to directly model complex phenomena like the lime alpha forest signal.

Implementing Differentiable Simulations

• 🚀 The core idea is to make hydrodynamical simulations themselves differentiable, enabling the solution of inverse problems and backpropagation through the hydrodynamics.
• 🛠️ This involves coding the Euler equations on a grid using JAX on GPUs, which significantly accelerates computations and allows for end-to-end differentiability.
• 🔍 Differentiable hydro simulations can be used to inverse solve for initial conditions, such as reconstructing the origins of a cluster from noisy mock data.
• ✨ The framework can be coupled to dark matter simulations (JAX-PM), jointly evolving gas density, temperature, and dark matter on a grid to form large scale structures.

Integrating Subgrid Physics

• 🌌 The goal is to make subgrid models, like star formation and feedback (AGN/supernova), differentiable to understand their effects on gas thermal states and lime alpha forest signals.
• 🧩 A simple model, inspired by IllustrisTNG, allows for star formation in cold, dense gas and subsequent energy injection from stars into the gas, made differentiable using a "scumble softmax trick."
• ⚠️ While effective, hydrodynamical simulations remain computationally expensive for large volumes, especially when incorporating advanced techniques like Adaptive Mesh Refinement (AMR).

The Solver-in-Loop Approach

• ⚡ To address computational cost, the solver-in-loop approach maps low-resolution to high-resolution simulations by applying corrections at every time step throughout the evolution.
• 🧠 This involves training a fully convolutional correction network that learns to push the low-resolution simulation trajectory towards the high-resolution benchmark.
• ✅ This method can capture qualitatively different physical phenomena (e.g., double swirl patterns in Kelvin-Helmholtz instability) at a negligible additional computational cost.

Astrophysical Applications and Future Work

• 🌱 The solver-in-loop method is being applied to cosmological astrophysical systems, such as modeling star formation physics and stellar mass field evolution using particle mesh simulations.
• 📈 Future work aims to train these models directly on observational data (e.g., light cones of galaxies) to jointly infer initial conditions and galaxy formation physics.
• 💡 Differentiable models and solver-in-loop approaches offer a novel way to accelerate traditional simulations and generate robust hydrodynamical simulations more efficiently, even for those skeptical of explicit field-level inference.