Nonlin models (#734)

* added first skeleton version

* small tweak

* first complete version of pt_power. still buggy

* bugs fixed. more work to be done

* fixed bug in assert error message that appears with newest numpy

* started PTWorkspace and flaked

* added some tests

* small refactor

* added matter tracer

* added IA. Missing IAxG

* fixed bug in bias combinations for P_gg. minor comment updates

* minor updates to comments

* iaxg

* fastpt installation

* back to no fastpt in travis

* changelog

* messed up changelog

* fastpt pip in travis

* in module

* fixed interdependency

* some docs, more install, fixed init

* more docs, fixed tests

* more docs, new structure

* Update install.sh

* stupid bug

* added regression benchmark

* cleanup

* switch order of ptc and tracer2

* amp_fix

* benchmark

* work on SPT in progress.

* added SPT for Pd1d1 functionality, also lightweight PTC initialization by default.

* minor cleanup

* bug fix

* things work

* pt with spt

* exhaustive tests

* few more tests

* BB stuff

* oneloop stuff

* changelog

* matt's comments

Co-authored-by: Jonathan Blazek <blazek@berkeley.edu>
Co-authored-by: Matthew R. Becker <beckermr@users.noreply.github.com>

.travis/install.sh

@@ -42,3 +42,8 @@ py36)
     cython "camb>=1"
   ;;
 esac;
+
+# we have to activate the cond env before we install this
+source activate test-environment
+pip install https://github.com/JoeMcEwen/FAST-PT/archive/v3.0.2.tar.gz --no-deps
+conda deactivate

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 ## Python library
+- Perturbation theory power spectra #734
 - Generalized Tinker mass function to all SO mass definitions (#736)
 - Deprecated old halo model (#736)
 - Halo model power spectra #717

benchmarks/data/codes/pt_bm.py

@@ -0,0 +1,99 @@
+import numpy as np
+import pyccl as ccl
+import fastpt as fpt
+
+cosmo = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, n_s=0.96, sigma8=0.8,
+                      transfer_function='bbks')
+lkmin = -4
+lkmax = 2
+n_per_decade = 20
+z_s = np.array([0., 1.])
+a_s = 1./(1 + z_s)
+ks = np.logspace(lkmin, lkmax, (lkmax - lkmin) * n_per_decade)
+pk = ccl.linear_matter_power(cosmo, ks, 1.)
+gf = ccl.growth_factor(cosmo, a_s)
+g4 = gf**4
+pklin = np.array([ccl.linear_matter_power(cosmo, ks, a)
+                  for a in a_s])
+
+C_window=.75
+P_window=None
+n_pad=int(0.5*len(ks))
+to_do = ['one_loop_dd', 'dd_bias', 'IA']
+pt_ob=fpt.FASTPT(ks,to_do=to_do,low_extrap=-5,high_extrap=3,n_pad=n_pad)
+oloop_dd = pt_ob.one_loop_dd(pk,
+                             P_window=P_window,
+                             C_window=C_window)
+Pd1d1 = pklin + g4[:, None] * oloop_dd[0][None, :]
+dd_bias = pt_ob.one_loop_dd_bias(pk,
+                                 P_window=P_window,
+                                 C_window=C_window)
+
+ia_ta = pt_ob.IA_ta(pk,
+                    P_window=P_window,
+                    C_window=C_window)
+ia_tt = pt_ob.IA_tt(pk,
+                    P_window=P_window,
+                    C_window=C_window)
+ia_mix = pt_ob.IA_mix(pk,
+                      P_window=P_window,
+                      C_window=C_window)
+
+g4 = g4[:, None]
+Pd1d2 = g4 * dd_bias[2][None, :]
+Pd2d2 = g4 * dd_bias[3][None, :]
+Pd1s2 = g4 * dd_bias[4][None, :]
+Pd2s2 = g4 * dd_bias[5][None, :]
+Ps2s2 = g4 * dd_bias[6][None, :]
+a00e = g4 * ia_ta[0][None, :]
+c00e = g4 * ia_ta[1][None, :]
+a0e0e = g4 * ia_ta[2][None, :]
+a0b0b = g4 * ia_ta[3][None, :]
+ae2e2 = g4 * ia_tt[0][None, :]
+ab2b2 = g4 * ia_tt[1][None, :]
+a0e2 = g4 * ia_mix[0][None, :]
+b0e2 = g4 * ia_mix[1][None, :]
+d0ee2 = g4 * ia_mix[2][None, :]
+d0bb2 = g4 * ia_mix[3][None, :]
+
+b1 = 1.3
+b2 = 1.5
+bs = 1.7
+c1 = 1.9
+c2 = 2.1
+cd = 2.3
+
+pgg = (b1**2 * Pd1d1 +
+       b1*b2 * Pd1d2 +
+       0.25 * b2**2 * Pd2d2 +
+       b1 * bs * Pd1s2 +
+       0.5 * b2 * bs * Pd2s2 +
+       0.25 * bs**2 * Ps2s2)
+
+pgm = (b1 * Pd1d1 +
+       0.5* b2 * Pd1d2 +
+       0.5 * bs * Pd1s2)
+pgi = b1 * (c1 * Pd1d1 +
+            cd * (a00e + c00e) +
+            c2 * (a0e2 + b0e2))
+pii = (c1**2 * Pd1d1 +
+       2 * c1 * cd * (a00e + c00e) +
+       cd**2 * a0e0e +
+       c2**2 * ae2e2 +
+       2 * c1 * c2 * (a0e2 + b0e2) +
+       2 * cd * c2 * d0ee2)
+pii_bb = (cd**2 * a0b0b +
+          c2**2 * ab2b2 +
+          2 * cd * c2 * d0bb2)
+pim = (c1 * Pd1d1 +
+       cd * (a00e + c00e) +
+       c2 * (a0e2 + b0e2))
+
+np.savetxt("../pt_bm_z0.txt",
+           np.transpose([ks, pgg[0], pgm[0], pgi[0],
+                         pii[0], pii_bb[0], pim[0]]),
+           header='[0]-k  [1]-GG [2]-GM [3]-GI [4]-II [5]-II_BB [6]-IM')
+np.savetxt("../pt_bm_z1.txt",
+           np.transpose([ks, pgg[1], pgm[1], pgi[1],
+                         pii[1], pii_bb[1], pim[1]]),
+           header='[0]-k  [1]-GG [2]-GM [3]-GI [4]-II [5]-II_BB [6]-IM')

benchmarks/test_ptpk.py

@@ -0,0 +1,58 @@
+import os
+import numpy as np
+import pyccl as ccl
+import pyccl.nl_pt as pt
+import pytest
+
+
+# Set cosmology
+COSMO = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05,
+                      h=0.7, sigma8=0.8, n_s=0.96,
+                      transfer_function='bbks')
+
+# Redshifts
+zs = np.array([0., 1.])
+
+# Tracers
+ptt = {}
+ptt['g'] = pt.PTNumberCountsTracer(b1=1.3, b2=1.5, bs=1.7)
+ptt['i'] = pt.PTIntrinsicAlignmentTracer(c1=1.9, c2=2.1, cdelta=2.3)
+ptt['m'] = pt.PTMatterTracer()
+
+# Calculator
+ptc = pt.PTCalculator(with_NC=True, with_IA=True,
+                      log10k_min=-4, log10k_max=2,
+                      nk_per_decade=20, pad_factor=0.5)
+
+# Read data
+dirdat = os.path.join(os.path.dirname(__file__), 'data')
+data = []
+data.append(np.loadtxt(os.path.join(dirdat, 'pt_bm_z0.txt'), unpack=True))
+data.append(np.loadtxt(os.path.join(dirdat, 'pt_bm_z1.txt'), unpack=True))
+order = ['gg', 'gm', 'gi', 'ii', 'ib', 'im']
+
+
+@pytest.mark.parametrize('comb', enumerate(order))
+def test_pt_pk(comb):
+    i_d, cc = comb
+    t1, t2 = cc
+
+    return_bb = False
+    if t2 == 'b':
+        t2 = 'i'
+        return_bb = True
+
+    ptt1 = ptt[t1]
+    ptt2 = ptt[t2]
+
+    a_arr = 1./(1+np.array([0., 0.25, 0.5, 0.75, 1.]))[::-1]
+    pk = pt.get_pt_pk2d(COSMO, ptt1, tracer2=ptt2, ptc=ptc,
+                        return_ia_bb=return_bb,
+                        nonlin_pk_type='spt',
+                        a_arr=a_arr)
+    for iz, z in enumerate(zs):
+        a = 1./(1+z)
+        k = data[iz][0]
+        dpk = data[iz][i_d+1]
+        tpk = pk.eval(k, a, COSMO)
+        assert np.all(np.fabs(tpk / dpk - 1) < 1E-5)

pyccl/__init__.py

@@ -13,7 +13,7 @@
     environ["CLASS_PARAM_DIR"] = path.dirname(path.abspath(__file__))
 
 from . import ccllib as lib
-from . import core, constants, background, power, halomodel, pk2d, haloprofile, halos, massfunction
+from . import core, constants, background, power, halomodel, pk2d, haloprofile, halos, massfunction, nl_pt
 
 # Core data structures
 from .core import Cosmology

pyccl/halos/halo_model.py

@@ -527,11 +527,7 @@ def halomod_Pk2D(cosmo, hmc, prof,
             :class:`~pyccl.pk2d.Pk2D`.
 
     Returns:
-        float or array_like: integral values evaluated at each
-        combination of `k` and `a`. The shape of the output will
-        be `(N_a, N_k)` where `N_k` and `N_a` are the sizes of
-        `k` and `a` respectively. If `k` or `a` are scalars, the
-        corresponding dimension will be squeezed out on output.
+        :class:`~pyccl.pk2d.Pk2D`: halo model power spectrum.
     """
     if lk_arr is None:
         status = 0

pyccl/nl_pt/__init__.py

@@ -0,0 +1,10 @@
+# Tracers
+from .tracers import (  # noqa
+    PTTracer, PTMatterTracer,
+    PTNumberCountsTracer,
+    PTIntrinsicAlignmentTracer)
+
+# Power spectra
+from .power import (  # noqa
+    PTCalculator,
+    get_pt_pk2d)