Inference Techniques for Cosmological Forward Modeling

• 1000 images each night, 15 TB/night for 10 years


• 18,000 square degrees, observed once every few days


• Tens of billions of objects, each one observed $\sim1000$ times

• Measure the ellipticity $\epsilon = \epsilon_i + \gamma$ of all galaxies
  $\Longrightarrow$ Noisy tracer of the weak lensing shear $\gamma$


• Compute summary statistics based on 2pt functions,
  e.g. the power spectrum


• Run an MCMC to recover a posterior on model parameters, using an analytic likelihood $$ p(\theta | x ) \propto \underbrace{p(x | \theta)}_{\mathrm{likelihood}} \ \underbrace{p(\theta)}_{\mathrm{prior}}$$

• Instead of trying to analytically evaluate the likelihood $p(x | \theta)$, let us build a forward model of the observables.
  $\Longrightarrow$ The simulator becomes the physical model.


• Each component of the model is now tractable, but at the cost of a large number of latent variables.

• Mixture Density Networks (MDN) \begin{equation} p(\theta | x) = \prod_i \pi_i(x) \ \mathcal{N}\left(\mu_i(x), \ \sigma_i(x) \right) \nonumber \end{equation}


• Flourishing Machine Learning literature on density estimators
  GLOW, (Kingma & Dhariwal, 2018)

• Deep Residual Network $y = f_\phi(x)$ followed by mixture density network $q_\phi(\theta | y)$


• Training on weak lensing maps simulated for different cosmologies






• Optimization of the variational lower bound: $$\mathbb{E}_{(x, \theta) \in \mathcal{D}} [ \log q_\phi(\theta | f_\phi(y) ) ]$$

				      													import tensorflow as tf
				      													import flowpm
				      													# Defines integration steps
				      													stages = np.linspace(0.1, 1.0, 10, endpoint=True)

				      													initial_conds = flowpm.linear_field(32,       # size of the cube
				      													                                   100,       # Physical size
				      													                                   ipklin,    # Initial powerspectrum
				      													                                   batch_size=16)

				      													# Sample particles and displace them by LPT
				      													state = flowpm.lpt_init(initial_conds, a0=0.1)

				      													# Evolve particles down to z=0
				      													final_state = flowpm.nbody(state, stages, 32)

				      													# Retrieve final density field
				      													final_field = flowpm.cic_paint(tf.zeros_like(initial_conditions),
				      													                               final_state[0])
				      												

• Seamless interfacing with deep learning components
• Mesh TensorFlow implementation for distribution on supercomputers

• We developed a Mesh TensorFlow implementation that can scale on GPU clusters (horovod+NCCL).



• For a $2048^3$ simulation:
  • Distributed on 256 NVIDIA V100 GPUs
  • Runtime: 3 mins



• Don't hesitate to reach out if you have a use case for model parallelism!
  
  
  https://github.com/DifferentiableUniverseInitiative/mesh

Reconstructed initial conditions $z$

Reconstructed dark matter distribution $x = f(z)$

Data
$x_{DM} = f(z_0)$

• Pressure waves in the primordial photon-baryon fluids


• Frozen at the time of decoupling
  $\rightarrow$ Clustering at preferred scale


• Damping due to non-linear evolution


• Fisher information
  
  $F_{ij} = V_{\rm{eff}}\int_{k_{min}}^{k_{max}} \frac{\partial {\rm ln} P(k)}{\partial p_i} \frac{\partial {\rm ln} P(k)}{\partial p_j}\frac{4 \pi k^2}{2 (2\pi)^3} dk \quad \mathbf{\propto \, r_c^4} $

							import jax.numpy as np
							import jax_cosmo as jc

							# Defining a Cosmology
							cosmo = jc.Planck15()

							# Define a redshift distribution with smail_nz(a, b, z0)
							nz = jc.redshift.smail_nz(1., 2., 1.)

							# Build a lensing tracer with a single redshift bin
							probe = probes.WeakLensing([nz])

							# Compute angular Cls for some ell
							ell = np.logspace(0.1,3)
							cls = angular_cl(cosmo_jax, ell, [probe])
						

Validating against the DESC Core Cosmology Library

• Current cosmological MCMC chains take days, and typically require access to large computer clusters.


• Gradients of the log posterior are required for modern efficient and scalable inference techniques:
  • Variational Inference
  • Hamiltonian Monte-Carlo


• In jax-cosmo, we can trivially obtain exact gradients:

  
  											def log_posterior( theta ):
  													return gaussian_likelihood( theta ) + log_prior(theta)
  
  											score = jax.grad(log_posterior)(theta)
  											



• On a DES Y1 analysis, we find convergence in 70,000 samples with vanilla HMC, 140,000 with Metropolis-Hastings