__add__, __mul__, __pow__ operators for Pk2D (#894)

* Allow addition of two Pk2D objects

* Make flake8 happy.

* Added __mul__ and __pow__ operators.

* Test for operator composition.

* Better test coverage

* pacify flake8

* Use TypeErrors, allow addition of constant, warn about powers.

* Added test for warning.

* More verbose error message.

* Added warning when interpolation is happening.

* More verbose warning.

pyccl/pk2d.py

@@ -1,8 +1,12 @@
-from . import ccllib as lib
+import warnings
 
-from .pyutils import check
 import numpy as np
 
+from . import ccllib as lib
+
+from .errors import CCLWarning
+from .pyutils import check, _get_spline2d_arrays
+
 
 class Pk2D(object):
     """A power spectrum class holding the information needed to reconstruct an
@@ -63,6 +67,7 @@ def __init__(self, pkfunc=None, a_arr=None, lk_arr=None, pk_arr=None,
                  is_logp=True, extrap_order_lok=1, extrap_order_hik=2,
                  cosmo=None, empty=False):
         if empty:
+            self.has_psp = False
             return
 
         status = 0
@@ -106,6 +111,9 @@ def __init__(self, pkfunc=None, a_arr=None, lk_arr=None, pk_arr=None,
                 pkflat[ia, :] = pkfunc(k=np.exp(lk_arr), a=a)
             pkflat = pkflat.flatten()
 
+        self.extrap_order_lok = extrap_order_lok
+        self.extrap_order_hik = extrap_order_hik
+
         self.psp, status = lib.set_pk2d_new_from_arrays(lk_arr, a_arr, pkflat,
                                                         int(extrap_order_lok),
                                                         int(extrap_order_hik),
@@ -252,13 +260,171 @@ def eval_dlogpk_dlogk(self, k, a, cosmo):
 
         return f
 
+    def get_spline_arrays(self):
+        """Get the spline data arrays.
+
+        Returns:
+            a_arr: array_like
+                Array of scale factors.
+            lk_arr: array_like
+                Array of logarithm of wavenumber k.
+            pk_arr: array_like
+                Array of the power spectrum P(k, z). The shape
+                is (a_arr.size, lk_arr.size).
+        """
+        if not self.has_psp:
+            raise ValueError("Pk2D object does not have data.")
+
+        a_arr, lk_arr, pk_arr = _get_spline2d_arrays(self.psp.fka)
+        if self.psp.is_log:
+            pk_arr = np.exp(pk_arr)
+
+        return a_arr, lk_arr, pk_arr
+
     def __del__(self):
         """Free memory associated with this Pk2D structure
         """
         if hasattr(self, 'has_psp'):
             if self.has_psp and hasattr(self, 'psp'):
                 lib.f2d_t_free(self.psp)
 
+    def _get_binary_operator_arrays(self, other):
+        if not isinstance(other, Pk2D):
+            raise TypeError("Binary operator of Pk2D objects is only defined "
+                            "for other Pk2D objects.")
+        if not (self.has_psp and other.has_psp):
+            raise ValueError("Pk2D object does not have data.")
+        if (self.psp.lkmin < other.psp.lkmin
+                or self.psp.lkmax > other.psp.lkmax):
+            raise ValueError("The 2nd operand has its data defined over a "
+                             "smaller k range than the 1st operand. To avoid "
+                             "extrapolation, this operation is forbidden. If "
+                             "you want to operate on the smaller support, "
+                             "try swapping the operands.")
+        if (self.psp.amin < other.psp.amin
+                or self.psp.amax > other.psp.amax):
+            raise ValueError("The 2nd operand has its data defined over a "
+                             "smaller a range than the 1st operand. To avoid "
+                             "extrapolation, this operation is forbidden. If "
+                             "you want to operate on the smaller support, "
+                             "try swapping the operands.")
+
+        a_arr_a, lk_arr_a, pk_arr_a = self.get_spline_arrays()
+        a_arr_b, lk_arr_b, pk_arr_b = other.get_spline_arrays()
+        if not (a_arr_a.size == a_arr_b.size and lk_arr_a.size == lk_arr_b.size
+                and np.allclose(a_arr_a, a_arr_b)
+                and np.allclose(lk_arr_a, lk_arr_b)):
+            warnings.warn("The arrays of the two Pk2D objects are defined at "
+                          "different points in k and/or a. The second operand "
+                          "will be interpolated for the operation.\n"
+                          "The resulting Pk2D object will be defined for "
+                          f"{self.psp.lkmin} <= log k <= {self.psp.lkmax} and "
+                          f"{self.psp.amin} <= a <= {self.psp.amax}.",
+                          category=CCLWarning)
+
+            # Since the power spectrum is evalulated on a smaller support than
+            # where it was defined, no extrapolation is necessary and the
+            # dependence on the cosmology in moot.
+            # CosmologyVanillaLCDM is being imported here instead of the top of
+            # the module due to circular import issues there.
+            from .core import CosmologyVanillaLCDM
+            dummy_cosmo = CosmologyVanillaLCDM()
+            pk_arr_b = np.array([other.eval(k=np.exp(lk_arr_a),
+                                            a=a_,
+                                            cosmo=dummy_cosmo)
+                                 for a_ in a_arr_a])
+        return a_arr_a, lk_arr_a, pk_arr_a, pk_arr_b
+
+    def __add__(self, other):
+        """Adds two Pk2D instances.
+
+        The a and k ranges of the 2nd operand need to be the same or smaller
+        than the 1st operand.
+        The returned Pk2D object uses the same a and k arrays as the first
+        operand.
+        """
+        if isinstance(other, (float, int)):
+            a_arr_a, lk_arr_a, pk_arr_a = self.get_spline_arrays()
+            pk_arr_new = pk_arr_a + other
+        elif isinstance(other, Pk2D):
+            a_arr_a, lk_arr_a, pk_arr_a, pk_arr_b = \
+                self._get_binary_operator_arrays(other)
+            pk_arr_new = pk_arr_a + pk_arr_b
+        else:
+            raise TypeError("Addition of Pk2D is only defined for "
+                            "floats, ints, and Pk2D objects.")
+
+        logp = np.all(pk_arr_new > 0)
+        if logp:
+            pk_arr_new = np.log(pk_arr_new)
+
+        new = Pk2D(a_arr=a_arr_a, lk_arr=lk_arr_a, pk_arr=pk_arr_new,
+                   is_logp=logp,
+                   extrap_order_lok=self.extrap_order_lok,
+                   extrap_order_hik=self.extrap_order_hik)
+
+        return new
+
+    def __radd__(self, other):
+        return self.__add__(other)
+
+    def __mul__(self, other):
+        """Multiply two Pk2D instances.
+
+        The a and k ranges of the 2nd operand need to be the same or smaller
+        than the 1st operand.
+        The returned Pk2D object uses the same a and k arrays as the first
+        operand.
+        """
+        if isinstance(other, (float, int)):
+            a_arr_a, lk_arr_a, pk_arr_a = self.get_spline_arrays()
+            pk_arr_new = other * pk_arr_a
+        elif isinstance(other, Pk2D):
+            a_arr_a, lk_arr_a, pk_arr_a, pk_arr_b = \
+                self._get_binary_operator_arrays(other)
+            pk_arr_new = pk_arr_a * pk_arr_b
+        else:
+            raise TypeError("Multiplication of Pk2D is only defined for "
+                            "floats, ints, and Pk2D objects.")
+
+        logp = np.all(pk_arr_new > 0)
+        if logp:
+            pk_arr_new = np.log(pk_arr_new)
+
+        new = Pk2D(a_arr=a_arr_a, lk_arr=lk_arr_a, pk_arr=pk_arr_new,
+                   is_logp=logp,
+                   extrap_order_lok=self.extrap_order_lok,
+                   extrap_order_hik=self.extrap_order_hik)
+        return new
+
+    def __rmul__(self, other):
+        return self.__mul__(other)
+
+    def __pow__(self, exponent):
+        """Take a Pk2D instance to a power.
+        """
+        if not isinstance(exponent, (float, int)):
+            raise TypeError("Exponentiation of Pk2D is only defined for "
+                            "floats and ints.")
+        a_arr_a, lk_arr_a, pk_arr_a = self.get_spline_arrays()
+        if np.any(pk_arr_a < 0) and exponent % 1 != 0:
+            warnings.warn("Taking a non-positive Pk2D object to a non-integer "
+                          "power may lead to unexpected results",
+                          category=CCLWarning)
+
+        pk_arr_new = pk_arr_a**exponent
+
+        logp = np.all(pk_arr_new > 0)
+        if logp:
+            pk_arr_new = np.log(pk_arr_new)
+
+        new = Pk2D(a_arr=a_arr_a, lk_arr=lk_arr_a, pk_arr=pk_arr_new,
+                   is_logp=logp,
+                   extrap_order_lok=self.extrap_order_lok,
+                   extrap_order_hik=self.extrap_order_hik)
+
+        return new
+
 
 def parse_pk2d(cosmo, p_of_k_a, is_linear=False):
     """ Return the C-level `f2d` spline associated with a

pyccl/tests/test_pk2d.py

@@ -4,6 +4,7 @@
     assert_,
     assert_raises, assert_almost_equal, assert_allclose)
 import pyccl as ccl
+from pyccl import CCLWarning
 
 
 def pk1d(k):
@@ -261,3 +262,105 @@ def test_pk2d_parsing():
     with pytest.raises(ValueError):
         ccl.angular_cl(cosmo, lens1, lens1, ells,
                        p_of_k_a=3)
+
+
+def test_pk2d_get_spline_arrays():
+    empty_pk2d = ccl.Pk2D(empty=True)
+
+    # Pk2D needs splines defined to get splines out
+    with pytest.raises(ValueError):
+        empty_pk2d.get_spline_arrays()
+
+
+def test_pk2d_add():
+    x = np.linspace(0.1, 1, 10)
+    log_y = np.linspace(-3, 1, 20)
+    zarr_a = np.outer(x, np.exp(log_y))
+    zarr_b = np.outer(-1*x, 4*np.exp(log_y))
+
+    empty_pk2d = ccl.Pk2D(empty=True)
+    pk2d_a = ccl.Pk2D(a_arr=x, lk_arr=log_y, pk_arr=np.log(zarr_a),
+                      is_logp=True)
+    pk2d_b = ccl.Pk2D(a_arr=2*x, lk_arr=log_y, pk_arr=zarr_b,
+                      is_logp=False)
+    pk2d_b2 = ccl.Pk2D(a_arr=x, lk_arr=log_y+0.5, pk_arr=zarr_b,
+                       is_logp=False)
+
+    # This raises an error because the a ranges don't match
+    with pytest.raises(ValueError):
+        pk2d_a + pk2d_b
+    # This raises an error because the k ranges don't match
+    with pytest.raises(ValueError):
+        pk2d_a + pk2d_b2
+    # This raises an error because addition with an empty Pk2D should not work
+    with pytest.raises(ValueError):
+        pk2d_a + empty_pk2d
+
+    pk2d_c = ccl.Pk2D(a_arr=x, lk_arr=log_y, pk_arr=zarr_b,
+                      is_logp=False)
+
+    pk2d_d = pk2d_a + pk2d_c
+    pk2d_d2 = pk2d_a + 1.0
+    xarr_d, yarr_d, zarr_d = pk2d_d.get_spline_arrays()
+    _, _, zarr_d2 = pk2d_d2.get_spline_arrays()
+
+    assert np.allclose(x, xarr_d)
+    assert np.allclose(log_y, yarr_d)
+    assert np.allclose(zarr_a + zarr_b, zarr_d)
+    assert np.allclose(zarr_a + 1.0, zarr_d2)
+
+    pk2d_e = ccl.Pk2D(a_arr=x[1:-1], lk_arr=log_y[1:-1],
+                      pk_arr=zarr_b[1:-1, 1:-1],
+                      is_logp=False)
+
+    # This raises a warning because the power spectra are not defined on the
+    # same support
+    with pytest.warns(CCLWarning):
+        pk2d_f = pk2d_e + pk2d_a
+
+    xarr_f, yarr_f, zarr_f = pk2d_f.get_spline_arrays()
+
+    assert np.allclose((zarr_a + zarr_b)[1:-1, 1:-1], zarr_f)
+
+
+def test_pk2d_mul_pow():
+    x = np.linspace(0.1, 1, 10)
+    log_y = np.linspace(-3, 1, 20)
+    zarr_a = np.outer(x, np.exp(log_y))
+    zarr_b = np.outer(-1*x, 4*np.exp(log_y))
+
+    pk2d_a = ccl.Pk2D(a_arr=x, lk_arr=log_y, pk_arr=np.log(zarr_a),
+                      is_logp=True)
+    pk2d_b = ccl.Pk2D(a_arr=x, lk_arr=log_y, pk_arr=zarr_b,
+                      is_logp=False)
+
+    # This raises an error because multiplication is only defined for
+    # float, int, and Pk2D
+    with pytest.raises(TypeError):
+        pk2d_a*np.array([0.1, 0.2])
+
+    # This raises an error because exponention is only defined for
+    # float and int
+    with pytest.raises(TypeError):
+        pk2d_a**pk2d_b
+
+    # This raises a warning because the power spectrum is non-negative and the
+    # power is non-integer
+    with pytest.warns(CCLWarning):
+        pk2d_b**0.5
+
+    pk2d_g = pk2d_a * pk2d_b
+    pk2d_h = 2*pk2d_a
+    pk2d_i = pk2d_a**1.8
+
+    _, _, zarr_g = pk2d_g.get_spline_arrays()
+    _, _, zarr_h = pk2d_h.get_spline_arrays()
+    _, _, zarr_i = pk2d_i.get_spline_arrays()
+
+    assert np.allclose(zarr_a * zarr_b, zarr_g)
+    assert np.allclose(2 * zarr_a, zarr_h)
+    assert np.allclose(zarr_a**1.8, zarr_i)
+
+    pk2d_j = (pk2d_a + 0.5*pk2d_i)**1.5
+    _, _, zarr_j = pk2d_j.get_spline_arrays()
+    assert np.allclose((zarr_a + 0.5*zarr_i)**1.5, zarr_j)