MADLens, a python package for fast and differentiable non-Gaussian lensing simulations

Measurements of the weak cosmic shear signal will be among the major experimental drivers for advancing cosmology in the next decade. Next generation surveys such as LSST (LSST Science Collaboration et al., 2009), the Roman Space Telescope (Spergel et al., 2015) and the EUCLID satellite (Euclid Collaboration et al., 2020) will provide an unprecedented amount of high resolution weak cosmic shear data, which creates a demand for novel data analysis and modeling techniques. The weak cosmic shear signal is sensitive to the evolution of matter clustering over several orders of scales, ranging from well within the linear to the highly non-linear regime. Tomographic lensing measurements are sensitive to the total matter content, $\Omega_{m0}$, the amplitude of clustering, $\sigma_{8}$, and the time evolution of clustering, which allows to constrain dark energy (Hu, 2002; Huterer, 2002; Song and Knox, 2004) and the sum of neutrino masses. Weak cosmic shear measurements can further be used to test general relativity (Heavens et al., 2007; Schmidt, 2008; Harnois-Déraps et al., 2015).

Traditional lensing analyses mostly rely on two-point statistics or related observables (Kitching et al., 2011; Heymans et al., 2013; Kitching et al., 2014; Alsing et al., 2016; Hildebrandt et al., 2017; DES Collaboration et al., 2018; Taylor et al., 2019; Hikage et al., 2019) to extract cosmological information. However, since the lensing convergence field is inherently and significantly non-Gaussian, two-point statistics do not exploit its full information content. In fact, a long list of studies have shown that non-Gaussian summary statistics, such as higher order correlation functions (Pen et al., 2003; Takada and Jain, 2003; Jarvis et al., 2004; Semboloni et al., 2011; Fu et al., 2014; Coulton et al., 2019) or peak statistics (Jain and Van Waerbeke, 2000; Dietrich and Hartlap, 2010; Maturi et al., 2011; Marian et al., 2012; Pires et al., 2012; Petri et al., 2013; Cardone et al., 2013; Lin and Kilbinger, 2015; Liu et al., 2015a, b; Kacprzak et al., 2016; Peel et al., 2017; Li et al., 2019) can break parameter degeneracies that cumber power spectrum analyses and lead to significantly tighter constraints. While inference from these summaries offers improvements over power spectra analyses, their choice is somewhat ad hoc and the question of how to best extract cosmological information from non-Gaussian lensing scales is still an active field of research.

A number of works have recently suggested machine-learning tools for identifying informative summaries (Gupta et al., 2018; Ribli et al., 2019) and even successfully applied them to real data (Fluri et al., 2019). Machine learning methods require a large amount of training data. If applied correctly, these methods learn which features in the lensing map are most informative about cosmological parameters. A recent study finds that these models are mostly sensitive to extreme values in the lensing field (Zorrilla Matilla et al., 2020). This underlines the importance of training data that accurately mimics real cosmic shear data and its dependence on cosmological parameters down to very small scales.

Another approach, and in principle the optimal one, is to build a differentiable non-linear data model that starts from the Gaussian initial conditions and forward models them accurately to the measured lensing signal. This forward model is used to model the posterior of the parameters of interest (these can be cosmological parameters, the modes of the initial field or the bandpowers of the initial power spectrum). Analyzing this posterior generally relies on powerful sampling or optimization schemes, which in turn require many model evaluations and the derivatives of the model with respect to the parameters of interest. Forward-model based inference schemes have been developed for a range of of observables in cosmology (Seljak, 1998; Seljak et al., 2017; Jasche and Kitaura, 2010; Jasche and Wandelt, 2013) including weak cosmic shear (Böhm et al., 2017; Porqueres et al., 2020).

All of these new avenues for lensing analyses create the need for fast and differentiable simulations of the lensing field that at the same time accurately capture the nonlinear features of that field.

Realistic lensing simulations are challenging because a range of scales in the three dimensional matter distribution contribute to a single angular scale in the lensing field making even intermediate lensing scales sensitive to the non-linearity of structure formation on small scales (Jain and Seljak, 1997). Accurate lensing simulations rely on lightcones constructed from high resolution N-body simulations. These N-body simulations must accurately resolve small scales, but must at the same time be large enough to produce lensing maps with an extent of several degrees.

A number of recent works have studied the applicability of deep generative models, in particular generative adversarial networks (GANs) for producing accurate lensing convergence maps at low computational costs (Mustafa et al., 2019; Perraudin et al., 2020). These models do not aim at simulating the underlying physics, but are trained to mimic the training data to a degree where their output becomes indistinguishable from the training data for a neural network. While these early studies look promising, future research will have to show that these models indeed learn the correct data distribution or that using their outputs for inference leads to unbiased parameter posteriors. Another, more safeguard approach is to use machine-learning inspired techniques to boost the accuracy of low resolution N-body simulations (Dai et al., 2020; Dai and Seljak, 2020) and to construct lightcones from those. This is the avenue we have chosen in this work to create high resolution lensing simulations from approximate N-body solvers.

In this publication we describe a new, weak gravitational lensing package, MADLens. MADLens is a python package that allows to compute fully nonlinear lensing convergence maps at different source redshifts and low computational cost while accurately modeling the non-Gaussianity of the field down to scales of several tens of arcseconds. MADLens is built on top of a particle-mesh solver that evolves an initial linear density field into non-linear late time density fields. It provides derivatives with respect to the initial conditions of the particle-mesh simulation and a number of cosmological parameters through automated differentiation. MADLens fills the gap in the accuracy-speed space between computationally expensive, high accuracy lensing simulations and fast approximate simulations. In particular, MADLens correctly captures scales down to $L{=}10000$ at a field of view (FOV) of $6.2^{\circ}$ with percent level precision at a runtime of 30 seconds on 32 processes. MADLens can be run at different levels of resolution and approximations, allowing the user to choose the speed to accuracy trade-off that is optimal for their application.

We begin this paper with a brief introduction of the cosmological-scale weak lensing formalism and our notation in Section 2. This is followed by a in depth discussion of the package design in Section 3. We demonstrate the packages abilities in a number of tests in Section 4 and conclude with a summary and outlook in Section 5. A provides details on the novel feature of differentiability with respect to cosmological parameters.

Weak gravitational lensing observations provide insight into the projected matter density distribution between an observer and a source through correlated image distortions. Here, we provide a brief overview and define our usage of the lensing kernel, lensing convergence, and power spectrum of cosmic shear used throughout MADLens (for a detailed discussion of weak gravitational lensing and especially cosmic shear, see Bartelmann and Maturi (2017); Kilbinger (2015); Bartelmann and Schneider (2001)).

The image of a source galaxy at a comoving distance $\chi_{s}$ is distorted along the line of sight by some lensing potential $\Psi$. The potential of an extended lens under the Born approximation representing all density fluctuations along a line of sight at some angular position $\vec{\theta}$ can be found by integrating individual Weyl potentials $\Phi$ up to the comoving distance of the galaxy,

where $\chi_{\vec{\theta}}$ is the angular perpendicular component of the potential and $\chi$ is the parallel component. The lensing deflection and convergence are defined as,

where the derivatives are taken with respect to $\vec{\theta}$. Using the Poisson equation and neglecting derivative along the line-of-sight direction, the lensing convergence can be rewritten as,

where $\delta$ represents the density contrast from the mean density, $H_{0}$ is the Hubble parameter and $q$ is the lensing kernel describing the projection of sources selected by the redshift selection function $n(\chi^{\prime})$. MADLens evaluates the integral in Eq. 3 numerically. The most commonly used summary statistic in lensing analyses that can also be computed analytically is the power spectrum. To compute the convergence power spectrum we use Limber’s approximation and the flat sky approximation which are both valid on intermediate and small scales,

Since the lensing power spectrum is very sensitive to non-linear corrections to the matter power spectrum, we use HaloFit (Takahashi et al., 2012) to model $P_{m}(k)$ throughout this paper.

The MADLens package is based on FastPM (Feng et al., 2016), a highly scalable particle-mesh solver, that evolves particle positions through a kick and drift scheme enforcing correct linear displacement in each step. FastPM has been implemented in C and Python and two versions of FastPM support automatic differentiation, including the MPI based version used in this work, and the recently published FlowPM (Modi et al., 2020) package, which is based on TensorFlow. A FastPM particle-mesh simulation requires the choice of a particle-mesh resolution, equivalent to the number of particles in the simulation, the force resolution, the resolution of the grid onto which the particles are painted to compute the forces, the box size of the simulation and the number of steps in which the particle positions are evolved. The initial conditions and particle evolution depend on the cosmological parameters $\sigma_{8}$ and $\Omega_{m0}$.

MADLens runs a FastPM simulation. As the simulation evolves, MADLens projects the particles in the simulation weighted by the lensing kernel at each simulation step to 2D meshes at the desired source redshifts. The field of view, i.e. the size of the convergence map, and its resolution can be set by the user given the simulation box covers the entire field of view at the most distant source redshift. Each FastPM step evolves particles in redshift steps, $\Delta z_{i}=z_{i+1}{-}z_{i}$. MADLens constructs the lightcone by translating these redshifts into distances $\Delta\chi_{i}=\chi_{i+1}{-}\chi_{i}$ and projecting particles at the correct evolution step corresponding to that distance onto the convergence map. If the distance between two simulation steps is larger than the extent of the box, the box is replicated at the same redshift as often as is needed to fill the entire extent. In order to avoid spurious correlations, the simulations box is rotated before being repeated. In these techniques, MADLens is constructed similarly to other lightcone packages. We provide an overview of all MADLens parameters that can be set by the user in table 1.

To reach extraordinary accuracy at low computational costs, MADLens employs two special techniques:

• 1.

  Particle Gradient Descent (PGD) (Dai et al., 2018) is an additional particle evolution step that corrects for the difference between particle distributions in a low resolution simulation and a high resolution simulation. The correction is applied after each simulation step. PGD introduces 5 additional nuisance parameters, which are fitted on training simulations. PGD allows simulations to run at lower resolution while still obtaining results that are comparable and highly correlated with a high resolution simulation.

• 2.

  A sub-evolution step allows for a massive reduction of the number of simulations steps. When using sub-evolution, particles are evolved according to the redshift of their position within the simulation box before projection, rather than by the redshift of the FastPM step.

Finally, MADLens is made differentiable through numerically accurate tape-based automatic differentiation. Specifically it uses the Virtual Machine Automated Differentiation package (VMAD). VMAD builds a graph that is traversed for the model evaluation, during which all operations are recorded on a sequential tape. Gradient graphs generated from the tape are used to compute Jacobian vector products ($J_{ij}v_{j}$) and vector Jacobian products ($v_{i}J_{ij}$, commonly referred to as back-propagation). MADLens is made available in two variants. In its main version it is built to provide differentiability with respect to the initial, Gaussian modes of the simulation. A second package version, that is included in this release, adds differentiability with respect to the cosmological parameters, $\Omega_{m0}$ and $\sigma_{8}$. Differentiability with respect to the PGD paramaters ($k_{l}$, $k_{s}$, $\alpha_{0}$, $\mu$) will be included in a future release.

We analyze the performance of MADLens, the PGD enhancement, the sub-evolution scheme and computation times as well as the accuracy of the gradient computation. Testing the accuracy of MADLens output is challenging, because of the lack of a ground truth. We will use theoretical power spectra based on Halofit matter power spectra and high resolution runs for comparison. These can serve as a reasonable baseline, but as should become evident from our analysis, should not be mistaken for the ground truth.

In Figure 1 we show an example of a convergence map produced with MADLens. The non-Gaussian structure is clearly visible by eye and also evident in the corresponding histogram in Figure 2. The map resolves sub arcmin scales and extends over more than $6^{\circ}$ on the sky showing that MADLens overcomes one of the key challenges in lensing simulations: the accurate modeling of both large and small scales.

We compare the histogram of pixel values in Figure 2 with a log-normal distribution. Log-normal distributions have been used in the past to model lensing PDFs, but are not strictly theoretically motivated. The log-normal fit captures the rough shape of the distribution, but underestimates the probability of high values and overestimates the probability of low values.

We have presented MADLens, a fully differentiable python package for producing non-Gaussian convergence maps of weak gravitational lensing on cosmological scales. MADLens reaches unprecedented accuracy even when compared to many non-differentiable lensing simulations, and operates at run times of the order of 2-20 below conventional N-body simulation based lightcone packages. These advancements are made possible by several features, including a machine learning inspired post processing step, that allows the N-body simulation to run at a lower resolution and with less steps without paying a significant penalty in accuracy. Taking the derivative through a MADLens simulation with respect to the initial modes of the N-body simulation and the two key cosmological parameters $\sigma_{8}$ and $\Omega_{m0}$ is made possible through back-propagation. This means that evaluating the derivatives has comparable computational cost as the forward simulation. With these features MADLens constitutes a milestone towards the development of fully differentiable inference pipelines for weak cosmic shear. In the future MADLens will be integrated into the tensorflow-based FlowPM framework. Package updates will also feature differentiability with respect to nuisance parameters, such as the PGD parameters.

In the interest of scientific advancement and reproducibility, we make the MADLens package publicly available on github ¹¹1https://github.com/VMBoehm/MADLens.

This research used resources of the National Energy Research Scientific Computing Center (NERSC), a U.S. Department of Energy Office of Science User Facility operated under Contract No. DE-AC02-05CH11231. The authors thank Uros Seljak for useful discussions and feedback and Francois Lanusse for testing of the package and helpful feedback.