Ia der term

pyccl/nl_pt/ept.py

@@ -13,19 +13,20 @@
     'm:m': 'm:m', 'm:b1': 'm:m', 'm:b2': 'm:b2',
     'm:b3nl': 'm:b3nl', 'm:bs': 'm:bs', 'm:bk2': 'm:bk2',
     'm:c1': 'm:m', 'm:c2': 'm:c2', 'm:cdelta': 'm:cdelta',
-    'b1:b1': 'm:m', 'b1:b2': 'm:b2', 'b1:b3nl': 'm:b3nl',
-    'b1:bs': 'm:bs', 'b1:bk2': 'm:bk2', 'b1:c1': 'm:m',
-    'b1:c2': 'm:c2', 'b1:cdelta': 'm:cdelta', 'b2:b2': 'b2:b2',
-    'b2:b3nl': 'zero', 'b2:bs': 'b2:bs', 'b2:bk2': 'zero',
-    'b2:c1': 'zero', 'b2:c2': 'zero', 'b2:cdelta': 'zero',
-    'b3nl:b3nl': 'zero', 'b3nl:bs': 'zero',
+    'm:ck': 'm:ck', 'b1:b1': 'm:m', 'b1:b2': 'm:b2',
+    'b1:b3nl': 'm:b3nl', 'b1:bs': 'm:bs', 'b1:bk2': 'm:bk2',
+    'b1:c1': 'm:m', 'b1:c2': 'm:c2', 'b1:cdelta': 'm:cdelta', 'b1:ck': 'm:ck',
+    'b2:b2': 'b2:b2', 'b2:b3nl': 'zero', 'b2:bs': 'b2:bs',
+    'b2:bk2': 'zero', 'b2:c1': 'zero', 'b2:c2': 'zero',
+    'b2:cdelta': 'zero', 'b3nl:b3nl': 'zero', 'b3nl:bs': 'zero',
     'b3nl:bk2': 'zero', 'b3nl:c1': 'zero', 'b3nl:c2':
     'zero', 'b3nl:cdelta': 'zero', 'bs:bs': 'bs:bs',
     'bs:bk2': 'zero', 'bs:c1': 'zero', 'bs:c2': 'zero',
     'bs:cdelta': 'zero', 'bk2:bk2': 'zero', 'bk2:c1': 'zero',
     'bk2:c2': 'zero', 'bk2:cdelta': 'zero', 'c1:c1': 'm:m',
-    'c1:c2': 'm:c2', 'c1:cdelta': 'm:cdelta', 'c2:c2': 'c2:c2',
-    'c2:cdelta': 'c2:cdelta', 'cdelta:cdelta': 'cdelta:cdelta'}
+    'c1:c2': 'm:c2', 'c1:cdelta': 'm:cdelta', 'c1:ck': 'm:ck',
+    'c2:c2': 'c2:c2', 'c2:cdelta': 'c2:cdelta',
+    'cdelta:cdelta': 'cdelta:cdelta', 'ck:ck': 'zero'}
 
 
 class EulerianPTCalculator(CCLAutoRepr):
@@ -47,7 +48,7 @@ class EulerianPTCalculator(CCLAutoRepr):
 
     .. math::
         s^I_{ij}=c_1\\,s_{ij}+c_2(s_{ik}s_{jk}-s^2\\delta_{ik}/3)
-        +c_\\delta\\,\\delta\\,s_{ij}
+        +c_\\delta\\,\\delta\\,s_{ij + c_k\\k^2\\,s_{ij}}
 
     (note that the higher-order terms are not divided by 2!).
 
@@ -115,6 +116,10 @@ class EulerianPTCalculator(CCLAutoRepr):
         bk2_pk_kind (:obj:`str`): power spectrum to use for the non-local
             bias terms in the expansion. Same options and default as
             ``b1_pk_kind``.
+        ak2_pk_kind (:obj:`str`): power spectrum to use for the derivative
+            term of the IA expansion. Same options and default as
+            ``b1_pk_kind``. Additionally, users may pass ``'fastpt'`` to
+            use the FAST-PT version.
         pad_factor (:obj:`float`): fraction of the :math:`\\log_{10}(k)`
              interval you to add as padding for FFTLog calculations.
         low_extrap (:obj:`float`): decimal logaritm of the minimum Fourier
@@ -141,12 +146,13 @@ def __init__(self, *, with_NC=False, with_IA=False,
                  log10k_min=-4, log10k_max=2, nk_per_decade=20,
                  a_arr=None, k_cutoff=None, n_exp_cutoff=4,
                  b1_pk_kind='nonlinear', bk2_pk_kind='nonlinear',
+                 ak2_pk_kind='nonlinear',
                  pad_factor=1.0, low_extrap=-5.0, high_extrap=3.0,
                  P_window=None, C_window=0.75, sub_lowk=False):
         self.with_matter_1loop = with_matter_1loop
         self.with_NC = with_NC
         self.with_IA = with_IA
-
+        self._ufpt_ak = ak2_pk_kind == 'fastpt'
         # Set FAST-PT parameters
         self.fastpt_par = {'pad_factor': pad_factor,
                            'low_extrap': low_extrap,
@@ -179,7 +185,25 @@ def __init__(self, *, with_NC=False, with_IA=False,
             self.exp_cutoff = 1
 
         # Call FAST-PT
-        import fastpt as fpt
+        try:
+            import fastpt as fpt
+        except Exception:
+            raise ImportError("Your attempted import of FAST-PT has failed. "
+                              "You either dont have fast-pt installed, "
+                              "or have the wrong version. Try running "
+                              "pip install fast-pt  or conda "
+                              "install fast-pt, then try again")
+
+        if (not hasattr(fpt, "IA_ta")):
+            raise ValueError("Your FAST-PT version lacks a required function. "
+                             "You may have the wrong fast-pt install. "
+                             "Try running pip install "
+                             "fast-pt or conda install fast-pt, "
+                             "then try again")
+
+        if (not hasattr(fpt, "IA_der") and self._ufpt_ak):
+            raise ValueError("You need a newer version of "
+                             "FAST-PT to use fpt for k2 term")
         n_pad = int(self.fastpt_par['pad_factor'] * len(self.k_s))
         self.pt = fpt.FASTPT(self.k_s, to_do=to_do,
                              low_extrap=self.fastpt_par['low_extrap'],
@@ -191,8 +215,11 @@ def __init__(self, *, with_NC=False, with_IA=False,
             raise ValueError(f"Unknown P(k) prescription {b1_pk_kind}")
         if bk2_pk_kind not in ['linear', 'nonlinear', 'pt']:
             raise ValueError(f"Unknown P(k) prescription {bk2_pk_kind}")
+        if ak2_pk_kind not in ['linear', 'nonlinear', 'pt']:
+            raise ValueError(f"Unknown P(k) prescription in {ak2_pk_kind}")
         self.b1_pk_kind = b1_pk_kind
         self.bk2_pk_kind = bk2_pk_kind
+        self.ak2_pk_kind = ak2_pk_kind
         if (self.b1_pk_kind == 'pt') or (self.bk2_pk_kind == 'pt'):
             self.with_matter_1loop = True
 
@@ -260,16 +287,20 @@ def reshape_fastpt(tupl):
             reshape_fastpt(self.ia_tt)
             self.ia_mix = self.pt.IA_mix(**kw)
             reshape_fastpt(self.ia_mix)
+            if (self._ufpt_ak):
+                self.ia_der = self.pt.IA_der(**kw)
+                reshape_fastpt(self.ia_der)
 
-        # b1/bk power spectrum
+        # b1/bk/ak power spectrum
         pks = {}
-        if 'nonlinear' in [self.b1_pk_kind, self.bk2_pk_kind]:
+        kinds = [self.b1_pk_kind, self.bk2_pk_kind, self.ak2_pk_kind]
+        if 'nonlinear' in kinds:
             pks['nonlinear'] = np.array([cosmo.nonlin_matter_power(self.k_s, a)
                                          for a in self.a_s])
-        if 'linear' in [self.b1_pk_kind, self.bk2_pk_kind]:
+        if 'linear' in kinds:
             pks['linear'] = np.array([cosmo.linear_matter_power(self.k_s, a)
                                       for a in self.a_s])
-        if 'pt' in [self.b1_pk_kind, self.bk2_pk_kind]:
+        if 'pt' in kinds:
             if 'linear' in pks:
                 pk = pks['linear']
             else:
@@ -280,6 +311,7 @@ def reshape_fastpt(tupl):
             pks['pt'] = pk
         self.pk_b1 = pks[self.b1_pk_kind]
         self.pk_bk = pks[self.bk2_pk_kind]
+        self.pk_ak = pks[self.ak2_pk_kind]
 
         # Reset template power spectra
         self._pk2d_temp = {}
@@ -364,6 +396,10 @@ def _get_pgi(self, trg, tri):
         Pd1d1 = self.pk_b1
         a00e, c00e, a0e0e, a0b0b = self.ia_ta
         a0e2, b0e2, d0ee2, d0bb2 = self.ia_mix
+        if (self._ufpt_ak):
+            Pak2 = self.ia_der
+        else:
+            Pak2 = self.pk_ak*(self.k_s**2)[None, :]
 
         # Get biases
         b1 = trg.b1(self.z_s)
@@ -379,10 +415,12 @@ def _get_pgi(self, trg, tri):
         c1 = tri.c1(self.z_s)
         c2 = tri.c2(self.z_s)
         cd = tri.cdelta(self.z_s)
+        ck = tri.ck(self.z_s)
 
         pgi = b1[:, None] * (c1[:, None] * Pd1d1 +
                              (self._g4*cd)[:, None] * (a00e + c00e) +
-                             (self._g4*c2)[:, None] * (a0e2 + b0e2))
+                             (self._g4*c2)[:, None] * (a0e2 + b0e2) +
+                             ck[:, None] * Pak2)
         return pgi*self.exp_cutoff
 
     def _get_pgm(self, trg):
@@ -442,14 +480,20 @@ def _get_pii(self, tr1, tr2, return_bb=False):
         a00e, c00e, a0e0e, a0b0b = self.ia_ta
         ae2e2, ab2b2 = self.ia_tt
         a0e2, b0e2, d0ee2, d0bb2 = self.ia_mix
+        if (self._ufpt_ak):
+            Pak2 = self.ia_der
+        else:
+            Pak2 = self.pk_ak*(self.k_s**2)[None, :]
 
         # Get biases
         c11 = tr1.c1(self.z_s)
         c21 = tr1.c2(self.z_s)
         cd1 = tr1.cdelta(self.z_s)
+        ck1 = tr1.ck(self.z_s)
         c12 = tr2.c1(self.z_s)
         c22 = tr2.c2(self.z_s)
         cd2 = tr2.cdelta(self.z_s)
+        ck2 = tr2.ck(self.z_s)
 
         if return_bb:
             pii = ((cd1*cd2*self._g4)[:, None]*a0b0b +
@@ -461,7 +505,8 @@ def _get_pii(self, tr1, tr2, return_bb=False):
                    (cd1*cd2*self._g4)[:, None]*a0e0e +
                    (c21*c22*self._g4)[:, None]*ae2e2 +
                    ((c11*c22+c21*c12)*self._g4)[:, None]*(a0e2+b0e2) +
-                   ((cd1*c22+cd2*c21)*self._g4)[:, None]*d0ee2)
+                   ((cd1*c22+cd2*c21)*self._g4)[:, None]*d0ee2 +
+                   (ck1*c12 + ck2*c11)[:, None] * (Pak2))
 
         return pii*self.exp_cutoff
 
@@ -483,15 +528,21 @@ def _get_pim(self, tri):
         Pd1d1 = self.pk_b1
         a00e, c00e, a0e0e, a0b0b = self.ia_ta
         a0e2, b0e2, d0ee2, d0bb2 = self.ia_mix
+        if (self._ufpt_ak):
+            Pak2 = self.ia_der
+        else:
+            Pak2 = self.pk_ak*(self.k_s**2)[None, :]
 
         # Get biases
         c1 = tri.c1(self.z_s)
         c2 = tri.c2(self.z_s)
         cd = tri.cdelta(self.z_s)
+        ck = tri.ck(self.z_s)
 
         pim = (c1[:, None] * Pd1d1 +
                (self._g4*cd)[:, None] * (a00e + c00e) +
-               (self._g4*c2)[:, None] * (a0e2 + b0e2))
+               (self._g4*c2)[:, None] * (a0e2 + b0e2) +
+               ck[:, None] * Pak2)
         return pim*self.exp_cutoff
 
     def _get_pmm(self):
@@ -656,6 +707,11 @@ def get_pk2d_template(self, kind, *, extrap_order_lok=1,
             pk = self._g4T * (self.ia_mix[0]+self.ia_mix[1])
         elif pk_name == 'm:cdelta':
             pk = self._g4T * (self.ia_ta[0]+self.ia_ta[1])
+        elif pk_name == 'm:ck':
+            if (self._ufpt_ak):
+                pk = self.ia_der
+            else:
+                pk = self.pk_ak*(self.k_s**2)
         elif pk_name == 'b2:b2':
             if self.fastpt_par['sub_lowk']:
                 s4 = self.dd_bias[7]

pyccl/nl_pt/tracers.py

@@ -10,7 +10,7 @@
 from ..pyutils import _check_array_params
 
 
-def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None,
+def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None, ak=None,
                       Om_m2_for_c2=False, Om_m_fid=0.3):
     """
     Function to convert from :math:`A_{ia}` values to :math:`c_{ia}` values,
@@ -40,8 +40,9 @@ def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None,
 
     Om_m = cosmo['Omega_m']
     rho_crit = physical_constants.RHO_CRITICAL
-    c1 = c1delta = c2 = None
+    c1 = c1delta = c2 = ck = None
     gz = cosmo.growth_factor(1./(1+z))
+    knorm = 1  # Units of Mpc/h, normalizes units out of ck term
 
     if a1 is not None:
         c1 = -1*a1*5e-14*rho_crit*Om_m/gz
@@ -54,8 +55,10 @@ def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None,
             c2 = a2*5*5e-14*rho_crit*Om_m**2/(Om_m_fid*gz**2)
         else:  # DES convention
             c2 = a2*5*5e-14*rho_crit*Om_m/(gz**2)
+    if ak is not None:
+        ck = ak*knorm**2*5e-14*rho_crit*Om_m/gz
 
-    return c1, c1delta, c2
+    return c1, c1delta, c2, ck
 
 
 class PTTracer(CCLAutoRepr):
@@ -212,7 +215,7 @@ class PTIntrinsicAlignmentTracer(PTTracer):
     """
     type = 'IA'
 
-    def __init__(self, c1, c2=None, cdelta=None):
+    def __init__(self, c1, c2=None, cdelta=None, ck=None):
 
         self.biases = {}
 
@@ -222,6 +225,8 @@ def __init__(self, c1, c2=None, cdelta=None):
         self.biases['c2'] = self._get_bias_function(c2)
         # Initialize cdelta
         self.biases['cdelta'] = self._get_bias_function(cdelta)
+        # Initialize cder
+        self.biases['ck'] = self._get_bias_function(ck)
 
     @property
     def c1(self):
@@ -240,3 +245,9 @@ def cdelta(self):
         """Internal overdensity bias function.
         """
         return self.biases['cdelta']
+
+    @property
+    def ck(self):
+        """Internal derivative bias function
+        """
+        return self.biases['ck']

pyccl/tests/test_ept_power.py

@@ -3,13 +3,12 @@
 import pytest
 from contextlib import nullcontext
 
-
 COSMO = ccl.Cosmology(
     Omega_c=0.27, Omega_b=0.045, h=0.67, sigma8=0.8, n_s=0.96,
     transfer_function='bbks', matter_power_spectrum='linear')
 TRS = {'TG': ccl.nl_pt.PTNumberCountsTracer(b1=2.0, b2=2.0, bs=2.0),
        'TI': ccl.nl_pt.PTIntrinsicAlignmentTracer(c1=2.0, c2=2.0,
-                                                  cdelta=2.0),
+                                                  cdelta=2.0, ck=2.0),
        'TM': ccl.nl_pt.PTMatterTracer()}
 PTC = ccl.nl_pt.EulerianPTCalculator(with_NC=True, with_IA=True,
                                      with_matter_1loop=True,
@@ -58,7 +57,7 @@ def test_ept_pk2d_bb_smoke():
 def test_ept_get_pk2d_nl(nl):
     ptc = ccl.nl_pt.EulerianPTCalculator(
         with_NC=True, with_IA=True, with_matter_1loop=True,
-        b1_pk_kind=nl, bk2_pk_kind=nl, cosmo=COSMO)
+        b1_pk_kind=nl, bk2_pk_kind=nl, ak2_pk_kind=nl, cosmo=COSMO)
     pk = ptc.get_biased_pk2d(TRS['TG'])
     assert isinstance(pk, ccl.Pk2D)
 
@@ -72,7 +71,7 @@ def test_ept_k2pk_types(typ_nlin, typ_nloc):
     tm = ccl.nl_pt.PTNumberCountsTracer(1., 0., 0.)
     ptc1 = ccl.nl_pt.EulerianPTCalculator(
         with_NC=True, with_IA=True, with_matter_1loop=True,
-        b1_pk_kind=typ_nlin, bk2_pk_kind=typ_nloc,
+        b1_pk_kind=typ_nlin, bk2_pk_kind=typ_nloc, ak2_pk_kind=typ_nloc,
         cosmo=COSMO)
     ptc2 = ccl.nl_pt.EulerianPTCalculator(
         with_NC=True, with_IA=True, with_matter_1loop=True,
@@ -92,7 +91,7 @@ def test_ept_deconstruction(kind):
                                          with_matter_1loop=True,
                                          cosmo=COSMO, sub_lowk=True)
     b_nc = ['b1', 'b2', 'b3nl', 'bs', 'bk2']
-    b_ia = ['c1', 'c2', 'cdelta']
+    b_ia = ['c1', 'c2', 'cdelta', 'ck']
     pk1 = ptc.get_pk2d_template(kind)
 
     def get_tr(tn):
@@ -110,14 +109,14 @@ def get_tr(tn):
             bdict[tn] = 1.0
             return ccl.nl_pt.PTIntrinsicAlignmentTracer(
                 c1=bdict['c1'], c2=bdict['c2'],
-                cdelta=bdict['cdelta'])
+                cdelta=bdict['cdelta'], ck=bdict['ck'])
 
     tn1, tn2 = kind.split(':')
     t1 = get_tr(tn1)
     t2 = get_tr(tn2)
 
     is_nl = tn1 in ["b2", "bs", "bk2", "b3nl"]
-    is_g = tn2 in ["c1", "c2", "cdelta"]
+    is_g = tn2 in ["c1", "c2", "cdelta", 'ck']
     with pytest.warns(ccl.CCLWarning) if is_nl and is_g else nullcontext():
         pk2 = ptc.get_biased_pk2d(t1, tracer2=t2)
     if pk1 is None:
@@ -135,15 +134,15 @@ def get_tr(tn):
 @pytest.mark.parametrize('kind',
                          ['c2:c2', 'c2:cdelta', 'cdelta:cdelta'])
 def test_ept_deconstruction_bb(kind):
-    b_ia = ['c1', 'c2', 'cdelta']
+    b_ia = ['c1', 'c2', 'cdelta', 'ck']
     pk1 = PTC.get_pk2d_template(kind, return_ia_bb=True)
 
     def get_tr(tn):
         bdict = {b: 0.0 for b in b_ia}
         bdict[tn] = 1.0
         return ccl.nl_pt.PTIntrinsicAlignmentTracer(
             c1=bdict['c1'], c2=bdict['c2'],
-            cdelta=bdict['cdelta'])
+            cdelta=bdict['cdelta'], ck=bdict['ck'])
 
     tn1, tn2 = kind.split(':')
     t1 = get_tr(tn1)
@@ -214,6 +213,10 @@ def test_ept_calculator_raises():
     with pytest.raises(ValueError):
         ccl.nl_pt.EulerianPTCalculator(bk2_pk_kind='non-linear')
 
+    # Wrong type of ak2 power spectra
+    with pytest.raises(ValueError):
+        ccl.nl_pt.EulerianPTCalculator(ak2_pk_kind="non-linear")
+
     # Uninitialized templates
     with pytest.raises(ccl.CCLError):
         ptc = ccl.nl_pt.EulerianPTCalculator(with_NC=True)