From a22defdfb1301b8b7b81354e66ce30a73e825c06 Mon Sep 17 00:00:00 2001 From: carlosggarcia Date: Tue, 24 Oct 2023 11:05:31 +0100 Subject: [PATCH] use quad_vec for trispectra integrals --- pyccl/halos/pk_4pt.py | 101 ++++++++++++++++++------------------------ 1 file changed, 44 insertions(+), 57 deletions(-) diff --git a/pyccl/halos/pk_4pt.py b/pyccl/halos/pk_4pt.py index bb27dfaa4..27f97b5fd 100644 --- a/pyccl/halos/pk_4pt.py +++ b/pyccl/halos/pk_4pt.py @@ -918,15 +918,6 @@ def halomod_trispectrum_3h(cosmo, hmc, k, a, prof, *, prof2=None, a_use = np.atleast_1d(a) k_use = np.atleast_1d(k) - # Romberg needs 1 + 2^n points - # Since the functions we average depend only on cos(theta) we can rewrite - # the integrals as \int_0^2pi dtheta f(cos theta) / 2pi as - # \int_0^pi dtheta f(cos theta) / pi - # Exclude theta = pi to avoid k + k' = 0 - theta = np.linspace(0, np.pi - 1e-5, 8193) - dtheta = theta[1] - theta[0] - cth = np.cos(theta) - # Check inputs prof, prof2, prof3, prof4, prof13_2pt, prof14_2pt, \ prof24_2pt, prof32_2pt = _allocate_profiles2(prof, prof2, prof3, prof4, @@ -950,9 +941,10 @@ def get_pk(k, a): # Compute bispectrum # Encapsulate code in a function - def get_kr_and_f2(): - kk = k_use[:, None, None] - kp = k_use[None, :, None] + def get_kr_and_f2(theta): + cth = np.cos(theta) + kk = k_use[:, None ] + kp = k_use[None, :] kr2 = kk ** 2 + kp ** 2 + 2 * kk * kp * cth kr = np.sqrt(kr2) @@ -965,20 +957,20 @@ def get_kr_and_f2(): return kr, f2 - kr, f2 = get_kr_and_f2() def get_Bpt(a): # We only need to compute the independent k * k * cos(theta) since Pk # only depends on the module of ki + kj pk = get_pk(k_use, a)[:, None] - pkr = get_pk(kr.flatten(), a).reshape(kr.shape) - P3 = scipy.integrate.romb(pkr * f2, dtheta, axis=-1) / np.pi + def integ(theta): + kr, f2 = get_kr_and_f2(theta) + pkr = get_pk(kr.flatten(), a).reshape(kr.shape) + return pkr * f2 + P3 = scipy.integrate.quad_vec(integ, 0, np.pi)[0] / np.pi Bpt = 6. / 7. * pk * pk.T + 2 * pk * P3 Bpt += Bpt.T - # TODO: We need the factor 2 to match Robert's code. Why? - # return 2 * Bpt return Bpt na = len(a_use) @@ -1028,9 +1020,6 @@ def get_Bpt(a): # Normalize out[ia, :, :] = tk_3h / norm - # np.savez_compressed('/home/ardok/Bpt_isotropized.npz', Bpt=get_Bpt(1), - # k=k_use) - if np.ndim(a) == 0: out = np.squeeze(out, axis=0) if np.ndim(k) == 0: @@ -1116,43 +1105,47 @@ def get_pk(k, a): return pk - # Romberg needs 1 + 2^n points - # Since the functions we average depend only on cos(theta) we can rewrite - # the integrals as \int_0^2pi dtheta f(cos theta) / 2pi as - # \int_0^pi dtheta f(cos theta) / pi - # Exclude theta = pi to avoid k + k' = 0 - theta = np.linspace(0, np.pi - 1e-5, 8193) - dtheta = theta[1] - theta[0] - cth = np.cos(theta) - - def isotropize(arr): - int_arr = scipy.integrate.romb(arr, dtheta, axis=-1) - return int_arr / np.pi - - def get_kr_f2_f2T_X(): - k = k_use[:, None, None] - kp = k_use[None, :, None] - kr2 = k ** 2 + kp ** 2 + 2 * k * kp * cth - kr = np.sqrt(kr2) + def get_P4A_P4X(a): + k = k_use[:, None] + kp = k_use[None, :] - f2 = 5./7. - 0.5 * (1 + k ** 2 / kr2) * (1 + kp / k * cth) + \ - 2/7. * k ** 2 / kr2 * (1 + kp / k * cth)**2 - f2[np.where(kr == 0)] = 13. / 28 + def integ(theta): + cth = np.cos(theta) + kr2 = k ** 2 + kp ** 2 + 2 * k * kp * cth + kr = np.sqrt(kr2) - # k <-> k' - f2T = np.transpose(f2, (1, 0, 2)) + f2 = 5./7. - 0.5 * (1 + k ** 2 / kr2) * (1 + kp / k * cth) + \ + 2/7. * k ** 2 / kr2 * (1 + kp / k * cth)**2 + f2[np.where(kr == 0)] = 13. / 28 + pkr = get_pk(kr.flatten(), a).reshape((nk, nk)) + return np.array([pkr * f2**2, pkr * f2 * f2.T]) + P4A, P4X = scipy.integrate.quad_vec(integ, 0, np.pi)[0] / np.pi + + return P4A, P4X + + def get_X(): + k = k_use[:, None] + kp = k_use[None, :] r = kp / k - intd = (5 * r + (7 - 2*r**2)*cth) / (1 + r**2 + 2*r*cth) * \ - (3/7. * r + 0.5 * (1 + r**2) * cth + 4/7. * r * cth**2) - # When kr = 0, r = 1 and intd = 0 - intd[np.where(kr == 0)] = 0 - X = -7./4. * (1 + r.reshape(nk, nk)**2) + isotropize(intd) + def integ(theta): + cth = np.cos(theta) + kr2 = k ** 2 + kp ** 2 + 2 * k * kp * cth + kr = np.sqrt(kr2) + intd = (5 * r + (7 - 2*r**2)*cth) / (1 + r**2 + 2*r*cth) * \ + (3/7. * r + 0.5 * (1 + r**2) * cth + 4/7. * r * cth**2) + # When kr = 0, r = 1 and intd = 0 + intd[np.where(kr == 0)] = 0 + return intd - return kr, f2, f2T, X + isotropized_integ = \ + scipy.integrate.quad_vec(integ, 0, np.pi)[0] / np.pi - kr, f2, f2T, X = get_kr_f2_f2T_X() + X = -7./4. * (1 + r**2) + isotropized_integ + return X + + X = get_X() out = np.zeros([na, nk, nk]) for ia, aa in enumerate(a_use): # Compute profile normalizations @@ -1161,10 +1154,7 @@ def get_kr_f2_f2T_X(): norm = norm1 * norm2 * norm3 * norm4 pk = get_pk(k_use, aa)[:, None] - pkr = get_pk(kr.flatten(), aa).reshape((nk, nk, theta.size)) - - P4A = isotropize(f2 ** 2 * pkr) - P4X = isotropize(f2 * f2T * pkr) + P4A, P4X = get_P4A_P4X(aa) t1113 = 4/9. * pk**2 * pk.T * X t1113 += t1113.T @@ -1189,9 +1179,6 @@ def get_kr_f2_f2T_X(): out = np.squeeze(out, axis=-1) out = np.squeeze(out, axis=-1) - # np.savez_compressed('/home/ardok/Tpt_isotropized.npz', Tpt=t1113+t1122, - # k=k_use) - return out