Enhancements for Pk2D, Tk3D, Tracer. (#923)

* first commit

* 'other' contains 'self'

* 'other' contains 'self'

* added bool dunder method

* correct implementation of non-commutative __dunders__ & simple psp checks

* tidy-up code for optional cosmo in Pk2D.eval

* querying C-level stuff via property

* flaked

* has_tsp property

* tests for all new features

* a few more tests

* straightforward memory optimization

* Pk2D.from_model

* removed spurious cosmo from testing

* *sigh* py37-compatibility loses 1-line elegance

* functools.wraps to preserve docs

* immutable Pk2D

* added bool for Tracer

* moved Tracer.__bool__ to pk2d_extensions

* Tracer chi_min and chi_max properties

* improved docstrings, simplified Pk2D.copy and Pk2D.__call__

* brought back vectorized calls in

* Tracer.[chi_min | chi_max] behavior

pyccl/core.py

@@ -825,7 +825,7 @@ def compute_linear_power(self):
         elif trf in ['bbks', 'eisenstein_hu', 'eisenstein_hu_nowiggles']:
             rescale_s8 = False
             rescale_mg = False
-            pk = Pk2D.pk_from_model(self, model=trf)
+            pk = Pk2D.from_model(self, model=trf)
 
         # Rescale by sigma8/mu-sigma if needed
         if pk:
@@ -934,7 +934,7 @@ def compute_nonlin_power(self):
                                "necessary input for halofit")
             pk = Pk2D.apply_halofit(self, pkl)
         elif mps == 'emu':
-            pk = Pk2D.pk_from_model(self, model='emu')
+            pk = Pk2D.from_model(self, model='emu')
         elif mps == 'linear':
             pk = self._pk_lin['delta_matter:delta_matter']
 

pyccl/parameters.py

@@ -38,7 +38,7 @@ def _new_setattr(self, key, value):
             # raise AttributeError(f"Direct assignment in {name} not supported.")  # noqa
             warnings.warn(
                 f"Direct assignment of {name} is deprecated "
-                "and an error will be raised in the next CCL release."
+                "and an error will be raised in the next CCL release. "
                 f"Set via `pyccl.{name}.{key}` before instantiation.",
                 CCLDeprecationWarning)
             object.__setattr__(self, key, value)

pyccl/tests/test_pk2d.py

@@ -76,7 +76,7 @@ def test_pk2d_from_model(model):
     cosmo = ccl.Cosmology(
         Omega_c=0.27, Omega_b=0.045, h=0.67, sigma8=0.8, n_s=0.96,
         transfer_function=model)
-    pk = ccl.Pk2D.pk_from_model(cosmo_fixed, model=model)
+    pk = ccl.Pk2D.from_model(cosmo_fixed, model=model)
     ks = np.geomspace(1E-3, 1E1, 128)
     for z in [0., 0.5, 2.]:
         a = 1./(1+z)
@@ -111,7 +111,7 @@ def test_pk2d_from_model_emu():
                           Omega_k=0,
                           transfer_function='bbks',
                           matter_power_spectrum='emu')
-    pk = ccl.Pk2D.pk_from_model(cosmo_fixed, model='emu')
+    pk = ccl.Pk2D.from_model(cosmo_fixed, model='emu')
     ks = np.geomspace(1E-3, 1E1, 128)
     for z in [0., 0.5, 2.]:
         a = 1./(1+z)
@@ -126,13 +126,13 @@ def test_pk2d_from_model_fails(model):
     cosmo = ccl.Cosmology(
         Omega_c=0.27, Omega_b=0.045, h=0.67, A_s=1E-10, n_s=0.96,
         transfer_function='boltzmann_class')
-    assert_raises(ccl.CCLError, ccl.Pk2D.pk_from_model,
+    assert_raises(ccl.CCLError, ccl.Pk2D.from_model,
                   cosmo, model=model)
 
 
 def test_pk2d_from_model_raises():
     cosmo = ccl.CosmologyVanillaLCDM()
-    assert_raises(ValueError, ccl.Pk2D.pk_from_model,
+    assert_raises(ValueError, ccl.Pk2D.from_model,
                   cosmo, model='bbkss')
 
 
@@ -368,3 +368,101 @@ def test_pk2d_pkfunc_init_without_cosmo():
     arr1 = ccl.Pk2D(pkfunc=lpk2d, cosmo=cosmo).get_spline_arrays()[-1]
     arr2 = ccl.Pk2D(pkfunc=lpk2d).get_spline_arrays()[-1]
     assert np.allclose(arr1, arr2, rtol=0)
+
+
+def test_pk2d_extrap_orders():
+    # Check that setting extrap orders propagates down to the `psp`.
+    x = np.linspace(0.1, 1, 10)
+    log_y = np.linspace(-3, 1, 20)
+    zarr_a = np.outer(x, np.exp(log_y))
+    pk = ccl.Pk2D(a_arr=x, lk_arr=log_y, pk_arr=np.log(zarr_a), is_logp=True)
+
+    assert pk.extrap_order_hik == pk.psp.extrap_order_hik
+    assert pk.extrap_order_lok == pk.psp.extrap_order_lok
+
+
+def test_pk2d_descriptor():
+    # Check that `apply_halofit` can be called as a class method or
+    # as an instance method.
+    cosmo = ccl.CosmologyVanillaLCDM(transfer_function="bbks")
+    cosmo.compute_linear_power()
+    pkl = cosmo.get_linear_power()
+    pk1 = ccl.Pk2D.apply_halofit(cosmo, pk_linear=pkl)
+    pk2 = pkl.apply_halofit(cosmo)
+    assert np.all(pk1.get_spline_arrays()[-1] == pk2.get_spline_arrays()[-1])
+
+
+def test_pk2d_eval_cosmo():
+    # Check that `eval` can be called without `cosmo` and that an error
+    # is raised when scale factor is out of interpolation range.
+    cosmo = ccl.CosmologyVanillaLCDM(transfer_function="bbks")
+    cosmo.compute_linear_power()
+    pk = cosmo.get_linear_power()
+    assert pk.eval(1., 1.) == pk.eval(1., 1., cosmo)
+
+    amin = pk.psp.amin
+    pk.eval(1., amin*0.99, cosmo)  # doesn't fail because cosmo is provided
+    with pytest.raises(TypeError):
+        pk.eval(1., amin*0.99)
+
+
+def test_pk2d_copy():
+    # Check that copying works as intended (also check `bool`).
+    x = np.linspace(0.1, 1, 10)
+    log_y = np.linspace(-3, 1, 20)
+    zarr_a = np.outer(x, np.exp(log_y))
+    pk = ccl.Pk2D(a_arr=x, lk_arr=log_y, pk_arr=np.log(zarr_a), is_logp=True)
+
+    pkc = pk.copy()
+    assert np.allclose(pk.get_spline_arrays()[-1],
+                       pkc.get_spline_arrays()[-1],
+                       rtol=1e-15)
+    assert bool(pk) is bool(pkc) is True  # they both have `psp`
+
+    pk = ccl.Pk2D(empty=True)
+    pkc = pk.copy()
+    assert bool(pk) is bool(pkc) is False
+
+
+def test_pk2d_operations():
+    # Everything is based on the already tested `add`, `mul`, and `pow`,
+    # so we don't need to test every accepted type separately.
+    x = np.linspace(0.1, 1, 10)
+    log_y = np.linspace(-3, 1, 20)
+    zarr_a = np.outer(x, np.exp(log_y))
+    pk0 = ccl.Pk2D(a_arr=x, lk_arr=log_y, pk_arr=np.log(zarr_a), is_logp=True)
+    pk1, pk2 = pk0.copy(), pk0.copy()
+
+    # sub, truediv
+    assert np.allclose((pk1 - pk2).get_spline_arrays()[-1], 0, rtol=1e-15)
+    assert np.allclose((pk1 / pk2).get_spline_arrays()[-1], 1, rtol=1e-15)
+
+    # rsub, rtruediv
+    assert np.allclose((1 - pk1).get_spline_arrays()[-1],
+                       1 - pk1.get_spline_arrays()[-1])
+    assert np.allclose((1 / pk1).get_spline_arrays()[-1],
+                       1 / pk1.get_spline_arrays()[-1])
+
+    # iadd, isub, imul, itruediv, ipow
+    pk1 += pk1
+    assert np.allclose((pk1 / pk2).get_spline_arrays()[-1], 2, rtol=1e-15)
+    pk1 -= pk2
+    assert np.allclose((pk1 / pk2).get_spline_arrays()[-1], 1, rtol=1e-15)
+    pk1 *= pk1
+    assert np.allclose(pk1.get_spline_arrays()[-1],
+                       pk2.get_spline_arrays()[-1]**2,
+                       rtol=1e-15)
+    pk1 /= pk2
+    assert np.allclose((pk1 / pk2).get_spline_arrays()[-1], 1, rtol=1e-15)
+    pk1 **= 2
+    assert np.allclose(pk1.get_spline_arrays()[-1],
+                       pk2.get_spline_arrays()[-1]**2,
+                       rtol=1e-15)
+
+
+def test_pk2d_from_model_smoke():
+    # Verify that both `from_model` methods are equivalent.
+    cosmo = ccl.CosmologyVanillaLCDM(transfer_function="bbks")
+    pk1 = ccl.Pk2D.from_model(cosmo, "bbks")
+    pk2 = ccl.Pk2D.pk_from_model(cosmo, "bbks")
+    assert np.all(pk1.get_spline_arrays()[-1] == pk2.get_spline_arrays()[-1])

pyccl/tests/test_tk3d.py

@@ -139,6 +139,15 @@ def test_tk3d_eval(is_product):
     assert_allclose(phere.flatten(), ptrue.flatten(), rtol=1E-6)
 
 
+def test_tk3d_call():
+    # Test `__call__` and `__bool__`
+    (a_arr, lk_arr, fka1_arr, fka2_arr, tkka_arr) = get_arrays()
+    tsp = ccl.Tk3D(a_arr, lk_arr, tkk_arr=tkka_arr)
+    assert bool(tsp) is tsp.has_tsp
+    assert np.allclose(np.array([tsp.eval(np.exp(lk_arr), a) for a in a_arr]),
+                       tsp(np.exp(lk_arr), a_arr), rtol=1e-15)
+
+
 @pytest.mark.parametrize('is_product', [True, False])
 def test_tk3d_spline_arrays(is_product):
     (a_arr, lk_arr, fka1_arr, fka2_arr, tkka_arr) = get_arrays()
@@ -163,14 +172,8 @@ def test_tk3d_spline_arrays_raises():
     (a_arr, lk_arr, fka1_arr, fka2_arr, tkka_arr) = get_arrays()
     tsp = ccl.Tk3D(a_arr, lk_arr, tkk_arr=tkka_arr)
 
-    # PR923 aims to change this bit of code; the assertion is there to remind
-    # us to uncomment what is commented out.
-    assert not hasattr(tsp.__class__, "has_tsp")
-    # ccl.lib.f3d_t_free(tsp.tsp)
-    # delattr(tsp, "tsp")
-
-    # fool `Tk3D` into believing it doesn't have a `tsp`
-    tsp.has_tsp = False
+    ccl.lib.f3d_t_free(tsp.tsp)
+    delattr(tsp, "tsp")
 
     with pytest.raises(ValueError):
         tsp.get_spline_arrays()

pyccl/tests/test_tracers.py

@@ -351,3 +351,22 @@ def test_tracer_n_sample_warn():
 
     with pytest.warns(CCLWarning):
         _ = ccl.WeakLensingTracer(COSMO, dndz=(z, n))
+
+
+def test_tracer_bool():
+    assert bool(ccl.Tracer()) is False
+    assert bool(ccl.CMBLensingTracer(COSMO, z_source=1100)) is True
+
+
+def test_tracer_chi_min_max():
+    # Test that it can access the C-level chi_min and chi_max.
+    tr = ccl.CMBLensingTracer(COSMO, z_source=1100)
+    assert tr.chi_min == tr._trc[0].chi_min
+    assert tr.chi_max == tr._trc[0].chi_max
+
+    # Returns the lowest/highest if chi_min or chi_max are not the same.
+    chi = np.linspace(tr.chi_min+0.05, tr.chi_max+0.05, 128)
+    wchi = np.ones_like(chi)
+    tr.add_tracer(COSMO, kernel=(chi, wchi))
+    assert tr.chi_min == tr._trc[0].chi_min
+    assert tr.chi_max == tr._trc[1].chi_max

pyccl/tk3d.py

@@ -130,7 +130,10 @@ def __init__(self, a_arr, lk_arr, tkk_arr=None,
                                                         int(extrap_order_lok),
                                                         int(is_logt), status)
         check(status)
-        self.has_tsp = True
+
+    @property
+    def has_tsp(self):
+        return 'tsp' in vars(self)
 
     def eval(self, k, a):
         """Evaluate trispectrum. If `k` is a 1D array with size `nk`, the
@@ -163,11 +166,20 @@ def eval(self, k, a):
         check(status)
         return f
 
+    def __call__(self, k, a):
+        """Callable vectorized instance."""
+        out = np.array([self.eval(k, aa)
+                        for aa in np.atleast_1d(a).astype(float)])
+        return out.squeeze()[()]
+
     def __del__(self):
         if hasattr(self, 'has_tsp'):
             if self.has_tsp and hasattr(self, 'tsp'):
                 lib.f3d_t_free(self.tsp)
 
+    def __bool__(self):
+        return self.has_tsp
+
     def get_spline_arrays(self):
         """Get the spline data arrays.
 

pyccl/tracers.py

@@ -182,6 +182,27 @@ def __init__(self):
         # Do nothing, just initialize list of tracers
         self._trc = []
 
+    def __bool__(self):
+        return bool(self._trc)
+
+    @property
+    def chi_min(self):
+        """Return ``chi_min`` for this ``Tracer``, if it exists. For more than
+        one tracers containing a ``chi_min`` in the tracer collection, the
+        lowest value is returned.
+        """
+        chis = [tr.chi_min for tr in self._trc]
+        return min(chis) if chis else None
+
+    @property
+    def chi_max(self):
+        """Return ``chi_max`` for this ``Tracer``, if it exists. For more than
+        one tracers containing a ``chi_max`` in the tracer collection, the
+        highest value is returned.
+        """
+        chis = [tr.chi_max for tr in self._trc]
+        return max(chis) if chis else None
+
     def _dndz(self, z):
         raise NotImplementedError("`get_dndz` not implemented for "
                                   "this `Tracer` type.")
@@ -749,8 +770,8 @@ class tSZTracer(Tracer):
             distance on which we sample the kernel.
     """
     def __init__(self, cosmo, z_max=6., n_chi=1024):
-        self.chi_max = comoving_radial_distance(cosmo, 1./(1+z_max))
-        chi_arr = np.linspace(0, self.chi_max, n_chi)
+        chi_max = comoving_radial_distance(cosmo, 1./(1+z_max))
+        chi_arr = np.linspace(0, chi_max, n_chi)
         a_arr = scale_factor_of_chi(cosmo, chi_arr)
         # This is \sigma_T / (m_e * c^2)
         prefac = 4.01710079e-06
@@ -783,9 +804,9 @@ class CIBTracer(Tracer):
             distance on which we sample the kernel.
     """
     def __init__(self, cosmo, z_min=0., z_max=6., n_chi=1024):
-        self.chi_max = comoving_radial_distance(cosmo, 1./(1+z_max))
-        self.chi_min = comoving_radial_distance(cosmo, 1./(1+z_min))
-        chi_arr = np.linspace(self.chi_min, self.chi_max, n_chi)
+        chi_max = comoving_radial_distance(cosmo, 1./(1+z_max))
+        chi_min = comoving_radial_distance(cosmo, 1./(1+z_min))
+        chi_arr = np.linspace(chi_min, chi_max, n_chi)
         a_arr = scale_factor_of_chi(cosmo, chi_arr)
 
         self._trc = []
@@ -818,8 +839,8 @@ class ISWTracer(Tracer):
             distance on which we sample the kernel.
     """
     def __init__(self, cosmo, z_max=6., n_chi=1024):
-        self.chi_max = comoving_radial_distance(cosmo, 1./(1+z_max))
-        chi = np.linspace(0, self.chi_max, n_chi)
+        chi_max = comoving_radial_distance(cosmo, 1./(1+z_max))
+        chi = np.linspace(0, chi_max, n_chi)
         a_arr = scale_factor_of_chi(cosmo, chi)
         H0 = cosmo['h'] / physical_constants.CLIGHT_HMPC
         OM = cosmo['Omega_c']+cosmo['Omega_b']