Optimize HM autocorrelation power spectra (#891)

* Optimized HM autocorr.

---------

Co-authored-by: David Alonso <dam.phys@gmail.com>

pyccl/halos/profiles.py

@@ -53,6 +53,10 @@ def __init__(self):
                                  'plaw_fourier': -1.5,
                                  'plaw_projected': -1.}
 
+    __eq__ = object.__eq__
+
+    __hash__ = object.__hash__  # TODO: remove once __eq__ is replaced.
+
     @unlock_instance(mutate=True)
     def update_precision_fftlog(self, **kwargs):
         """ Update any of the precision parameters used by
@@ -660,7 +664,7 @@ def __init__(self, c_M_relation,
                  cumul2d_analytic=False,
                  truncated=True):
         if not isinstance(c_M_relation, Concentration):
-            raise TypeError("c_M_relation must be of type `Concentration`)")
+            raise TypeError("c_M_relation must be of type `Concentration`")
 
         self.cM = c_M_relation
         self.truncated = truncated
@@ -852,7 +856,7 @@ class HaloProfileEinasto(HaloProfile):
 
     def __init__(self, c_M_relation, truncated=True, alpha='cosmo'):
         if not isinstance(c_M_relation, Concentration):
-            raise TypeError("c_M_relation must be of type `Concentration`)")
+            raise TypeError("c_M_relation must be of type `Concentration`")
 
         self.cM = c_M_relation
         self.truncated = truncated
@@ -967,7 +971,7 @@ def __init__(self, c_M_relation,
                  projected_analytic=False,
                  cumul2d_analytic=False):
         if not isinstance(c_M_relation, Concentration):
-            raise TypeError("c_M_relation must be of type `Concentration`)")
+            raise TypeError("c_M_relation must be of type `Concentration`")
 
         self.cM = c_M_relation
         self.truncated = truncated
@@ -1483,7 +1487,7 @@ def __init__(self, c_M_relation,
                  bg_0=1., bg_p=0., bmax_0=1., bmax_p=0.,
                  a_pivot=1., ns_independent=False):
         if not isinstance(c_M_relation, Concentration):
-            raise TypeError("c_M_relation must be of type `Concentration`)")
+            raise TypeError("c_M_relation must be of type `Concentration`")
 
         self.cM = c_M_relation
         self.lMmin_0 = lMmin_0

pyccl/halos/profiles_2pt.py

@@ -28,6 +28,10 @@ class Profile2pt(CCLHalosObject):
     def __init__(self, r_corr=0.):
         self.r_corr = r_corr
 
+    __eq__ = object.__eq__
+
+    __hash__ = object.__hash__  # TODO: remove once __eq__ is replaced.
+
     def update_parameters(self, r_corr=None):
         """ Update any of the parameters associated with this 1-halo
         2-point correlator. Any parameter set to `None` won't be updated.
@@ -71,16 +75,16 @@ def fourier_2pt(self, prof, cosmo, k, M, a,
         """
         if not isinstance(prof, HaloProfile):
             raise TypeError("prof must be of type `HaloProfile`")
+        if prof2 is None:
+            prof2 = prof
+        elif not isinstance(prof2, HaloProfile):
+            raise TypeError("prof2 must be of type `HaloProfile` or None")
 
         uk1 = prof.fourier(cosmo, k, M, a, mass_def=mass_def)
 
-        if prof2 is None:
+        if prof == prof2:
             uk2 = uk1
         else:
-            if not isinstance(prof2, HaloProfile):
-                raise TypeError("prof2 must be of type "
-                                "`HaloProfile` or `None`")
-
             uk2 = prof2.fourier(cosmo, k, M, a, mass_def=mass_def)
 
         return uk1 * uk2 * (1 + self.r_corr)
@@ -111,7 +115,7 @@ def fourier_2pt(self, prof, cosmo, k, M, a,
             k (float or array_like): comoving wavenumber in Mpc^-1.
             M (float or array_like): halo mass in units of M_sun.
             a (float): scale factor.
-            prof2 (:class:`~pyccl.halos.profiles.HaloProfile`):
+            prof2 (:class:`~pyccl.halos.profiles.HaloProfileHOD` or None):
                 second halo profile for which the second-order moment
                 is desired. If `None`, the assumption is that you want
                 an auto-correlation. Note that only auto-correlations
@@ -126,11 +130,13 @@ def fourier_2pt(self, prof, cosmo, k, M, a,
             respectively. If `k` or `M` are scalars, the
             corresponding dimension will be squeezed out on output.
         """
-        if not isinstance(prof, HaloProfileHOD):
-            raise TypeError("prof must be of type `HaloProfileHOD`")
+        if prof2 is None:
+            prof2 = prof
 
-        if prof2 is not None:
-            if prof2 is not prof:
-                raise ValueError("prof2 must be the same as prof")
+        if not (isinstance(prof, HaloProfileHOD)
+                and isinstance(prof2, HaloProfileHOD)):
+            raise TypeError("prof and prof2 should be HaloProfileHOD")
+        if prof != prof2:
+            raise ValueError("prof and prof2 must be equivalent")
 
         return prof._fourier_variance(cosmo, k, M, a, mass_def)

pyccl/halos/profiles_cib.py

@@ -86,7 +86,7 @@ def __init__(self, c_M_relation, nu_GHz, alpha=0.36, T0=24.4, beta=1.75,
                  gamma=1.7, s_z=3.6, log10meff=12.6, sigLM=0.707, Mmin=1E10,
                  L0=6.4E-8):
         if not isinstance(c_M_relation, Concentration):
-            raise TypeError("c_M_relation must be of type `Concentration`)")
+            raise TypeError("c_M_relation must be of type `Concentration`")
 
         self.nu = nu_GHz
         self.alpha = alpha

pyccl/tests/test_cclobject.py

@@ -119,15 +119,16 @@ def test_CCLHalosObject():
     CM2 = ccl.halos.ConcentrationDuffy08()
     assert CM1 == CM2
 
-    P1 = ccl.halos.HaloProfileHOD(c_M_relation=CM1)
-    P2 = ccl.halos.HaloProfileHOD(c_M_relation=CM2)
-    assert P1 == P2
+    # TODO: uncomment once __eq__ methods are implemented.
+    # P1 = ccl.halos.HaloProfileHOD(c_M_relation=CM1)
+    # P2 = ccl.halos.HaloProfileHOD(c_M_relation=CM2)
+    # assert P1 == P2
 
-    PCOV1 = ccl.halos.Profile2pt(r_corr=1.5)
-    PCOV2 = ccl.halos.Profile2pt(r_corr=1.0)
-    assert PCOV1 != PCOV2
-    PCOV2.update_parameters(r_corr=1.5)
-    assert PCOV1 == PCOV2
+    # PCOV1 = ccl.halos.Profile2pt(r_corr=1.5)
+    # PCOV2 = ccl.halos.Profile2pt(r_corr=1.0)
+    # assert PCOV1 != PCOV2
+    # PCOV2.update_parameters(r_corr=1.5)
+    # assert PCOV1 == PCOV2
 
 
 def test_CCLObject_immutable():
@@ -153,12 +154,13 @@ def test_CCLObject_immutable():
     prof.update_parameters(mass_bias=0.7)
     assert prof.mass_bias == 0.7
 
+    # TODO: uncomment once __eq__ methods are implemented.
     # Check that the hash repr is deleted as required.
-    prof1 = ccl.halos.HaloProfilePressureGNFW(mass_bias=0.5)
-    prof2 = ccl.halos.HaloProfilePressureGNFW(mass_bias=0.5)
-    assert prof1 == prof2                   # repr is cached
-    prof2.update_parameters(mass_bias=0.7)  # cached repr is deleted
-    assert prof1 != prof2                   # repr is cached again
+    # prof1 = ccl.halos.HaloProfilePressureGNFW(mass_bias=0.5)
+    # prof2 = ccl.halos.HaloProfilePressureGNFW(mass_bias=0.5)
+    # assert prof1 == prof2                   # repr is cached
+    # prof2.update_parameters(mass_bias=0.7)  # cached repr is deleted
+    # assert prof1 != prof2                   # repr is cached again
 
 
 def test_CCLObject_default_behavior():

pyccl/tests/test_pkhm.py

@@ -9,8 +9,8 @@
 M200 = ccl.halos.MassDef200m()
 HMF = ccl.halos.MassFuncTinker10(COSMO, mass_def=M200)
 HBF = ccl.halos.HaloBiasTinker10(COSMO, mass_def=M200)
-P1 = ccl.halos.HaloProfileNFW(ccl.halos.ConcentrationDuffy08(M200),
-                              fourier_analytic=True)
+CON = ccl.halos.ConcentrationDuffy08(M200)
+P1 = ccl.halos.HaloProfileNFW(CON, fourier_analytic=True)
 P2 = P1
 PKC = ccl.halos.Profile2pt()
 KK = np.geomspace(1E-3, 10, 32)
@@ -186,6 +186,48 @@ def test_pkhm_pk2d():
     assert np.all(np.fabs((pk_arr / pk_arr_2 - 1)).flatten()
                   < 1E-4)
 
+    # Testing profiles which are not equivalent (but very close)
+    G1 = ccl.halos.HaloProfileHOD(CON, lMmin_0=12.00000)
+    G2 = ccl.halos.HaloProfileHOD(CON, lMmin_0=11.99999)
+    assert G1 != G2
+
+    # I_1_1
+    pk0 = ccl.halos.halomod_power_spectrum(COSMO, hmc, k_arr, a_arr, G1,
+                                           prof2=G1, normprof1=False,
+                                           normprof2=False)
+    pk1 = ccl.halos.halomod_power_spectrum(COSMO, hmc, k_arr, a_arr, G1,
+                                           prof2=G2, normprof1=False,
+                                           normprof2=False)
+    assert np.allclose(pk1, pk0, rtol=1e-4)
+
+    # Profile normalization
+    pk0 = ccl.halos.halomod_power_spectrum(COSMO, hmc, k_arr, a_arr, G1,
+                                           prof2=G1, normprof1=True,
+                                           normprof2=True)
+    pk1 = ccl.halos.halomod_power_spectrum(COSMO, hmc, k_arr, a_arr, G1,
+                                           prof2=G2, normprof1=True,
+                                           normprof2=True)
+    assert np.allclose(pk1, pk0, rtol=1e-4)
+
+    # I_0_2 & I_1_2
+    assert np.allclose(hmc.I_0_2(COSMO, KK, AA, P1, prof_2pt=PKC),
+                       hmc.I_0_2(COSMO, KK, AA, P1, prof2=P1, prof_2pt=PKC),
+                       rtol=0)
+    assert np.allclose(hmc.I_1_2(COSMO, KK, AA, P1, prof_2pt=PKC),
+                       hmc.I_1_2(COSMO, KK, AA, P1, prof2=P1, prof_2pt=PKC),
+                       rtol=0)
+    # I_0_22
+    I0 = hmc.I_0_22(COSMO, KK, AA, P1, prof2=P1, prof3=P1, prof4=P1,
+                    prof12_2pt=PKC, prof34_2pt=PKC)
+    assert np.allclose(hmc.I_0_22(COSMO, KK, AA,
+                                  P1, prof2=P1, prof3=P1, prof4=P1,
+                                  prof12_2pt=PKC, prof34_2pt=None),
+                       I0, rtol=0)
+    assert np.allclose(hmc.I_0_22(COSMO, KK, AA,
+                                  P1, prof2=P1, prof3=None, prof4=None,
+                                  prof12_2pt=PKC, prof34_2pt=PKC),
+                       I0, rtol=0)
+
     # 1h/2h transition
     def alpha0(a):  # no smoothing
         return 1.
@@ -260,6 +302,10 @@ def test_pkhm_errors():
         ccl.halos.halomod_bias_1pt(COSMO, hmc, KK, AA, None)
     with pytest.raises(TypeError):
         ccl.halos.halomod_power_spectrum(COSMO, hmc, KK, AA, None)
+    with pytest.raises(TypeError):
+        ccl.halos.halomod_Tk3D_SSC(COSMO, hmc, P1,
+                                   prof2=P1, prof3=P1, prof4="hello",
+                                   normprof1=True)
 
     # Wrong prof2
     with pytest.raises(TypeError):

pyccl/tests/test_profiles.py

@@ -45,6 +45,45 @@ def smoke_assert_prof_real(profile, method='_real'):
         assert np.shape(p) == sh
 
 
+# def test_profiles_equal():
+#     M200m = ccl.halos.MassDef200m()
+#     cm = ccl.halos.ConcentrationDuffy08(mdef=M200m)
+#     p1 = ccl.halos.HaloProfileHOD(c_M_relation=cm, lMmin_0=12.)
+
+#     # different profile types
+#     p2 = ccl.halos.HaloProfilePressureGNFW()
+#     assert p1 != p2
+
+#     # equal profiles
+#     p2 = p1
+#     assert p1 == p2
+
+#     # equivalent profiles
+#     cm2 = ccl.halos.ConcentrationDuffy08(mdef=M200m)
+#     p2 = ccl.halos.HaloProfileHOD(c_M_relation=cm2, lMmin_0=12.)
+#     assert p1 == p2
+
+#     # different parameters
+#     p2 = ccl.halos.HaloProfileHOD(c_M_relation=cm, lMmin_0=11.)
+#     assert p1 != p2
+
+#     # different mass-concentration
+#     cm2 = ccl.halos.ConcentrationConstant()
+#     p2 = ccl.halos.HaloProfileHOD(c_M_relation=cm2, lMmin_0=12.)
+#     assert p1 != p2
+
+#     # different mass-concentration mass definition
+#     M200c = ccl.halos.MassDef200c()
+#     cm2 = ccl.halos.ConcentrationDuffy08(mdef=M200c)
+#     p2 = ccl.halos.HaloProfileHOD(c_M_relation=cm2, lMmin_0=12.)
+#     assert p1 != p2
+
+#     # different FFTLog
+#     p2 = ccl.halos.HaloProfileHOD(c_M_relation=cm, lMmin_0=12.)
+#     p2.update_precision_fftlog(**{"plaw_fourier": -2.0})
+#     assert p1 != p2
+
+
 @pytest.mark.parametrize('prof_class',
                          [ccl.halos.HaloProfileNFW,
                           ccl.halos.HaloProfileHernquist,
@@ -75,6 +114,11 @@ def test_empirical_smoke(prof_class):
     smoke_assert_prof_real(p, method='fourier')
 
 
+def test_empirical_smoke_CIB():
+    with pytest.raises(TypeError):
+        ccl.halos.HaloProfileCIBShang12(None, nu_GHz=417)
+
+
 def test_cib_smoke():
     c = ccl.halos.ConcentrationDuffy08(M200)
     p = ccl.halos.HaloProfileCIBShang12(c, 217)
@@ -206,23 +250,41 @@ def func(prof):
     assert p1.ns_independent is True
 
 
-def test_hod_2pt_raises():
+def test_hod_2pt():
     pbad = ccl.halos.HaloProfilePressureGNFW()
     c = ccl.halos.ConcentrationDuffy08(M200)
-    pgood = ccl.halos.HaloProfileHOD(c_M_relation=c)
-    pgood_b = ccl.halos.HaloProfileHOD(c_M_relation=c)
+    pgood = ccl.halos.HaloProfileHOD(c_M_relation=c, lM0_0=11.)
+    pgood_b = ccl.halos.HaloProfileHOD(c_M_relation=c, lM0_0=11.)
     p2 = ccl.halos.Profile2ptHOD()
+    F0 = p2.fourier_2pt(pgood, COSMO, 1., 1e13, 1.,
+                        prof2=pgood,
+                        mass_def=M200)
+    assert np.allclose(p2.fourier_2pt(pgood, COSMO, 1., 1e13, 1.,
+                                      prof2=None, mass_def=M200),
+                       F0, rtol=0)
+
     with pytest.raises(TypeError):
         p2.fourier_2pt(pbad, COSMO, 1., 1E13, 1.,
                        mass_def=M200)
 
-    with pytest.raises(ValueError):
-        p2.fourier_2pt(pgood, COSMO, 1., 1E13, 1.,
-                       prof2=pgood_b, mass_def=M200)
+    with pytest.raises(TypeError):
+        p2.fourier_2pt(pgood, COSMO, 1., 1e13, 1.,
+                       prof2=pbad,
+                       mass_def=M200)
 
+    # doesn't raise because profiles are equivalent
     p2.fourier_2pt(pgood, COSMO, 1., 1E13, 1.,
                    prof2=pgood, mass_def=M200)
 
+    # TODO: bring back when proper __eq__s are implemented
+    # p2.fourier_2pt(pgood, COSMO, 1., 1E13, 1.,
+    #                prof2=pgood_b, mass_def=M200)
+
+    with pytest.raises(ValueError):
+        pgood_b.update_parameters(lM0_0=10.)
+        p2.fourier_2pt(pgood, COSMO, 1., 1E13, 1.,
+                       prof2=pgood_b, mass_def=M200)
+
 
 def test_2pt_rcorr_smoke():
     c = ccl.halos.ConcentrationDuffy08(M200)

pyccl/tests/test_tkk1h.py

@@ -142,8 +142,6 @@ def test_tkk1h_tk3d():
 
 
 def test_tkk1h_errors():
-    from pyccl.pyutils import assert_warns
-
     hmc = ccl.halos.HMCalculator(COSMO, HMF, HBF, mass_def=M200)
     k_arr = KK
     a_arr = np.array([0.1, 0.4, 0.7, 1.0])
@@ -170,7 +168,8 @@ def test_tkk1h_errors():
                                          P1, prof34_2pt=P2)
 
     # Negative profile in logspace
-    assert_warns(ccl.CCLWarning, ccl.halos.halomod_Tk3D_1h,
-                 COSMO, hmc, P3, prof2=Pneg,
-                 lk_arr=np.log(k_arr), a_arr=a_arr,
-                 use_log=True)
+    with pytest.warns(ccl.CCLWarning):
+        ccl.halos.halomod_Tk3D_1h(COSMO, hmc, P3, prof2=Pneg,
+                                  prof3=P3, prof4=P3,
+                                  lk_arr=np.log(k_arr), a_arr=a_arr,
+                                  use_log=True)

pyccl/tests/test_tkkssc.py

@@ -149,10 +149,11 @@ def test_tkkssc_errors():
         ccl.halos.halomod_Tk3D_SSC(COSMO, hmc, P1, prof4=P2, normprof4=False)
 
     # Negative profile in logspace
-    assert_warns(ccl.CCLWarning, ccl.halos.halomod_Tk3D_1h,
-                 COSMO, hmc, P3, prof2=Pneg,
-                 lk_arr=np.log(k_arr), a_arr=a_arr,
-                 use_log=True)
+    with pytest.warns(ccl.CCLWarning):
+        ccl.halos.halomod_Tk3D_SSC(COSMO, hmc, P3, prof2=Pneg,
+                                   prof3=P3, prof4=P3,
+                                   lk_arr=np.log(k_arr), a_arr=a_arr,
+                                   use_log=True)
 
 
 @pytest.mark.parametrize('kwargs', [
@@ -235,7 +236,7 @@ def test_tkkssc_counterterms_gc(kwargs):
     if kwargs['prof3'] is None:
         kwargs['prof3'] = kwargs['prof1']
     if kwargs['prof4'] is None:
-        kwargs['prof4'] = kwargs['prof3']
+        kwargs['prof4'] = kwargs['prof2']
 
     # Tk's of the clustering terms
     tkc12 = []