Repository: skypyproject/skypy
Pull request #455 (2021-03-29 10:25:01)

ENH: Power spectrum base classes
## Description
Implements PowerSpectrum (and TabulatedPowerSpectrum) classes to provide a common interface for power spectrum calculations (using interpolation). Also implements EisensteinHu, CAMB and CLASS classes to replace previous raw functions.

N.b. data used when testing CLASS was found to be inaccurate and has been recalculated. However, even the new "correct" data does not match our interpolation method at the previous relative tolerance of 1e-4 so this has been reduced to 1e-3.

## Checklist
- [x] Follow the [Contributor Guidelines](https://github.com/skypyproject/skypy/blob/main/CONTRIBUTING.rst)
- [x] Write unit tests
- [ ] Write documentation strings
- [ ] Assign someone from your working team to review this pull request
- [ ] Assign someone from the infrastructure team to review this pull request



Patch:
diff --git a/skypy/power_spectrum/_base.py b/skypy/power_spectrum/_base.py
new file mode 100644
index 0000000..56e52ea
--- /dev/null
+++ b/skypy/power_spectrum/_base.py
@@ -0,0 +1,33 @@
+from abc import ABCMeta, abstractmethod
+from inspect import signature
+import numpy as np
+from scipy.interpolate import RectBivariateSpline
+
+
+class PowerSpectrum(metaclass=ABCMeta):
+    '''Base class for power spectrum calculation'''
+
+    @abstractmethod
+    def __init__(self):
+        raise NotImplementedError
+
+
+class TabulatedPowerSpectrum(PowerSpectrum):
+    '''Base class for power spectrum interpolation'''
+
+    def __init__(self, wavenumber, redshift, power_spectrum):
+        kx = min(np.size(redshift)-1, signature(RectBivariateSpline).parameters['kx'].default)
+        ky = min(np.size(wavenumber)-1, signature(RectBivariateSpline).parameters['ky'].default)
+        self.interpolator = RectBivariateSpline(redshift, np.log10(wavenumber),
+                                                np.log10(power_spectrum), kx=kx, ky=ky)
+
+    def __call__(self, wavenumber, redshift):
+        shape = np.shape(redshift) + np.shape(wavenumber)
+        sort_k = np.argsort(wavenumber)
+        sort_z = np.argsort(redshift)
+        unsort_k = np.argsort(sort_k)
+        unsort_z = np.argsort(sort_z)
+        interp = self.interpolator(np.atleast_1d(redshift)[sort_z],
+                                   np.log10(np.atleast_1d(wavenumber)[sort_k]))
+        pzk = np.power(10, interp[unsort_z[:, np.newaxis], unsort_k]).reshape(shape)
+        return pzk.item() if np.isscalar(wavenumber) and np.isscalar(redshift) else pzk
diff --git a/skypy/power_spectrum/_camb.py b/skypy/power_spectrum/_camb.py
index 95f8a52..fefe01f 100644
--- a/skypy/power_spectrum/_camb.py
+++ b/skypy/power_spectrum/_camb.py
@@ -1,9 +1,11 @@
 import numpy as np
 from astropy import units
 from inspect import signature
+from ._base import TabulatedPowerSpectrum
 
 
 __all__ = [
+    'CAMB',
     'camb',
 ]
 
@@ -89,3 +91,36 @@ def camb(wavenumber, redshift, cosmology, A_s, n_s):
     pzk = pk_interp.P(redshift[redshift_order], wavenumber[wavenumber_order])
 
     return pzk[redshift_order].reshape(return_shape)
+
+
+class CAMB(TabulatedPowerSpectrum):
+
+    def __init__(self, kmax, redshift, cosmology, A_s, n_s, **kwargs):
+
+        try:
+            from camb import CAMBparams, model, get_results
+        except ImportError:
+            raise Exception("camb is required to use skypy.power_spectrum.camb")
+
+        h2 = cosmology.h * cosmology.h
+
+        pars = CAMBparams()
+        pars.set_cosmology(H0=cosmology.H0.value,
+                           ombh2=cosmology.Ob0 * h2,
+                           omch2=cosmology.Odm0 * h2,
+                           omk=cosmology.Ok0,
+                           TCMB=cosmology.Tcmb0.value,
+                           mnu=np.sum(cosmology.m_nu.to_value(units.eV)),
+                           standard_neutrino_neff=cosmology.Neff
+                           )
+
+        pars.InitPower.ns = n_s
+        pars.InitPower.As = A_s
+        pars.NonLinear = model.NonLinear_none
+        k_per_logint, var1, var2, hubble_units, nonlinear = None, None, None, False, False
+        pars.set_matter_power(redshifts=redshift, kmax=kmax, k_per_logint=k_per_logint, silent=True)
+        results = get_results(pars)
+        k, z, p = results.get_linear_matter_power_spectrum(var1, var2, hubble_units,
+                                                           nonlinear=nonlinear)
+
+        super().__init__(k*cosmology.h, z, p)
diff --git a/skypy/power_spectrum/_classy.py b/skypy/power_spectrum/_classy.py
index cf705ee..ffc89c6 100644
--- a/skypy/power_spectrum/_classy.py
+++ b/skypy/power_spectrum/_classy.py
@@ -1,6 +1,8 @@
 import numpy as np
+from ._base import TabulatedPowerSpectrum
 
 __all__ = [
+    'CLASSY',
     'classy',
 ]
 
@@ -73,3 +75,35 @@ def classy(wavenumber, redshift, cosmology, **kwargs):
         pzk = pzk.item()
 
     return pzk
+
+
+class CLASSY(TabulatedPowerSpectrum):
+
+    def __init__(self, kmax, redshift, cosmology, **kwargs):
+        try:
+            from classy import Class
+        except ImportError:
+            raise Exception("classy is required to use skypy.linear.classy")
+
+        h2 = cosmology.h * cosmology.h
+
+        params = {
+            'output': 'mPk',
+            'P_k_max_1/Mpc':  kmax,
+            'z_pk': ', '.join(str(z) for z in np.atleast_1d(redshift)),
+            'H0':        cosmology.H0.value,
+            'omega_b':   cosmology.Ob0 * h2,
+            'omega_cdm': cosmology.Odm0 * h2,
+            'T_cmb':     cosmology.Tcmb0.value,
+            'N_eff':     cosmology.Neff,
+        }
+
+        params.update(kwargs)
+
+        classy_obj = Class()
+        classy_obj.set(params)
+        classy_obj.compute()
+
+        p, k, z = classy_obj.get_pk_and_k_and_z(nonlinear=False)
+        z_order = np.argsort(z)
+        super().__init__(k, z[z_order], p.T[z_order, :])
diff --git a/skypy/power_spectrum/_eisenstein_hu.py b/skypy/power_spectrum/_eisenstein_hu.py
index 1bdc993..8605861 100644
--- a/skypy/power_spectrum/_eisenstein_hu.py
+++ b/skypy/power_spectrum/_eisenstein_hu.py
@@ -7,9 +7,12 @@ formula for the matter power spectrum.
 from astropy.utils import isiterable
 import numpy as np
 from astropy import constants
+from ._base import PowerSpectrum
+from ._growth import growth_function_carroll
 
 
 __all__ = [
+    'EisensteinHu',
     'eisenstein_hu',
     'transfer_with_wiggles',
     'transfer_no_wiggles',
@@ -253,3 +256,21 @@ def eisenstein_hu(wavenumber, A_s, n_s, cosmology, kwmap=0.02, wiggle=True):
         np.pi**2 / (wavenumber)**3
 
     return power_spectrum
+
+
+class EisensteinHu(PowerSpectrum):
+
+    def __init__(self, A_s, n_s, cosmology, kwmap=0.02, wiggle=True):
+        self.A_s = A_s
+        self.n_s = n_s
+        self.cosmology = cosmology
+        self.kwmap = kwmap
+        self.wiggle = wiggle
+
+    def __call__(self, wavenumber, redshift):
+        growth_function = growth_function_carroll(np.atleast_1d(redshift), self.cosmology)
+        power_spectrum = eisenstein_hu(wavenumber, self.A_s, self.n_s, self.cosmology,
+                                       kwmap=self.kwmap, wiggle=self.wiggle)
+        shape = np.shape(redshift) + np.shape(wavenumber)
+        pzk = (np.square(growth_function)[:, np.newaxis] * power_spectrum).reshape(shape)
+        return pzk.item() if np.isscalar(wavenumber) and np.isscalar(redshift) else pzk


Test patch:
diff --git a/skypy/power_spectrum/tests/test_camb.py b/skypy/power_spectrum/tests/test_camb.py
index 23c8002..fc2762c 100644
--- a/skypy/power_spectrum/tests/test_camb.py
+++ b/skypy/power_spectrum/tests/test_camb.py
@@ -22,18 +22,23 @@ def test_camb():
     '''
     Test a default astropy cosmology
     '''
-    from skypy.power_spectrum import camb
+    from skypy.power_spectrum import CAMB
+
+    # Setup CAMB interpolator
+    k_max, z_grid = 10, np.array([0, 1])
+    A_s, n_s = 2.e-9, 0.965
+    ps = CAMB(k_max, z_grid, Planck15, A_s, n_s)
 
     # test shape and compare with the mocked power spectrum
     redshift = [0.0, 1.0]
     wavenumber = np.logspace(-4.0, np.log10(2.0), 200)
-    pzk = camb(wavenumber, redshift, Planck15, 2.e-9, 0.965)
+    pzk = ps(wavenumber, redshift)
     assert pzk.shape == (len(redshift), len(wavenumber))
     assert allclose(pzk, test_pzk, rtol=1.e-4)
 
     # also check redshifts are ordered correctly
     redshift = [1.0, 0.0]
-    pzk = camb(wavenumber, redshift, Planck15, 2.e-9, 0.965)
+    pzk = ps(wavenumber, redshift)
     assert pzk.shape == (len(redshift), len(wavenumber))
     assert allclose(pzk, test_pzk[np.argsort(redshift), :], rtol=1.e-4)
 
@@ -44,10 +49,16 @@ def test_camb_redshift_zero():
     Regression test for #438
     Test that camb runs succesfully with redshift=0.0
     """
-    from skypy.power_spectrum import camb
+    from skypy.power_spectrum import CAMB
+
+    # Setup CAMB interpolator
+    k_max, z_grid = 10, np.array([0, 1])
+    A_s, n_s = 2.e-9, 0.965
+    ps = CAMB(k_max, z_grid, Planck15, A_s, n_s)
+
     redshift = 0.0
     wavenumber = np.logspace(-4.0, np.log10(2.0), 200)
-    pzk = camb(wavenumber, redshift, Planck15, 2.e-9, 0.965)
+    pzk = ps(wavenumber, redshift)
     assert allclose(pzk, test_pzk[0], rtol=1.e-4)
 
 
diff --git a/skypy/power_spectrum/tests/test_eisenstein_hu.py b/skypy/power_spectrum/tests/test_eisenstein_hu.py
index e080659..25a36ba 100644
--- a/skypy/power_spectrum/tests/test_eisenstein_hu.py
+++ b/skypy/power_spectrum/tests/test_eisenstein_hu.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 from astropy.cosmology import Planck15, FlatLambdaCDM
-from skypy.power_spectrum import eisenstein_hu
+from skypy.power_spectrum import eisenstein_hu, EisensteinHu, growth_function_carroll
 
 
 def test_eisenstein_hu():
@@ -45,6 +45,16 @@ def test_eisenstein_hu():
     assert np.allclose(pk_eisensteinhu_w, pk_pre_w)
     assert np.allclose(pk_eisensteinhu_nw, pk_pre_nw)
 
+    # Also test using the EisensteinHu PowerSpectrum class
+    ps_w = EisensteinHu(A_s, n_s, cosmology, kwmap=kwmap, wiggle=True)
+    ps_nw = EisensteinHu(A_s, n_s, cosmology, kwmap=kwmap, wiggle=False)
+    redshift = 0
+    pk_eisensteinhu_w = ps_w(wavenumber, redshift)
+    pk_eisensteinhu_nw = ps_nw(wavenumber, redshift)
+    growth_function = growth_function_carroll(redshift, cosmology)
+    assert np.allclose(pk_eisensteinhu_w, pk_pre_w * np.square(growth_function))
+    assert np.allclose(pk_eisensteinhu_nw, pk_pre_nw * np.square(growth_function))
+
     # Test for failure when wavenumber <= 0
     negative_wavenumber_scalar = 0
     with pytest.raises(ValueError):