Hybrid Physical-Neural ODEs for fast differentiable N-body Simulations

Denise Lanzieri


Supervisors: François Lanusse, Jean-Luc Starck














slides at github.com/dlanzieri/talks/LSST_France_Nov2021

Kiessling et al. (2015)

Outline of this talk





  • Fast realization of complex processes

    • $\Longrightarrow$ Take the effective physics description and combine it with a ML approach


  • Cosmological simulations are based on physical processes

    • $\Longrightarrow$ these impose symmetries and constraints


  • Cosmological differentiable simulations

    • $\Longrightarrow$ JAX and TensorFlow framework




The Particle-Mesh scheme for N-body simulations

The idea: approximate gravitational forces by estimating densities on a grid.
  • The numerical scheme:

    • Estimate the density of particles on a mesh
      $\Longrightarrow$ use the cloud-in-cell (CIC) interpolation scheme

    • Apply a Fourier transform to obtain the over-density field $\delta_k$ in Fourier space.

    • Compute gravitational forces $\Longrightarrow$ related to the density field via a transfer function ($\nabla \nabla^{-2}$)

  • Only a series of FFTs and interpolations.
$\Longrightarrow$ Fast and simple, at the cost of approximating short range interactions.

Fill the gap in the accuracy-speed space

    N-body PM simulation:

    • Fast (limited number of time steps while enforcing the correct linear growth)

    • Cannot give accurate halo matter profiles or matter power spectrum

    The correction idea :mimics the physics that is missing
    $\Longrightarrow$ Halo virialization

Solving the N-body differential equations

  • Time integration from a system of ordinary differential equations (ODE) and treat the ODEsolver as a black box


  • We can use this parametrisation to complement the physical ODE with neural networks.
$$ \left\{ \begin{array}{ll} \frac{d \mathbf{x}}{d a} & = \frac{1}{a^3 E(a)} \mathbf{v} \\ \frac{d \mathbf{v}}{d a} & = \frac{1}{a^2 E(a)} F_\theta(\mathbf{x}, a), \\ \end{array} \right. $$
With:
$$ F_\theta(\mathbf{x}, a) = \frac{3 \Omega_m}{2} \nabla \left[ \phi_{PM} (\mathbf{x}) \ast (1 + \color{orange}{f_\theta(a)}) \right] $$


$\Longrightarrow$ Correction integrated to the gravitational potential as a Fourier-based isotropic filter $f_{\theta}$

Hybrid Physical-Neural ODE

  • $\color{orange}{f_\theta(a)}$ implemented as a Fourier-based isotropic filter defined by a b-spline $\Longrightarrow$ incorporates translation and rotation symmetries


  • The NN takes as input the scale factor and returns the parameters defining the b-spline

Hybrid Physical-Neural ODE

$$\mathcal{L} = \sum_{i}^{snapshots} \lambda_1|| \mathbf{x}^{ref}_i - \mathbf{x}_i ||_2^2 + \lambda_2 || \frac{p_i(k)}{p_i^{ref}(k)} -1 ||_2^2 $$
  • We are following the technique from Neural ODEs to backpropagate through an ODE solver.
    Neural Ordinary Differential Equations, Chen et al. 2018
  • Benefits:
    • Tune the accuracy-speed by defining a tolerance of the ODE Solver
    • Reduced memory usage: no need to store intermediate steps for backpropagation.

Without neural correction

With neural correction

Potential Gradient Descent (PGD)

  • Additional displacements to sharpen the halos

  • The direction of the displacements points towards the halo center (local potential minimum).

  • The gravitational force: \begin{equation} \mathbf{F}=-\nabla\Phi \end{equation}
  • The PGD correction displacement: \begin{equation} \mathbf{S}=-\alpha\nabla \hat{O}_{h}\hat{O}_{l}\Phi \end{equation} High pass filter prevents the large scale growth, low pass filter reduces the numerical effect

Projections of final density field



Camels simulations
PM simulations
PM+NN correction
PM+PGD correction

Results

Results


Different resolution

Different resolution and volume

Different Cosmology!

Results


Different resolution

Different Cosmology!

First application: Differentiable Lensing Lightcone

DLL (Differentiable lensing light cone)
  • Exports the density planes fom the output of the N-body simulation

  • Performs ray-tracing through the lightcone by implementing the Born approximation

  • Provides derivatives with respect to cosmological and nuisance parameters through automatic differentiation
⟹ ⟹
TensorFlow-based weak gravitational lensing package
\[\begin{equation} \kappa_{born}(\boldsymbol{\theta},\chi_s)= \frac{3H_0^2 \Omega_m}{2c^2} \int_0^{\chi_s} d\chi \frac{\chi}{a(\chi)} W(\chi,\chi_s) \delta(\chi \boldsymbol{\theta},\chi). \end{equation} \]
$y = a * x + b$
↓ $$ y = u + b \qquad u = a * x $$ ↓ $$\frac{\partial y}{\partial x} = \frac{\partial y}{\partial u} \frac{ \partial u}{\partial x} = 1 \times a = a$$

A first application: Fisher forecasts

  • Comparison of Fisher constraints from two map-based weak lensing statistics.

    • Conclusion




    takeaway message
    • Powerful ways to simulate fast and differentiable cosmological simulations

    • Data driven way of complementing a physical models
      • correction scheme to compensate for the small-scales approximations

    • Differentiable physical models for fast inference
      • Likelihood-free inference approach
      • Inference of the full posterior distribution by using the Hamiltonian Monte Carlo (HMC) method

    Thank you !