Emulating cosmological growth functions with B-Splines

Field level inference for cosmological analysis simulates the survey observations data at the level of full field by starting all the way from the initial density distribution at the beginning of the Universe and then evolving dark matter particles under gravity with Particle Mesh (PM) N-Body simulations. The goal then is to infer the cosmological parameters along with the initial conditions at all points in the Universe. This challenging high dimensional inference relies on using differentiable simulations [1, 2, 3, 4] and coupling them with gradient based algorithms such as Hamiltonian Monte Carlo (HMC) [5]. Due to the iterative nature of these algorithms, it is crucial for these simulators to be simultaneously fast and accurate.

Recent works have used advances in automatic differentiation libraries such Tensorflow and Jax to build these simulators like FlowPM [2] and pmwd [4]. An added advantage of this is that the simulations can now exploit efficient GPU parallelizations and accelerations for significant speed-ups. However this has inadvertently made other sequential operations a bottleneck to fully realize potential gains. One such example is solving ordinary differential equations (ODE) wherein gradient calculation has to backpropagate through all the sequential computations of the integrator [6]. In cosmological simulations, the growth factor of density fields and distance functions are estimated by solving a system of ODEs as a function of time. We find that for small to intermediate PM simulations, backpropogating through growth function ODE can be the majority time cost. Thus in this work, we seek to replace this ODE solution with trained emulators.

In the past couple years, many works have built emulators for time intensive operations in cosmological analysis pipeline [7, 8]. One way these works learn the emulator is by training a multi-layer perceptron (MLP) to fit the quantity of interest at some fixed points in the domain. Then during analysis, they construct an interpolating function through these points. However, it is redundant to query an MLP and then the interpolating function. For accurate interpolation, these output points also need to be densely sampled, making output of MLP high dimensional and challenging to learn. An alternative approach is to learn the coefficients for principal components basis but that limits the accuracy based on the number of principal components used. Depending on the analysis, it also does not necessarily solve the interpolation redundancy.

In this work, we take a different approach and learn the emulator by directly learning the components of an interpolating function. Specifically, we learn a B-Spline emulator [9] that can be parameterized by a small number of knot points and weights to model a smooth, high order differentiable function. We begin by setting up the growth function ODE in section 2 and follow it with discussing our emulator in section 3. In section 4, we show that the emulator meets desired accuracy and quantify achieved gains before concluding with a brief discussion in section 5.

Cosmological simulations evolve dark matter over-density field under gravitational force in an expanding universe. This evolution is governed by a solving a system of coupled non-linear partial differential equations called Vaslov Poisson equations. However at the linear order and after making a number of simplifying assumptions not detailed here for brevity, this evolution is governed by the following ODE [10]

$$\frac{\partial^{2}\delta(x,t)}{\partial t^{2}}+2{H(t)}\frac{\partial\delta(x,t)}{\partial t}=\frac{3}{2}\Omega_{m}(t)H_{0}^{2}\delta(x,t)$$

(1)

where ${H}(a)=H_{0}(\frac{\Omega_{m}}{a^{3}}+\Omega_{\Lambda})^{0.5}$ is the Hubble parameter, $H_{0}$ is the Hubble constant, $\Omega_{m}(a)=\frac{\Omega_{m}}{a^{3}}\frac{H_{0}^{2}}{H^{2}(a)}$ is the matter density, $\Omega_{\Lambda}$ is the dark energy density and $a$ is the scale factor of the Universe, defined as $H(a)=\frac{1}{a}\frac{da}{dt}$. For simplicity, in the following we use $a$ as a measure of time instead of $t$ since that is the convention in the particle mesh simulations of interest here.

At linear order, the evolution of the density field can be decoupled into spatial and time contributions by writing $\delta(x,a)=D_{1}(a)\delta(x,a_{0})$ for some “reference" time $a_{0}$ and $D_{1}(a)$ called the linear growth factor [11]. Then the time dependence of the growth factor is

$$a^{2}\frac{d^{2}D_{1}(a)}{da^{2}}+\Bigg{(}\Omega_{\Lambda}(a)+\frac{\Omega_{m}(a)}{2}+2\Bigg{)}a\frac{dD_{1}(a)}{da}=\frac{3}{2}\Omega_{m}(a)D_{1}(a)$$

(2)

Being a second order equation, this has two solutions. One of them increases with time and is referred to as growing mode ($D_{1}^{+}$) while the other decays with time and hence is generally ignored in the simulations. At higher orders in the Taylor expansion of the density field, we can define a similar growth function at the second order for $\delta^{(2)}(x,a)$ that follows the ODE

$$a^{2}\frac{d^{2}D_{2}(a)}{da^{2}}+\Bigg{(}\Omega_{\Lambda}(a)+\frac{\Omega_{m}(a)}{2}+2\Bigg{)}a\frac{dD_{2}(a)}{da}=\frac{3}{2}\Omega_{m}(a)\big{[}D_{2}(a)-(D_{1}^{+}(a))^{2}\big{]}$$

(3)

We are interested in constructing an emulator for the growth factor and its derivatives as a function of cosmology. In the standard $\Lambda$CDM model, Eq. 2 tells us that the growth factor is only a function of $\Omega_{m}$ since $\Omega_{m}+\Omega_{\Lambda}=1$. Hence our emulator has only one input.

We learn a B-spline (or basis spline) function [9] to model the growth function. B-splines are powerful since they are the maximally differentiable interpolative basis function and any spline function of a given degree $n$ can be expressed as a unique linear combination of B-splines of that degree. B-splines are defined by the number of interior ‘knots’ N. Thus let $t_{0},t_{1},...,t_{N},t_{N+1}$ be a non-decreasing sequence of knots. These knots are augmented by repeating exterior knots $t_{0}$ and $t_{N+1}$ $n$ times and for each augmented knot $t_{i}$, a set of basis functions $B_{i,j},\,\forall j=0,1,...,n$ is defined recursively from $j=0$ to $n$. A B-spline function is then the linear combinations of these basis functions with weights $\beta_{i}$

$$B(x)=\sum_{i=0}^{i=N+n}\beta_{i}B_{i,n}(x),\quad x\in[t_{0},t_{N+1}]$$

(5)

Thus a B-spline function $B(x)$ is completely defined given a sequence of knots $t_{i}$ and weights $\beta_{i}$. In our emulator, we train an MLP to predict these knots and weights for an input cosmology.

We train our emulator as a function $\Omega_{m}$ to predict the knot points and weights of cubic $(n=3)$ B-spline function. The emulator works in the range of $\Omega_{m}\in[0.1,0.5]$. $a$ is in the range of 0 to 1. 1000 random values of $\Omega_{m}$ from 0.05 to 0.55 are generated by uniform sampling for training. There are 256 points evenly distributed in the range of $a$, including 0 and 1. The training, validation and testing data sets are of size 800, 100 and 100 respectively.

We find that using only 8 knots for $a\in[0,1]$ is sufficient for achieving sub-percent accuracy. For the MLP architecture, there is first an input layer. The output of the input layer is duplicated and fed into two parts in parallel, each having a hidden layer and an output layer. The two parts account for the calculation of knot points and weights respectively. The input layer and two hidden layers have 64 neurons each. The neurons of output layers depend on the number of knots. In our case, the MLP outputs 8 knot points and the corresponding weights. One knot is fixed at $a=0$. We implement our emulator as part of pmwd code [4]. We begin by showing the accuracy of our emulator in Figure 2.

We show errors in value and gradient for a reference cosmology of $\Omega_{m}=0.3$. In addition, for every cosmology, we also show the maximum error among all $a\in[0,1]$ in the last two panels. The error in value is always less than 0.002 and in gradient is always less than 0.05. Next, we show that this level of agreement does not impact the accuracy of our inference.

As noted in the beginning, the objective of using differentiable PM simulations is to be able to access the gradients of cosmology parameters for field level inference. We mock this pipeline with a simulated data to infer $\Omega_{m}$ and $A_{s}$ (scalar amplitude) parameters. Figure 3 shows the results for two gradient based approaches when using our emulator versus ODE formulation for growth function. In the left panel, we show the trajectory followed when doing optimization to find the MAP (maximum-a-posteriori) estimate of the cosmology parameters. The last two panels show the marginal posteriors sampled with HMC for the same problem. The agreement between our emulator and ODE solution validates that the accuracy of the emulator and its gradients is sufficient for correct inference.

Finally, in Figure 4, we show the timings for gradient evaluation when using our emulator versus the adaptive Dormand-Prince ODE solver and Runge Kutta (order 4, solved at 256 points in $a\in[0,1]$). For $N<128$, increasing the simulation size does not change the time cost, demonstrating that ODE poses the bottleneck. In this case, we can gain up to an order of magnitude in time with the emulator. For larger simulations, PM cost starts to dominate as expected and we gain up to a factor of 2.

In this work, we demonstrate that for highly optimized PM simulations exploiting GPU accelerations, sequential operations such as solving ODE for growth function can be the majority time cost for gradient evaluations. We replace them with novel B-Spline emulators which can lead to an order of magnitude gains for small to intermediate simulations. While the gains are less remarkable for large $N\geq 256$ simulations, it is important to note that most of the methodology development is done on the small simulations with large runs primarily done only at the final analysis. Hence these gains will still lead to non-trivial time savings in developing novel methods for cosmological inference. We also anticipate similar bottleneck for distance calculation in weak lensing simulation and plan to explore that in the future. In the next work, we plan to extend the emulator to account for other cosmological parameters, such as the curvature $\Omega_{k}$, and dark energy parameters $w_{0}$ and $w_{a}$.