added ptc output option (#763)

* added ptc output option

* added documentation and unit tests

* update object type in output

* flaked

* bug fix

* bug fix, part 2

Co-authored-by: David Alonso <dam.phys@gmail.com>

pyccl/nl_pt/power.py

@@ -381,7 +381,8 @@ def get_pmm(self, Pd1d1_lin, g4):
 def get_pt_pk2d(cosmo, tracer1, tracer2=None, ptc=None,
                 sub_lowk=False, nonlin_pk_type='nonlinear',
                 a_arr=None, extrap_order_lok=1, extrap_order_hik=2,
-                return_ia_bb=False, return_ia_ee_and_bb=False):
+                return_ia_bb=False, return_ia_ee_and_bb=False,
+                return_ptc=False):
     """Returns a :class:`~pyccl.pk2d.Pk2D` object containing
     the PT power spectrum for two quantities defined by
     two :class:`~pyccl.nl_pt.tracers.PTTracer` objects.
@@ -429,9 +430,16 @@ def get_pt_pk2d(cosmo, tracer1, tracer2=None, ptc=None,
             :class:`~pyccl.nl_pt.tracers.PTIntrinsicAlignmentTracer`)
             If `False` (default) E-mode power spectrum is returned.
             Supersedes `return_ia_bb`.
+        return_ptc (bool): if `True`, the fastpt object used as the PT
+            calculator (ptc) will also be returned. This feature may
+            be useful if an input ptc is not specified and one is
+            initialized when this function is called. If `False` (default)
+            the ptc is not output, whether or not it is initialized as
+            part of the function call.
 
     Returns:
         :class:`~pyccl.pk2d.Pk2D`: PT power spectrum.
+        :class:`~pyccl.nl_pt.power.PTCalculator`: PT Calc [optional]
     """
     if a_arr is None:
         status = 0
@@ -581,10 +589,16 @@ def get_pt_pk2d(cosmo, tracer1, tracer2=None, ptc=None,
                         lk_arr=np.log(ptc.ks),
                         pk_arr=p_pt[1].T,
                         is_logp=False)
-        return pt_pk_ee, pt_pk_bb
+        if return_ptc:
+            return pt_pk_ee, pt_pk_bb, ptc
+        else:
+            return pt_pk_ee, pt_pk_bb
     else:
         pt_pk = Pk2D(a_arr=a_arr,
                      lk_arr=np.log(ptc.ks),
                      pk_arr=p_pt.T,
                      is_logp=False)
-        return pt_pk
+        if return_ptc:
+            return pt_pk, ptc
+        else:
+            return pt_pk

pyccl/tests/test_pt_power.py

@@ -37,6 +37,8 @@
 c_2_t = a_2*5*5e-14*rho_crit*COSMO['Omega_m']**2/(Om_m_fid*gz**2)
 c_2_t_des = a_2*5*5e-14*rho_crit*COSMO['Omega_m']/(gz**2)
 
+ks = np.logspace(-3, 2, 512)
+
 
 def test_pt_tracer_smoke():
     ccl.nl_pt.PTTracer()
@@ -233,3 +235,31 @@ def test_translate_IA_norm_raises():
     with pytest.raises(ValueError):
         c_1, c_d, c_2 = ccl.nl_pt.translate_IA_norm(COSMO, ZZ2, a1=a_1_v,
                                                     Om_m2_for_c2=False)
+
+
+def test_return_ptc():
+    # if no ptc is input, check that returned pk and ptc objects have
+    # the correct properties.
+    pk, ptc1 = ccl.nl_pt.get_pt_pk2d(COSMO, TRS['TG'], return_ptc=True)
+    assert isinstance(pk, ccl.Pk2D)
+    assert isinstance(ptc1, ccl.nl_pt.power.PTCalculator)
+    # same test with EE/BB output
+    pee2, pbb2, ptc2 = ccl.nl_pt.get_pt_pk2d(COSMO, TRS['TI'],
+                                             return_ia_ee_and_bb=True,
+                                             return_ptc=True)
+    assert isinstance(pee2, ccl.Pk2D)
+    assert isinstance(pbb2, ccl.Pk2D)
+    assert isinstance(ptc2, ccl.nl_pt.power.PTCalculator)
+    # check that returned ptc matches input ptc.
+    pk_2, ptc1_2 = ccl.nl_pt.get_pt_pk2d(COSMO, TRS['TG'], ptc=PTC,
+                                         return_ptc=True)
+    pee2_2, pbb2_2, ptc2_2 = ccl.nl_pt.get_pt_pk2d(COSMO, TRS['TI'], ptc=PTC,
+                                                   return_ia_ee_and_bb=True,
+                                                   return_ptc=True)
+    assert ptc1_2 is PTC
+    assert ptc2_2 is PTC
+    # check that the result outputs are the same
+    # for the internally initialized ptc.
+    assert np.allclose(pk_2.eval(ks, 1., COSMO), pk.eval(ks, 1., COSMO))
+    assert np.allclose(pee2_2.eval(ks, 1., COSMO), pee2.eval(ks, 1., COSMO))
+    assert np.allclose(pbb2_2.eval(ks, 1., COSMO), pbb2.eval(ks, 1., COSMO))