From bbb8be6a42716e7365d972fa793aa961531a9397 Mon Sep 17 00:00:00 2001
From: Carter <cdw7303@gmail.com>
Date: Tue, 17 Oct 2023 15:20:57 -0500
Subject: [PATCH 01/14] Added a cder term (not in cross correlations)

---
 pyccl/nl_pt/ept.py     |  2 ++
 pyccl/nl_pt/tracers.py | 19 +++++++++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index a50acc6a0..15fdba86b 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -260,6 +260,8 @@ def reshape_fastpt(tupl):
             reshape_fastpt(self.ia_tt)
             self.ia_mix = self.pt.IA_mix(**kw)
             reshape_fastpt(self.ia_mix)
+            self.ia_der = self.pt.IA_der(**kw)
+            reshape_fastpt(self.ia_der)
 
         # b1/bk power spectrum
         pks = {}
diff --git a/pyccl/nl_pt/tracers.py b/pyccl/nl_pt/tracers.py
index b2e5ae313..703bab2ad 100644
--- a/pyccl/nl_pt/tracers.py
+++ b/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, ader=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,10 @@ 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 = cder = None
     gz = cosmo.growth_factor(1./(1+z))
+    knorm = 1 #Units of Mpc/h, normalizes units out of derivative term
+    
 
     if a1 is not None:
         c1 = -1*a1*5e-14*rho_crit*Om_m/gz
@@ -54,8 +56,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 ader is not None:
+        cder= ader*knorm**2*5e-14*rho_crit*Om_m/gz
 
-    return c1, c1delta, c2
+    return c1, c1delta, c2, cder
 
 
 class PTTracer(CCLAutoRepr):
@@ -212,7 +216,7 @@ class PTIntrinsicAlignmentTracer(PTTracer):
     """
     type = 'IA'
 
-    def __init__(self, c1, c2=None, cdelta=None):
+    def __init__(self, c1, c2=None, cdelta=None, cder=None):
 
         self.biases = {}
 
@@ -222,6 +226,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['cder'] = self._get_bias_function(cder)
 
     @property
     def c1(self):
@@ -240,3 +246,8 @@ def cdelta(self):
         """Internal overdensity bias function.
         """
         return self.biases['cdelta']
+    
+    @property
+    def cder(self):
+        """Internal derivative bias function
+        """

From 67b0985ff3cf0d5e4c8b1a9e70838fe6add23ee6 Mon Sep 17 00:00:00 2001
From: Carter <cdw7303@gmail.com>
Date: Thu, 19 Oct 2023 15:32:01 -0500
Subject: [PATCH 02/14] Added derivative contribution to tracers

---
 pyccl/nl_pt/ept.py | 40 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index 15fdba86b..f7ca80c8b 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -141,6 +141,7 @@ 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
@@ -191,9 +192,12 @@ 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
-        if (self.b1_pk_kind == 'pt') or (self.bk2_pk_kind == 'pt'):
+        self.ak2_pk_kind = bk2_pk_kind
+        if (self.b1_pk_kind == 'pt') or (self.bk2_pk_kind == 'pt') or (self.ak2_pk_kind == 'pt'):
             self.with_matter_1loop = True
 
         # Initialize all expensive arrays to ``None``.
@@ -283,6 +287,24 @@ def reshape_fastpt(tupl):
         self.pk_b1 = pks[self.b1_pk_kind]
         self.pk_bk = pks[self.bk2_pk_kind]
 
+        # a1/ak power spectrum
+        pksa = {}
+        if 'nonlinear' in [self.ak2_pk_kind]:
+            pksa['nonlinear'] = np.array([cosmo.nonlin_matter_power(self.k_s, a)
+                                         for a in self.a_s])
+        if 'linear' in [self.ak2_pk_kind]:
+            pksa['linear'] = np.array([cosmo.linear_matter_power(self.k_s, a)
+                                      for a in self.a_s])
+        if 'pt' in [self.ak2_pk_kind]:
+            if 'linear' in pksa:
+                pka = pks['linear']
+            else:
+                pka = np.array([cosmo.linear_matter_power(self.k_s, a) for a in self.a_s])
+            pk += self._g4T*self.one_loop_dd[0]
+            pks['pt'] = pk
+        self.pk_ak=pks[self.ak2_pk_kind]
+        
+    
         # Reset template power spectra
         self._pk2d_temp = {}
         self._cosmo = cosmo
@@ -366,6 +388,7 @@ def _get_pgi(self, trg, tri):
         Pd1d1 = self.pk_b1
         a00e, c00e, a0e0e, a0b0b = self.ia_ta
         a0e2, b0e2, d0ee2, d0bb2 = self.ia_mix
+        Pak2 = self.pk_ak*(self.k_s**2)[None, :]
 
         # Get biases
         b1 = trg.b1(self.z_s)
@@ -381,10 +404,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.cder(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):
@@ -444,14 +469,17 @@ 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
+        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.cder(self.z_s)
         c12 = tr2.c1(self.z_s)
         c22 = tr2.c2(self.z_s)
         cd2 = tr2.cdelta(self.z_s)
+        ck2 = tr2.cder(self.z_s)
 
         if return_bb:
             pii = ((cd1*cd2*self._g4)[:, None]*a0b0b +
@@ -463,7 +491,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
 
@@ -485,15 +514,18 @@ def _get_pim(self, tri):
         Pd1d1 = self.pk_b1
         a00e, c00e, a0e0e, a0b0b = self.ia_ta
         a0e2, b0e2, d0ee2, d0bb2 = self.ia_mix
+        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.cder(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):

From 0a1222ea4485aec4ca2a0d14575bb5cc7b2a783d Mon Sep 17 00:00:00 2001
From: Carter <cdw7303@gmail.com>
Date: Mon, 23 Oct 2023 11:20:23 -0500
Subject: [PATCH 03/14] Changed cder to ck

---
 pyccl/nl_pt/ept.py     | 20 ++++++++++----------
 pyccl/nl_pt/tracers.py | 15 ++++++++-------
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index f7ca80c8b..3ffc404f4 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -196,8 +196,8 @@ def __init__(self, *, with_NC=False, with_IA=False,
             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 = bk2_pk_kind
-        if (self.b1_pk_kind == 'pt') or (self.bk2_pk_kind == 'pt') or (self.ak2_pk_kind == 'pt'):
+        self.ak2_pk_kind = ak2_pk_kind
+        if (self.b1_pk_kind == 'pt') or (self.bk2_pk_kind == 'pt'):
             self.with_matter_1loop = True
 
         # Initialize all expensive arrays to ``None``.
@@ -297,12 +297,12 @@ def reshape_fastpt(tupl):
                                       for a in self.a_s])
         if 'pt' in [self.ak2_pk_kind]:
             if 'linear' in pksa:
-                pka = pks['linear']
+                pka = pksa['linear']
             else:
                 pka = np.array([cosmo.linear_matter_power(self.k_s, a) for a in self.a_s])
-            pk += self._g4T*self.one_loop_dd[0]
-            pks['pt'] = pk
-        self.pk_ak=pks[self.ak2_pk_kind]
+            pka += self._g4T*self.one_loop_dd[0]
+            pksa['pt'] = pka
+        self.pk_ak=pksa[self.ak2_pk_kind]
         
     
         # Reset template power spectra
@@ -404,7 +404,7 @@ 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.cder(self.z_s)
+        ck = tri.ck(self.z_s)
 
         pgi = b1[:, None] * (c1[:, None] * Pd1d1 +
                              (self._g4*cd)[:, None] * (a00e + c00e) +
@@ -475,11 +475,11 @@ def _get_pii(self, tr1, tr2, return_bb=False):
         c11 = tr1.c1(self.z_s)
         c21 = tr1.c2(self.z_s)
         cd1 = tr1.cdelta(self.z_s)
-        ck1 = tr1.cder(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.cder(self.z_s)
+        ck2 = tr2.ck(self.z_s)
 
         if return_bb:
             pii = ((cd1*cd2*self._g4)[:, None]*a0b0b +
@@ -520,7 +520,7 @@ def _get_pim(self, tri):
         c1 = tri.c1(self.z_s)
         c2 = tri.c2(self.z_s)
         cd = tri.cdelta(self.z_s)
-        ck = tri.cder(self.z_s)
+        ck = tri.ck(self.z_s)
 
         pim = (c1[:, None] * Pd1d1 +
                (self._g4*cd)[:, None] * (a00e + c00e) +
diff --git a/pyccl/nl_pt/tracers.py b/pyccl/nl_pt/tracers.py
index 703bab2ad..1973c789c 100644
--- a/pyccl/nl_pt/tracers.py
+++ b/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, ader=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,
@@ -56,10 +56,10 @@ def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None, ader=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 ader is not None:
-        cder= ader*knorm**2*5e-14*rho_crit*Om_m/gz
+    if ak is not None:
+        ck= ak*knorm**2*5e-14*rho_crit*Om_m/gz
 
-    return c1, c1delta, c2, cder
+    return c1, c1delta, c2, ck
 
 
 class PTTracer(CCLAutoRepr):
@@ -216,7 +216,7 @@ class PTIntrinsicAlignmentTracer(PTTracer):
     """
     type = 'IA'
 
-    def __init__(self, c1, c2=None, cdelta=None, cder=None):
+    def __init__(self, c1, c2=None, cdelta=None, ck=None):
 
         self.biases = {}
 
@@ -227,7 +227,7 @@ def __init__(self, c1, c2=None, cdelta=None, cder=None):
         # Initialize cdelta
         self.biases['cdelta'] = self._get_bias_function(cdelta)
         # Initialize cder
-        self.biases['cder'] = self._get_bias_function(cder)
+        self.biases['ck'] = self._get_bias_function(ck)
 
     @property
     def c1(self):
@@ -248,6 +248,7 @@ def cdelta(self):
         return self.biases['cdelta']
     
     @property
-    def cder(self):
+    def ck(self):
         """Internal derivative bias function
         """
+        return self.biases['ck']

From b3f3d918e27f3b90e4716afdf93fd1d8b5b699ca Mon Sep 17 00:00:00 2001
From: Carter <cdw7303@gmail.com>
Date: Thu, 9 Nov 2023 12:56:24 -0600
Subject: [PATCH 04/14] Added unit test functionality for der term

---
 pyccl/nl_pt/ept.py            | 68 ++++++++++++++++++++++++-----------
 pyccl/nl_pt/tracers.py        |  2 +-
 pyccl/tests/test_ept_power.py | 21 ++++++-----
 3 files changed, 61 insertions(+), 30 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index 3ffc404f4..25949b7ed 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -13,19 +13,19 @@
     '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 +47,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 +115,9 @@ 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``.
         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
@@ -131,6 +134,8 @@ class EulerianPTCalculator(CCLAutoRepr):
              documentation for more details.
         sub_lowk (:obj:`bool`): if ``True``, the small-scale white noise
              contribution to some of the terms will be subtracted.
+        usefptk (::obj:`bool`):) if `True``, will use the FASTPT IA k2
+        term instead of the CCL IA k2 term
     """
     __repr_attrs__ = __eq_attrs__ = ('with_NC', 'with_IA', 'with_matter_1loop',
                                      'k_s', 'a_s', 'exp_cutoff', 'b1_pk_kind',
@@ -143,11 +148,11 @@ def __init__(self, *, with_NC=False, with_IA=False,
                  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):
+                 P_window=None, C_window=0.75, sub_lowk=False, usefptk=False):
         self.with_matter_1loop = with_matter_1loop
         self.with_NC = with_NC
         self.with_IA = with_IA
-
+        self.ufpt=usefptk
         # Set FAST-PT parameters
         self.fastpt_par = {'pad_factor': pad_factor,
                            'low_extrap': low_extrap,
@@ -180,7 +185,16 @@ 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:
+            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(f"Your FAST-PT version lacks a required attribute. 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):
+            raise ValueError(f"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'],
@@ -264,8 +278,9 @@ def reshape_fastpt(tupl):
             reshape_fastpt(self.ia_tt)
             self.ia_mix = self.pt.IA_mix(**kw)
             reshape_fastpt(self.ia_mix)
-            self.ia_der = self.pt.IA_der(**kw)
-            reshape_fastpt(self.ia_der)
+            if(self.ufpt):
+                self.ia_der = self.pt.IA_der(**kw)
+                reshape_fastpt(self.ia_der)
 
         # b1/bk power spectrum
         pks = {}
@@ -287,7 +302,7 @@ def reshape_fastpt(tupl):
         self.pk_b1 = pks[self.b1_pk_kind]
         self.pk_bk = pks[self.bk2_pk_kind]
 
-        # a1/ak power spectrum
+        # ak power spectrum
         pksa = {}
         if 'nonlinear' in [self.ak2_pk_kind]:
             pksa['nonlinear'] = np.array([cosmo.nonlin_matter_power(self.k_s, a)
@@ -388,7 +403,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
-        Pak2 = self.pk_ak*(self.k_s**2)[None, :]
+        if(self.ufpt):
+            Pak2 = self.ia_der
+        else:
+            Pak2 = self.pk_ak*(self.k_s**2)[None, :]
 
         # Get biases
         b1 = trg.b1(self.z_s)
@@ -409,7 +427,7 @@ def _get_pgi(self, trg, tri):
         pgi = b1[:, None] * (c1[:, None] * Pd1d1 +
                              (self._g4*cd)[:, None] * (a00e + c00e) +
                              (self._g4*c2)[:, None] * (a0e2 + b0e2) + 
-                             ck[:, None]*(Pak2))
+                             ck[:, None] * Pak2)
         return pgi*self.exp_cutoff
 
     def _get_pgm(self, trg):
@@ -469,7 +487,10 @@ 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
-        Pak2 = self.pk_ak*(self.k_s**2)[None, :]
+        if(self.ufpt):
+            Pak2 = self.ia_der
+        else:
+            Pak2 = self.pk_ak*(self.k_s**2)[None, :]
 
         # Get biases
         c11 = tr1.c1(self.z_s)
@@ -492,7 +513,7 @@ def _get_pii(self, tr1, tr2, return_bb=False):
                    (c21*c22*self._g4)[:, None]*ae2e2 +
                    ((c11*c22+c21*c12)*self._g4)[:, None]*(a0e2+b0e2) +
                    ((cd1*c22+cd2*c21)*self._g4)[:, None]*d0ee2 +
-                   (ck1*c12 + ck2*c11)[:,None]* (Pak2))
+                   (ck1*c12 + ck2*c11)[:,None] * (Pak2))
 
         return pii*self.exp_cutoff
 
@@ -514,7 +535,10 @@ def _get_pim(self, tri):
         Pd1d1 = self.pk_b1
         a00e, c00e, a0e0e, a0b0b = self.ia_ta
         a0e2, b0e2, d0ee2, d0bb2 = self.ia_mix
-        Pak2 = self.pk_ak*(self.k_s**2)[None, :]
+        if(self.ufpt):
+            Pak2 = self.ia_der
+        else:
+            Pak2 = self.pk_ak*(self.k_s**2)[None, :]
 
         # Get biases
         c1 = tri.c1(self.z_s)
@@ -525,7 +549,7 @@ def _get_pim(self, tri):
         pim = (c1[:, None] * Pd1d1 +
                (self._g4*cd)[:, None] * (a00e + c00e) +
                (self._g4*c2)[:, None] * (a0e2 + b0e2) +
-               (ck)[:,None]*Pak2)
+               ck[:,None] * Pak2)
         return pim*self.exp_cutoff
 
     def _get_pmm(self):
@@ -690,6 +714,8 @@ 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':
+            pk = self.pk_ak*(self.k_s**2)
         elif pk_name == 'b2:b2':
             if self.fastpt_par['sub_lowk']:
                 s4 = self.dd_bias[7]
diff --git a/pyccl/nl_pt/tracers.py b/pyccl/nl_pt/tracers.py
index 1973c789c..48608b4a0 100644
--- a/pyccl/nl_pt/tracers.py
+++ b/pyccl/nl_pt/tracers.py
@@ -42,7 +42,7 @@ def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None, ak=None,
     rho_crit = physical_constants.RHO_CRITICAL
     c1 = c1delta = c2 = cder = None
     gz = cosmo.growth_factor(1./(1+z))
-    knorm = 1 #Units of Mpc/h, normalizes units out of derivative term
+    knorm = 1 #Units of Mpc/h, normalizes units out of ck term
     
 
     if a1 is not None:
diff --git a/pyccl/tests/test_ept_power.py b/pyccl/tests/test_ept_power.py
index 9f4a95224..b9cf7703d 100644
--- a/pyccl/tests/test_ept_power.py
+++ b/pyccl/tests/test_ept_power.py
@@ -3,13 +3,14 @@
 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 +59,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 +73,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 +93,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 +111,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,7 +136,7 @@ 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):
@@ -143,7 +144,7 @@ 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)
@@ -213,6 +214,10 @@ def test_ept_calculator_raises():
     # Wrong type of b1 and bk2 power spectra
     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):

From 159f3330d3efbe3d65406338656f74e086f2ab2a Mon Sep 17 00:00:00 2001
From: Carter <cdw7303@gmail.com>
Date: Tue, 14 Nov 2023 10:17:13 -0600
Subject: [PATCH 05/14] changed cder to ck in tracers.py

---
 pyccl/nl_pt/tracers.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/pyccl/nl_pt/tracers.py b/pyccl/nl_pt/tracers.py
index 48608b4a0..73c6daa8b 100644
--- a/pyccl/nl_pt/tracers.py
+++ b/pyccl/nl_pt/tracers.py
@@ -40,10 +40,9 @@ def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None, ak=None,
 
     Om_m = cosmo['Omega_m']
     rho_crit = physical_constants.RHO_CRITICAL
-    c1 = c1delta = c2 = cder = 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
-    
+    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
@@ -57,7 +56,7 @@ def translate_IA_norm(cosmo, *, z, a1=1.0, a1delta=None, a2=None, ak=None,
         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
+        ck = ak*knorm**2*5e-14*rho_crit*Om_m/gz
 
     return c1, c1delta, c2, ck
 
@@ -246,7 +245,7 @@ def cdelta(self):
         """Internal overdensity bias function.
         """
         return self.biases['cdelta']
-    
+
     @property
     def ck(self):
         """Internal derivative bias function

From a0fdac50ca7f4717f4a9ef4a68be75b6d77efa8a Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Tue, 28 Nov 2023 10:12:29 -0600
Subject: [PATCH 06/14] Update test_ept_power.py

Attempted reformatting to fit flake8 format
---
 pyccl/tests/test_ept_power.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/pyccl/tests/test_ept_power.py b/pyccl/tests/test_ept_power.py
index b9cf7703d..0eb3765fe 100644
--- a/pyccl/tests/test_ept_power.py
+++ b/pyccl/tests/test_ept_power.py
@@ -3,8 +3,8 @@
 import pytest
 from contextlib import nullcontext
 
-    
-
+#Delete    
+#Delete
 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')
@@ -59,7 +59,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, ak2_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)
 
@@ -214,8 +214,8 @@ def test_ept_calculator_raises():
     # Wrong type of b1 and bk2 power spectra
     with pytest.raises(ValueError):
         ccl.nl_pt.EulerianPTCalculator(bk2_pk_kind='non-linear')
-    
-    #wrong type of ak2 power spectra
+
+    # Wrong type of ak2 power spectra
     with pytest.raises(ValueError):
         ccl.nl_pt.EulerianPTCalculator(ak2_pk_kind="non-linear")
 

From 71eb1aaddc8d1592117d0d7b7188f2f6532b2aeb Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Tue, 28 Nov 2023 10:13:47 -0600
Subject: [PATCH 07/14] Update test_ept_power.py

---
 pyccl/tests/test_ept_power.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pyccl/tests/test_ept_power.py b/pyccl/tests/test_ept_power.py
index 0eb3765fe..34b427824 100644
--- a/pyccl/tests/test_ept_power.py
+++ b/pyccl/tests/test_ept_power.py
@@ -3,8 +3,6 @@
 import pytest
 from contextlib import nullcontext
 
-#Delete    
-#Delete
 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')

From 037f7d80a632a52eba409eaed73e93e9006d15f6 Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Tue, 28 Nov 2023 10:24:31 -0600
Subject: [PATCH 08/14] Reformatted ept.py

Attempted to better fit flake8 formatting requirements
---
 pyccl/nl_pt/ept.py | 62 +++++++++++++++++++++++++---------------------
 1 file changed, 34 insertions(+), 28 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index 25949b7ed..8ccdf37f8 100644
--- a/pyccl/nl_pt/ept.py
+++ b/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',
-    'm:ck' : 'm:ck', 'b1:b1': 'm:m', 'b1:b2': 'm:b2',
+    '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',
+    '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: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', 'c1:ck' : 'm:ck', 'c2:c2': 'c2:c2',
-    'c2:cdelta': 'c2:cdelta', 'cdelta:cdelta': 'cdelta:cdelta', 'ck:ck' : 'zero'}
+    '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):
@@ -116,7 +117,7 @@ class EulerianPTCalculator(CCLAutoRepr):
             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 
+            term of the IA expansion. Same options and default as
             ``b1_pk_kind``.
         pad_factor (:obj:`float`): fraction of the :math:`\\log_{10}(k)`
              interval you to add as padding for FFTLog calculations.
@@ -152,7 +153,7 @@ def __init__(self, *, with_NC=False, with_IA=False,
         self.with_matter_1loop = with_matter_1loop
         self.with_NC = with_NC
         self.with_IA = with_IA
-        self.ufpt=usefptk
+        self.ufpt = usefptk
         # Set FAST-PT parameters
         self.fastpt_par = {'pad_factor': pad_factor,
                            'low_extrap': low_extrap,
@@ -187,14 +188,19 @@ def __init__(self, *, with_NC=False, with_IA=False,
         # Call FAST-PT
         try:
             import fastpt as fpt
-        except:
-            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(f"Your FAST-PT version lacks a required attribute. 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):
-            raise ValueError(f"You need a newer version of FAST-PT to use fpt for k2 term")
+        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 attribute.
+            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):
+            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'],
@@ -278,7 +284,7 @@ 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):
+            if (self.ufpt):
                 self.ia_der = self.pt.IA_der(**kw)
                 reshape_fastpt(self.ia_der)
 
@@ -305,8 +311,8 @@ def reshape_fastpt(tupl):
         # ak power spectrum
         pksa = {}
         if 'nonlinear' in [self.ak2_pk_kind]:
-            pksa['nonlinear'] = np.array([cosmo.nonlin_matter_power(self.k_s, a)
-                                         for a in self.a_s])
+            pksa['nonlinear'] = np.array(
+                  cosmo.nonlin_matter_power(self.k_s, a) for a in self.a_s])
         if 'linear' in [self.ak2_pk_kind]:
             pksa['linear'] = np.array([cosmo.linear_matter_power(self.k_s, a)
                                       for a in self.a_s])
@@ -314,12 +320,12 @@ def reshape_fastpt(tupl):
             if 'linear' in pksa:
                 pka = pksa['linear']
             else:
-                pka = np.array([cosmo.linear_matter_power(self.k_s, a) for a in self.a_s])
+                pka = np.array([cosmo.linear_matter_power(self.k_s, a) 
+                                for a in self.a_s]) #Shorten
             pka += self._g4T*self.one_loop_dd[0]
             pksa['pt'] = pka
-        self.pk_ak=pksa[self.ak2_pk_kind]
-        
-    
+        self.pk_ak = pksa[self.ak2_pk_kind]
+
         # Reset template power spectra
         self._pk2d_temp = {}
         self._cosmo = cosmo
@@ -403,7 +409,7 @@ 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):
+        if (self.ufpt):
             Pak2 = self.ia_der
         else:
             Pak2 = self.pk_ak*(self.k_s**2)[None, :]
@@ -426,7 +432,7 @@ def _get_pgi(self, trg, tri):
 
         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
 
@@ -487,7 +493,7 @@ 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):
+        if (self.ufpt):
             Pak2 = self.ia_der
         else:
             Pak2 = self.pk_ak*(self.k_s**2)[None, :]
@@ -513,7 +519,7 @@ def _get_pii(self, tr1, tr2, return_bb=False):
                    (c21*c22*self._g4)[:, None]*ae2e2 +
                    ((c11*c22+c21*c12)*self._g4)[:, None]*(a0e2+b0e2) +
                    ((cd1*c22+cd2*c21)*self._g4)[:, None]*d0ee2 +
-                   (ck1*c12 + ck2*c11)[:,None] * (Pak2))
+                   (ck1*c12 + ck2*c11)[:, None] * (Pak2))
 
         return pii*self.exp_cutoff
 
@@ -535,7 +541,7 @@ 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):
+        if (self.ufpt):
             Pak2 = self.ia_der
         else:
             Pak2 = self.pk_ak*(self.k_s**2)[None, :]
@@ -549,7 +555,7 @@ def _get_pim(self, tri):
         pim = (c1[:, None] * Pd1d1 +
                (self._g4*cd)[:, None] * (a00e + c00e) +
                (self._g4*c2)[:, None] * (a0e2 + b0e2) +
-               ck[:,None] * Pak2)
+               ck[:, None] * Pak2)
         return pim*self.exp_cutoff
 
     def _get_pmm(self):

From ed65114a0e9c1775c283779e718d6eaac048bbf3 Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Tue, 28 Nov 2023 15:58:33 -0600
Subject: [PATCH 09/14] Update ept.py

Changed the statements from lines 192-203 so that it fits flake8
---
 pyccl/nl_pt/ept.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index 8ccdf37f8..b67872e1a 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -189,18 +189,18 @@ def __init__(self, *, with_NC=False, with_IA=False,
         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")
+            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 attribute.
-            You may have the wrong fast-pt install.
-            Try running pip install fast-pt or conda install fast-pt, then try again")
+            raise ValueError("Your FAST-PT version lacks a required attribute. "
+            "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):
-            raise ValueError("You need a newer version of 
-            FAST-PT to use fpt for k2 term")
+            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'],

From fa48fd1f91ac2d32a3ce8dd88e73a8b14d13840e Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Tue, 28 Nov 2023 23:53:30 -0600
Subject: [PATCH 10/14] Update ept.py

---
 pyccl/nl_pt/ept.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index b67872e1a..b8d7fe2e8 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -312,7 +312,7 @@ def reshape_fastpt(tupl):
         pksa = {}
         if 'nonlinear' in [self.ak2_pk_kind]:
             pksa['nonlinear'] = np.array(
-                  cosmo.nonlin_matter_power(self.k_s, a) for a in self.a_s])
+                  [cosmo.nonlin_matter_power(self.k_s, a) for a in self.a_s])
         if 'linear' in [self.ak2_pk_kind]:
             pksa['linear'] = np.array([cosmo.linear_matter_power(self.k_s, a)
                                       for a in self.a_s])

From 69dbb5f2d3929ca714f8d214e9d90dd5594b2fd0 Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Wed, 29 Nov 2023 17:19:30 -0600
Subject: [PATCH 11/14] Formatted ept.py

---
 pyccl/nl_pt/ept.py | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index b8d7fe2e8..70e48a8a4 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -25,7 +25,7 @@
     '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', 'c1:ck': 'm:ck',
-    'c2:c2': 'c2:c2','c2:cdelta': 'c2:cdelta',
+    'c2:c2': 'c2:c2', 'c2:cdelta': 'c2:cdelta',
     'cdelta:cdelta': 'cdelta:cdelta', 'ck:ck': 'zero'}
 
 
@@ -190,17 +190,21 @@ def __init__(self, *, with_NC=False, with_IA=False,
             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")
+                              "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 attribute. "
-            "You may have the wrong fast-pt install. "
-            "Try running pip install fast-pt or conda install fast-pt, then try again")
+            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):
-            raise ValueError("You need a newer version of " 
-            "FAST-PT to use fpt for k2 term")
+            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'],
@@ -320,8 +324,8 @@ def reshape_fastpt(tupl):
             if 'linear' in pksa:
                 pka = pksa['linear']
             else:
-                pka = np.array([cosmo.linear_matter_power(self.k_s, a) 
-                                for a in self.a_s]) #Shorten
+                pka = np.array([cosmo.linear_matter_power(self.k_s, a)
+                                for a in self.a_s])
             pka += self._g4T*self.one_loop_dd[0]
             pksa['pt'] = pka
         self.pk_ak = pksa[self.ak2_pk_kind]

From 998a474e465c7d8417951916e015efa81b16792f Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Thu, 7 Dec 2023 12:41:11 -0600
Subject: [PATCH 12/14] Formatted ept.py

---
 pyccl/nl_pt/ept.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index 70e48a8a4..635fcc37f 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -316,7 +316,7 @@ def reshape_fastpt(tupl):
         pksa = {}
         if 'nonlinear' in [self.ak2_pk_kind]:
             pksa['nonlinear'] = np.array(
-                  [cosmo.nonlin_matter_power(self.k_s, a) for a in self.a_s])
+            [cosmo.nonlin_matter_power(self.k_s, a) for a in self.a_s])
         if 'linear' in [self.ak2_pk_kind]:
             pksa['linear'] = np.array([cosmo.linear_matter_power(self.k_s, a)
                                       for a in self.a_s])

From ffcdef86f557e295bc5ffa85c7d5a1716c68f5c9 Mon Sep 17 00:00:00 2001
From: CarterDW <140671500+CarterDW@users.noreply.github.com>
Date: Thu, 7 Dec 2023 14:34:45 -0600
Subject: [PATCH 13/14] Reformatted ept.py for flake8

---
 pyccl/nl_pt/ept.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index 635fcc37f..7871b585d 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -315,8 +315,8 @@ def reshape_fastpt(tupl):
         # ak power spectrum
         pksa = {}
         if 'nonlinear' in [self.ak2_pk_kind]:
-            pksa['nonlinear'] = np.array(
-            [cosmo.nonlin_matter_power(self.k_s, a) for a in self.a_s])
+            pksa['nonlinear'] = np.array([cosmo.nonlin_matter_power(
+                                          self.k_s, a)for a in self.a_s])
         if 'linear' in [self.ak2_pk_kind]:
             pksa['linear'] = np.array([cosmo.linear_matter_power(self.k_s, a)
                                       for a in self.a_s])

From c551530432439df6ba00ade3159851bb0ff7be5f Mon Sep 17 00:00:00 2001
From: damonge <dam.phys@gmail.com>
Date: Wed, 4 Dec 2024 08:20:24 +0000
Subject: [PATCH 14/14] ak options simplify

---
 pyccl/nl_pt/ept.py | 52 +++++++++++++++++-----------------------------
 1 file changed, 19 insertions(+), 33 deletions(-)

diff --git a/pyccl/nl_pt/ept.py b/pyccl/nl_pt/ept.py
index 7871b585d..372ac5fa6 100644
--- a/pyccl/nl_pt/ept.py
+++ b/pyccl/nl_pt/ept.py
@@ -118,7 +118,8 @@ class EulerianPTCalculator(CCLAutoRepr):
             ``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``.
+            ``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
@@ -135,8 +136,6 @@ class EulerianPTCalculator(CCLAutoRepr):
              documentation for more details.
         sub_lowk (:obj:`bool`): if ``True``, the small-scale white noise
              contribution to some of the terms will be subtracted.
-        usefptk (::obj:`bool`):) if `True``, will use the FASTPT IA k2
-        term instead of the CCL IA k2 term
     """
     __repr_attrs__ = __eq_attrs__ = ('with_NC', 'with_IA', 'with_matter_1loop',
                                      'k_s', 'a_s', 'exp_cutoff', 'b1_pk_kind',
@@ -149,11 +148,11 @@ def __init__(self, *, with_NC=False, with_IA=False,
                  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, usefptk=False):
+                 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 = usefptk
+        self._ufpt_ak = ak2_pk_kind == 'fastpt'
         # Set FAST-PT parameters
         self.fastpt_par = {'pad_factor': pad_factor,
                            'low_extrap': low_extrap,
@@ -202,7 +201,7 @@ def __init__(self, *, with_NC=False, with_IA=False,
                              "fast-pt or conda install fast-pt, "
                              "then try again")
 
-        if (not hasattr(fpt, "IA_der") and self.ufpt):
+        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))
@@ -288,19 +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):
+            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:
@@ -311,24 +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]
-
-        # ak power spectrum
-        pksa = {}
-        if 'nonlinear' in [self.ak2_pk_kind]:
-            pksa['nonlinear'] = np.array([cosmo.nonlin_matter_power(
-                                          self.k_s, a)for a in self.a_s])
-        if 'linear' in [self.ak2_pk_kind]:
-            pksa['linear'] = np.array([cosmo.linear_matter_power(self.k_s, a)
-                                      for a in self.a_s])
-        if 'pt' in [self.ak2_pk_kind]:
-            if 'linear' in pksa:
-                pka = pksa['linear']
-            else:
-                pka = np.array([cosmo.linear_matter_power(self.k_s, a)
-                                for a in self.a_s])
-            pka += self._g4T*self.one_loop_dd[0]
-            pksa['pt'] = pka
-        self.pk_ak = pksa[self.ak2_pk_kind]
+        self.pk_ak = pks[self.ak2_pk_kind]
 
         # Reset template power spectra
         self._pk2d_temp = {}
@@ -413,7 +396,7 @@ 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):
+        if (self._ufpt_ak):
             Pak2 = self.ia_der
         else:
             Pak2 = self.pk_ak*(self.k_s**2)[None, :]
@@ -497,7 +480,7 @@ 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):
+        if (self._ufpt_ak):
             Pak2 = self.ia_der
         else:
             Pak2 = self.pk_ak*(self.k_s**2)[None, :]
@@ -545,7 +528,7 @@ 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):
+        if (self._ufpt_ak):
             Pak2 = self.ia_der
         else:
             Pak2 = self.pk_ak*(self.k_s**2)[None, :]
@@ -725,7 +708,10 @@ def get_pk2d_template(self, kind, *, extrap_order_lok=1,
         elif pk_name == 'm:cdelta':
             pk = self._g4T * (self.ia_ta[0]+self.ia_ta[1])
         elif pk_name == 'm:ck':
-            pk = self.pk_ak*(self.k_s**2)
+            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]