tidy; more benchmarks

Author: alexander-mead

notebooks/demo.ipynb

@@ -7,7 +7,7 @@
    "source": [
     "## HMcode demo\n",
     "\n",
-    "Here we compare the HMcode power spectrum from pyHMcode to that from CAMB and to the halo-model.\n",
+    "Here we compare the `HMcode` power spectrum from `pyHMcode` to that from `CAMB` and to the halo-model (from `pyhalomodel`).\n",
     "\n",
     "First, import the necessary libraries"
    ]
@@ -67,7 +67,6 @@
     "\n",
     "# Redshifts\n",
     "zs = [3., 2., 1., 0.5, 0.]\n",
-    "#zs = [0.]\n",
     "zs = np.array(zs)\n",
     "\n",
     "# Halo masses [Msun/h] (for halo model only)\n",
@@ -81,7 +80,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Run CAMB to get the linear spectrum and $\\sigma(R)$"
+    "Run `CAMB` to get the linear spectrum and $\\sigma(R)$"
    ]
   },
   {
@@ -130,7 +129,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Get the HMcode spectrum, note that it takes the results data structure from CAMB as an argument"
+    "Get the `HMcode` spectrum, note that it takes the results data structure from `CAMB` as an argument"
    ]
   },
   {
@@ -148,7 +147,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Calculate the halo-model matter-matter power spectrum using pyhalomodel"
+    "Calculate the halo-model matter-matter power spectrum using `pyhalomodel`"
    ]
   },
   {

pyhmcode/cosmology.py

@@ -12,7 +12,7 @@
 # Parameters
 xmin_Tk = 1e-5    # Scale at which to switch to Taylor expansion approximation in tophat Fourier functions
 eps_sigmaR = 1e-4 # Accuracy of the sigmaR integration
-eps_sigmaV = 1e-4 # Accuracy of the sigmaV integration NOTE: Seems to fail with higher accuracy
+eps_sigmaV = 1e-4 # Accuracy of the sigmaV integration NOTE: Seems to fail with higher accuracy (1e-4; when zs=[0] only?!?)
 eps_neff = 1e-4   # Accuracy of the neff integration (d ln sigma^2/d ln k)
 
 ### Backgroud ###
@@ -110,14 +110,14 @@ def _transformed_integrand(t:np.ndarray, R:float, Pk:callable, integrand:callabl
 def _sigmaR_integrand(k:np.ndarray, R:float, Pk:callable) -> np.ndarray:
     '''
     Integrand for calculating sigma(R)
-    Note that k can be a float or an array here
     '''
     return Pk(k)*(k**2)*_Tophat_k(k*R)**2
 
 
-def sigmaR(R:np.ndarray, Pk:callable, eps=eps_sigmaR, transform_integrand=False) -> np.ndarray:
+def sigmaR(R:np.ndarray, Pk:callable, kmin=0., kmax=np.inf, eps=eps_sigmaR, transform_integrand=False) -> np.ndarray:
     '''
-    Quad integration TODO: This is slow, there must be something better to do
+    Quad integration to get sigma(R)
+    TODO: This is slow, there must be something better to do
     args:
         R: Comoving Lagrangian radius [Mpc/h]
         Pk: Function of k to evaluate the linear power spectrum
@@ -127,7 +127,6 @@ def sigmaR_vec(R:float, Pk:callable):
              sigmaR_squared, _ = integrate.quad(lambda t: _transformed_integrand(t, R, Pk, _sigmaR_integrand), 
                                                 0., 1., epsabs=eps, epsrel=eps)
         else:
-            kmin, kmax = 0., np.inf
             sigmaR_squared, _ = integrate.quad(lambda k: _sigmaR_integrand(k, R, Pk), 
                                                kmin, kmax, epsabs=eps, epsrel=eps)
         sigmaR = np.sqrt(sigmaR_squared/(2.*np.pi**2))
@@ -144,12 +143,12 @@ def _sigmaV_integrand(k:np.ndarray, R:float, Pk:callable) -> np.ndarray:
     return integrand
 
 
-def sigmaV(R:float, Pk:callable, eps=eps_sigmaV, transform_integrand=False) -> float:
+def sigmaV(R:float, Pk:callable, kmin=0., kmax=np.inf, eps=eps_sigmaV, transform_integrand=False) -> float:
     '''
     Integration to get the 1D RMS in the linear displacement field
-    TODO: This generates a warning sometimes, there must be a cleverer way to integrate here.
+    TODO: This generates a warning sometimes, and is slow, there must be a cleverer way to integrate here.
     Unless eps_sigmaV > 1e-3 the integration fails for z=0 sometimes, but not after being called for
-    z > 0. I really don't understand this, but it's annoying and should be fixed.
+    z > 0. I really don't understand this, but it's annoying and should be fixed (kmin /= 0. helps).
     I should look at how CAMB deals with these type of integrals (e.g., sigmaR).
     args:
         Pk: Function of k to evaluate the linear power spectrum
@@ -159,7 +158,6 @@ def sigmaV(R:float, Pk:callable, eps=eps_sigmaV, transform_integrand=False) -> f
         sigmaV_squared, _ = integrate.quad(lambda t: _transformed_integrand(t, R, Pk, _sigmaV_integrand), 
                                            0., 1., epsabs=eps, epsrel=eps)
     else:
-        kmin, kmax = 0., np.inf
         sigmaV_squared, _ = integrate.quad(lambda k: _sigmaV_integrand(k, R, Pk), 
                                            kmin, kmax, epsabs=eps, epsrel=eps)
     sigmaV = np.sqrt(sigmaV_squared/(2.*np.pi**2))
@@ -171,17 +169,16 @@ def _dsigmaR_integrand(k:float, R:float, Pk:callable) -> float:
     return Pk(k)*(k**3)*_Tophat_k(k*R)*_dTophat_k(k*R)
 
 
-def _dlnsigma2_dlnR(R:float, Pk:callable, eps=eps_neff, transform_integrand=False) -> float:
+def _dlnsigma2_dlnR(R:float, Pk:callable, kmin=0., kmax=np.inf, eps=eps_neff, transform_integrand=False) -> float:
     '''
     Calculates d(ln sigma^2)/d(ln R) by integration
     3+neff = -d(ln sigma^2) / dR
-    TODO: Vectorize
+    TODO: This is slow, regardless of whether transform_integrand is set or not
     '''
     if transform_integrand:
         dsigma, _ = integrate.quad(lambda t: _transformed_integrand(t, R, Pk, _dsigmaR_integrand), 
                                    0., 1., epsabs=eps, epsrel=eps)
     else:
-        kmin, kmax = 0., np.inf # Evaluate the integral and convert to a nicer form
         dsigma, _ = integrate.quad(lambda k: _dsigmaR_integrand(k, R, Pk), 
                                    kmin, kmax, epsabs=eps, epsrel=eps)
     dsigma = R*dsigma/(np.pi*sigmaR(R, Pk))**2

pyhmcode/linear_growth.py

@@ -67,6 +67,7 @@ def get_growth_interpolator(CAMB_results:camb.CAMBdata, LCDM=False) -> callable:
     '''
     Solve the linear growth ODE and returns an interpolating function for the solution
     LCDM = True forces w = -1 and imposes flatness by modifying the dark-energy density
+    TODO: w dependence for initial conditions; f here is correct for w=0 only
     '''
     from scipy.integrate import solve_ivp
     from scipy.interpolate import interp1d as interp

pyhmcode/pyhmcode.py

@@ -58,7 +58,11 @@ def hmcode(k:np.array, zs:np.array, CAMB_results:camb.CAMBdata,
         print()
 
     # Linear power interpolator
-    Pk_lin_interp = CAMB_results.get_matter_power_interpolator(nonlinear=False, extrap_kmax=Pk_lin_extrap_kmax).P
+    interp, _, k_interp = CAMB_results.get_matter_power_interpolator(nonlinear=False, 
+                                                                     return_z_k=True,
+                                                                     extrap_kmax=Pk_lin_extrap_kmax)
+    Pk_lin_interp = interp.P
+    kmin = k_interp[0] # Minimum wavenumber used for the CAMB interpolator [h/Mpc]
 
     # Loop over redshift
     Pk_HMcode = np.zeros((len(zs), len(k)))
@@ -111,9 +115,10 @@ def hmcode(k:np.array, zs:np.array, CAMB_results:camb.CAMBdata,
             print('Peak height range: {:.4} -> {:.4}'.format(nu[0], nu[-1]))
 
         # Parameters of the linear spectrum pertaining to non-linear growth
+        # TODO: I think the sigmaV integral is quite time-consuming for some reason
         Rnl = _get_nonlinear_radius(R[0], R[-1], dc, iz, CAMB_results, cold=True) # Non-linear Lagrangian radius
         sigma8 = _get_sigmaR(8., iz, CAMB_results, cold=True)                     # RMS in the linear cold matter field at 8 Mpc/h
-        sigmaV = cosmology.sigmaV(0., lambda k: Pk_lin_interp(z, k))              # RMS in the linear displacement field
+        sigmaV = cosmology.sigmaV(0., lambda k: Pk_lin_interp(z, k), kmin=kmin)   # RMS in the linear displacement field
         neff = _get_effective_index(Rnl, R, sigmaM)                               # Effective index of spectrum at collapse scale
         if verbose:
             print('Non-linear Lagrangian radius: {:.4} Mpc/h'.format(Rnl))

tests/benchmarks/cosmologies.txt

@@ -9,3 +9,18 @@
 3.529997714832051603e-01 5.404154960001772423e-02 5.357940065799579599e-03 7.236828642981654180e-01 9.303950098062612195e-01 7.061635669135878368e-01 -1.037969566460582582e+00 -1.084100134813217098e+00 4.085286437246361535e-01
 3.706806146536332225e-01 4.201818457596022405e-02 -4.416972583109340239e-02 6.125535568087986382e-01 9.293593757766683838e-01 8.323833029453789889e-01 -9.657807085952330173e-01 6.295336094103043401e-01 6.643135403273875417e-01
 2.812773722880141380e-01 5.942061153998104417e-02 -3.330270800922961044e-02 5.090848292535441466e-01 9.090047860775641864e-01 8.444718701192900356e-01 -1.022873661849167615e+00 -1.244571945108858602e+00 5.010447751033635377e-01
+2.304624205426336891e-01 5.588961125233208455e-02 -5.384372442596929709e-03 6.524084904385929473e-01 9.301512089147876416e-01 8.260565186237777136e-01 -1.082912433667965857e+00 -1.466173742858535967e+00 1.180059021205153158e-01
+3.923795329099029061e-01 6.225742072122820991e-02 1.997071338107495042e-02 6.063479845838077864e-01 9.969176377347723772e-01 8.557501807931588989e-01 -8.698658865046026145e-01 -3.774218785471961102e-01 2.722415618451590191e-01
+2.192781924306998753e-01 6.207807189631525668e-02 -4.422371016638892982e-03 5.809453459180921042e-01 9.305956624150653056e-01 8.158439137883791714e-01 -1.193936330236460952e+00 8.484089951180502354e-01 7.585195298352100934e-01
+3.438925911901873356e-01 4.796279119325311258e-02 1.273088407024432211e-02 7.336391875650942751e-01 9.649846601554820102e-01 7.168888642279778178e-01 -1.050515558697634244e+00 -1.604741336675703511e+00 4.939908192445189306e-01
+2.659722424665570961e-01 3.933572566598141085e-02 -3.965970322774484136e-02 7.350578288710848396e-01 9.170592968536885881e-01 8.850240236753594836e-01 -9.513633161797629967e-01 -6.859218883501405628e-01 5.909154914814167547e-01
+2.045607742059395162e-01 6.375677639724336432e-02 -1.769656305709979560e-03 8.130940909001145300e-01 9.082729999922438680e-01 7.973316661676320694e-01 -1.005575803387287470e+00 -9.098562348013317091e-03 4.734894010569538070e-01
+2.533951326183787600e-01 4.494706992027656944e-02 2.067240247153781518e-03 6.755645841220186965e-01 9.021612079880330093e-01 8.652583848388715193e-01 -7.623035368961399438e-01 -1.307850242114181549e+00 5.540361435390493972e-01
+2.217151482270887097e-01 5.516720279119435855e-02 -2.187662161609917055e-02 7.637690538767607507e-01 9.726994614286882390e-01 8.537294983835314044e-01 -1.235355432426462174e+00 1.027195653864199532e+00 2.302139908948808156e-01
+2.074825112352359779e-01 5.164557408174449971e-02 -1.290777161375612769e-02 8.319158972529652196e-01 9.808251472064302146e-01 7.634277785645430248e-01 -7.282603629581529958e-01 -8.543373071982429678e-01 5.150571292317145389e-01
+2.511930181135205720e-01 6.308130710146890585e-02 -3.353921824179818822e-02 5.179642477569316172e-01 9.435097060003038161e-01 8.984751128111674312e-01 -7.649936402470516317e-01 5.233101385654168602e-01 8.907924908785248563e-01
+3.786893279395726619e-01 5.056575081159347257e-02 -1.840709481692070354e-02 8.088049728443952313e-01 9.661661263167761193e-01 7.747315457747420142e-01 -1.243319999163090772e+00 5.178367301605679707e-01 2.624605159228646789e-01
+3.873626301067558675e-01 4.222911725017054752e-02 -3.772420675885140157e-02 8.324450688499624995e-01 9.153284316624493711e-01 7.358536616315477508e-01 -9.403703250874939634e-01 9.024317429207679986e-01 1.964346657145732422e-01
+2.620647345800189987e-01 5.832214514723532695e-02 4.718264260609673910e-02 7.002964744809369035e-01 9.143897502551251444e-01 7.027872575416402867e-01 -1.162206382000686933e+00 -1.333215124462571621e+00 6.776586736128574762e-01
+2.243665009257062304e-01 5.018989794861899306e-02 1.942624356428865329e-02 7.324466436883609877e-01 9.199775651660058173e-01 8.608249052364524978e-01 -8.707557222305190026e-01 4.943418517857804950e-01 1.310577515573132512e-01
+2.247507607300689436e-01 6.282687653019522478e-02 -1.024218061750593878e-02 6.203794767123759213e-01 9.488584045351533547e-01 8.325728425527164633e-01 -7.266260457718179966e-01 -8.677968570850144170e-01 9.248084293120271138e-01