Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ia der term #1137

Open
wants to merge 17 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 16 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 85 additions & 15 deletions pyccl/nl_pt/ept.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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!).

Expand Down Expand Up @@ -115,6 +116,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
Expand All @@ -131,6 +135,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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not require that people have the latest version of FastPT? My first reaction is that this adds unnecessary complexity that will need to be deprecated afterwards anyway (with the headache it entails, since it changes the API).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there are two issues. One, for the basic k^2 term, we don't need FAST-PT, and we want the flexibility to choose the P_delta from the options CCL has access to. There may be future versions of this term that need some calculation to be done in FAST-PT.

Two, more generally, do we want to demand the latest (or fairly recent) FAST-PT for all CCL ept users? Probably good in the long term. Right now, FAST-PT is also changing rapidly, and is actually a bit behind in some functionality, so we don't expect all CCL users to actually have access to the bleeding edge.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, I see, thanks @jablazek .

Then, how about:

  1. Remove the usefptk option from this function.
  2. Allow for ak2_pk_kind to take an additional option (e.g. 'pt_IA_k2'), in which case the calculator uses the fast-pt prediction.

Then, if users select this option, you check if the fast-pt version allows for it and throw an error if it doesn't. My main aim here is not polluting the API with arguments that we may need to deprecate soon-ish, since they become a headache.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As my message above, I've implemented what I meant in a new commit here. Please take a look and feel free to disagree with me!

"""
__repr_attrs__ = __eq_attrs__ = ('with_NC', 'with_IA', 'with_matter_1loop',
'k_s', 'a_s', 'exp_cutoff', 'b1_pk_kind',
Expand All @@ -141,12 +147,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):
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,
Expand Down Expand Up @@ -179,7 +186,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):
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'],
Expand All @@ -191,8 +216,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

Expand Down Expand Up @@ -260,6 +288,9 @@ 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):
self.ia_der = self.pt.IA_der(**kw)
reshape_fastpt(self.ia_der)

# b1/bk power spectrum
pks = {}
Expand All @@ -281,6 +312,24 @@ def reshape_fastpt(tupl):
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]

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe I missed something, but wouldn't you achieve the same just adding self.ak2_pk_kind to the previous if-else statement, and then doing

self.pk_ak = pks[self.ak2_pk_kind]

?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jablazek @CarterDW , I have just submitted a commit that implements this the way I meant it. Sorry this was easier than explaining here with words. Check out the differences here, and feel free to disagree with me!

# Reset template power spectra
self._pk2d_temp = {}
self._cosmo = cosmo
Expand Down Expand Up @@ -364,6 +413,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):
Pak2 = self.ia_der
else:
Pak2 = self.pk_ak*(self.k_s**2)[None, :]

# Get biases
b1 = trg.b1(self.z_s)
Expand All @@ -379,10 +432,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):
Expand Down Expand Up @@ -442,14 +497,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):
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 +
Expand All @@ -461,7 +522,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

Expand All @@ -483,15 +545,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):
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):
Expand Down Expand Up @@ -656,6 +724,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]
Expand Down
19 changes: 15 additions & 4 deletions pyccl/nl_pt/tracers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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 = {}

Expand All @@ -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):
Expand All @@ -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']
21 changes: 12 additions & 9 deletions pyccl/tests/test_ept_power.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand All @@ -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,
Expand All @@ -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):
Expand All @@ -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:
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
Loading