Theta bin average for Real-space covariance

environment.yml

@@ -2,7 +2,7 @@ name: tjpcov
 channels: 
   - conda-forge
 dependencies:
-  - python>=3.10
+  - python>=3.10,<3.14
   - pip
   - wheel
   - black

pyproject.toml

@@ -13,7 +13,7 @@ authors = [
 ]
 description = "Covariances for LSST DESC"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.10,<3.14"
 classifiers = [
     "Development Status :: 4 - Beta",
     "Programming Language :: Python :: 3",

run_tjpcov.py

@@ -2,7 +2,6 @@
 
 from tjpcov.covariance_calculator import CovarianceCalculator
 
-
 if __name__ == "__main__":
     import argparse
 

tests/test_covariance_calculator.py

@@ -8,7 +8,6 @@
 import sacc
 import shutil
 
-
 INPUT_YML = "./tests/data/conf_covariance_calculator.yml"
 OUTDIR = "./tests/tmp/"
 

tests/test_covariance_fourier_gaussian_nmt.py

@@ -13,7 +13,6 @@
 from tjpcov.covariance_fourier_gaussian_nmt import FourierGaussianNmt
 from tjpcov.covariance_io import CovarianceIO
 
-
 ROOT = "tests/benchmarks/32_DES_tjpcov_bm/"
 OUTDIR = "tests/tmp/"
 INPUT_YML = "tests/data/conf_covariance_gaussian_fourier_nmt.yaml"

tests/test_covariance_gaussian_fsky.py

@@ -118,13 +118,14 @@ def test_Fourier_get_covariance_block(cov_fg_fsky, mock_cosmo):
             tracer_comb1=trs, tracer_comb2=trs, for_real=True
         )
     # 2. Check block
-    cov2 = cov_fg_fsky.get_covariance_block(
+    cov2, SN = cov_fg_fsky.get_covariance_block(
         tracer_comb1=trs, tracer_comb2=trs, for_real=True, lmax=30
     )
     ell = np.arange(30 + 1)
     ccltr = ccl_tracers["src0"]
     cl = ccl.angular_cl(mock_cosmo, ccltr, ccltr, ell) + tracer_noise["src0"]
     cov = np.diag(2 * cl**2)
+    cov2 = cov2 + np.diag(np.ones_like(cl) * SN)
     assert cov2.shape == (ell.size, ell.size)
     np.testing.assert_allclose(cov2, cov)
 
@@ -150,13 +151,14 @@ def test_Fourier_get_covariance_block(cov_fg_fsky, mock_cosmo):
 def test_Real_get_fourier_block(
     cov_rg_fsky, cov_fg_fsky, tracer_comb1, tracer_comb2
 ):
-    cov = cov_rg_fsky._get_fourier_block(tracer_comb1, tracer_comb2)
-    cov2 = cov_fg_fsky.get_covariance_block(
+    cov, SN = cov_rg_fsky._get_fourier_block(tracer_comb1, tracer_comb2)
+    cov2, SN2 = cov_fg_fsky.get_covariance_block(
         tracer_comb1, tracer_comb2, for_real=True, lmax=cov_rg_fsky.lmax
     )
 
     norm = np.pi * 4 * cov_rg_fsky.fsky
     assert np.all(cov == cov2 / norm)
+    assert np.all(SN == SN2 / norm)
 
 
 def test_smoke_get_covariance(cov_fg_fsky, cov_rg_fsky):

tests/test_covariance_real_base.py

@@ -5,7 +5,7 @@
 
 from functools import partial
 
-from tjpcov.wigner_transform import bin_cov, WignerTransform
+from tjpcov.wigner_transform import WignerTransform
 from tjpcov.covariance_builder import (
     CovarianceProjectedReal,
     CovarianceReal,
@@ -115,7 +115,7 @@ def test_get_Wigner_transform(cov_prj_real):
     wt = cov_prj_real.get_Wigner_transform()
 
     assert isinstance(wt, WignerTransform)
-    assert np.all(wt.ell == np.arange(2, cov_prj_real.lmax + 1))
+    assert np.all(wt.ell == np.arange(0, cov_prj_real.lmax + 1))
     assert np.all(wt.theta == cov_prj_real.get_binning_info()[0])
     assert wt.s1_s2s == [(2, 2), (2, -2), (0, 2), (2, 0), (0, 0)]
 
@@ -158,12 +158,12 @@ def test_build_matrix_from_blocks(cov_prj_real):
 )
 def test_get_covariance_block(cov_prj_real, tracer_comb1, tracer_comb2):
     lmax = cov_prj_real.lmax
-    ell = np.arange(2, lmax + 1)
+    ell = np.arange(0, lmax + 1)
     fourier_block = np.random.rand(lmax + 1, lmax + 1)
 
     # Dynamically override the method on this instance.
     def override_fourier_block(self, tracer_comb1, tracer_comb2):
-        return fourier_block
+        return fourier_block, None
 
     cov_prj_real._get_fourier_block = partial(
         override_fourier_block, cov_prj_real
@@ -180,23 +180,15 @@ def override_fourier_block(self, tracer_comb1, tracer_comb2):
         ell_cl=ell,
         s1_s2=s1_s2_1,
         s1_s2_cross=s1_s2_2,
-        cl_cov=fourier_block[2:][:, 2:],
+        cl_cov=fourier_block,
     )
 
     gcov_xi_1 = cov_prj_real.get_covariance_block(
-        tracer_comb1=tracer_comb1, tracer_comb2=tracer_comb2, binned=False
+        tracer_comb1=tracer_comb1, tracer_comb2=tracer_comb2
     )
 
     assert np.max(np.abs((gcov_xi_1 + 1e-100) / (cov + 1e-100) - 1)) < 1e-5
-
-    gcov_xi_1 = cov_prj_real.get_covariance_block(
-        tracer_comb1=tracer_comb1, tracer_comb2=tracer_comb2, binned=True
-    )
-
-    theta, _, theta_edges = cov_prj_real.get_binning_info(in_radians=False)
-    _, cov = bin_cov(r=theta, r_bins=theta_edges, cov=cov)
     assert gcov_xi_1.shape == (20, 20)
-    assert np.max(np.abs((gcov_xi_1 + 1e-100) / (cov + 1e-100) - 1)) < 1e-5
 
 
 @pytest.mark.parametrize(

tests/test_wigner_transform.py

@@ -7,11 +7,13 @@
 
 def get_WT_kwargs():
     lmax = 96
-    ell = np.arange(2, lmax + 1)
-    theta = np.sort(np.pi / ell)[::2]  # Force them to have different sizes
+    ell = np.arange(0, lmax + 1)
+    theta = np.sort(np.pi / ell[1:])[::2]  # Force them to have different sizes
+    theta_edges = np.append(theta, [np.max(theta) * 1.05])
     WT_kwargs = {
         "ell": ell,
         "theta": theta,
+        "theta_edges": theta_edges,
         "s1_s2": [(2, 2), (2, -2), (0, 2), (2, 0), (0, 0)],
     }
     return WT_kwargs
@@ -103,14 +105,9 @@ def test_projected_covariance(s1_s2, s1_s2_cross):
         wt.projected_covariance(wt.ell[10:], mat, s1_s2, s1_s2_cross)
 
     th, matb = wt.projected_covariance(wt.ell, mat, s1_s2, s1_s2_cross)
-    wd_a = wigner_transform.wigner_d(*s1_s2, wt.theta, wt.ell)
-    wd_b = wigner_transform.wigner_d(*s1_s2_cross, wt.theta, wt.ell)
-    matb_2 = (
-        (wd_a * np.sqrt(wt.norm) * wt.grad_ell)
-        @ mat
-        @ (wd_b * np.sqrt(wt.norm)).T
-    )
-
+    wd_a = wt.wig_d[*s1_s2]
+    wd_b = wt.wig_d[*s1_s2_cross]
+    matb_2 = (wd_a * wt.grad_ell) @ mat @ wd_b.T
     assert np.all(th == wt.theta)
     assert np.max(np.abs(matb / matb_2) - 1) < 1e-5
 

tjpcov/covariance_builder.py

@@ -6,7 +6,7 @@
 import pyccl as ccl
 import sacc
 
-from .wigner_transform import bin_cov, WignerTransform
+from .wigner_transform import WignerTransform
 from . import tools
 from .covariance_io import CovarianceIO
 
@@ -995,6 +995,7 @@ def __init__(self, config):
                 "You need to specify the lmax you want to "
                 "compute the Fourier covariance up to"
             )
+        self.cov_type = None
 
     @property
     def fourier(self):
@@ -1086,13 +1087,14 @@ def get_Wigner_transform(self):
             :obj:`~tjpcov.wigner_transform.WignerTransform` instance
         """
         if self.WT is None:
-            # Removing ell <= 1 (following original implementation)
-            ell = np.arange(2, self.lmax + 1)
-            theta, _, _ = self.get_binning_info(in_radians=True)
+            # Removing ell <= 1 is done in legendre.py
+            ell = np.arange(0, self.lmax + 1)
+            theta, _, theta_edges = self.get_binning_info(in_radians=True)
 
             WT_kwargs = {
                 "ell": ell,
                 "theta": theta,
+                "theta_edges": theta_edges,
                 "s1_s2": [(2, 2), (2, -2), (0, 2), (2, 0), (0, 0)],
             }
 
@@ -1110,7 +1112,6 @@ def get_covariance_block(
         tracer_comb2,
         xi_plus_minus1="plus",
         xi_plus_minus2="plus",
-        binned=True,
     ):
         """Compute a single covariance matrix for a given pair of xi.
 
@@ -1127,7 +1128,7 @@ def get_covariance_block(
         """
         # For now we just use the EE block which should be dominant over the
         # EB, BE and BB pieces
-        cov = self._get_fourier_block(tracer_comb1, tracer_comb2)
+        cov, SN = self._get_fourier_block(tracer_comb1, tracer_comb2)
 
         WT = self.get_Wigner_transform()
 
@@ -1137,15 +1138,60 @@ def get_covariance_block(
             s1_s2_1 = s1_s2_1[xi_plus_minus1]
         if isinstance(s1_s2_2, dict):
             s1_s2_2 = s1_s2_2[xi_plus_minus2]
-        # Remove ell <= 1 for WT (following original implementation)
-        ell = np.arange(2, self.lmax + 1)
-        cov = cov[2:][:, 2:]
+
+        # Project sample variance term and mixed term.
+        # Removing ell <= 1 is done in legendre.py
+        ell = np.arange(0, self.lmax + 1)
         th, cov = WT.projected_covariance(
             ell_cl=ell, s1_s2=s1_s2_1, s1_s2_cross=s1_s2_2, cl_cov=cov
         )
-        if binned:
-            theta, _, theta_edges = self.get_binning_info(in_radians=False)
-            thb, cov = bin_cov(r=theta, r_bins=theta_edges, cov=cov)
+
+        if self.cov_type == "gauss":
+            # denominator in average
+            dcost = np.cos(WT.theta_edges[1:]) - np.cos(WT.theta_edges[:-1])
+
+            if SN is not None:
+                # Load Npair
+                # see https://github.com/LSSTDESC/TXPipe/blob/
+                # a9dfdb7809ac7ed6c162fd3930c643a67afcd881/txpipe/twopoint_plots.py#L215
+                if tracer_comb1 == tracer_comb2:
+                    sacc_file = self.io.get_sacc_file()
+                    data_type = self.get_tracer_comb_data_types(tracer_comb1)[
+                        0
+                    ]
+                    D = sacc_file.get_data_points(
+                        data_type, (tracer_comb1[0], tracer_comb1[1])
+                    )
+                    Npair = np.array([d.get_tag("npair") for d in D])
+
+                    if Npair[0] is None:
+                        # assuming no survey boundaries.
+                        if np.abs(s1_s2_1[0]) == np.abs(s1_s2_1[1]) == 2:
+                            SN *= 2
+
+                        cov_sn = SN / (np.pi * dcost)
+                    else:
+                        # catalog level N_pair from treecorr
+                        T_sn = 1
+                        if tracer_comb1[0] in self.sigma_e:
+                            T_sn *= self.sigma_e[tracer_comb1[0]] ** 2
+                        if tracer_comb1[1] in self.sigma_e:
+                            T_sn *= self.sigma_e[tracer_comb1[1]] ** 2
+
+                        if (tracer_comb1[0] in self.sigma_e) and (
+                            tracer_comb1[1] in self.sigma_e
+                        ):
+                            T_sn *= 2
+                        # Eq. 64 of https://arxiv.org/abs/2410.06962
+                        cov_sn = 2 * T_sn / Npair
+
+                    # Add pure shot/shape noise contribution.
+                    cov += np.diag(cov_sn)
+
+            else:
+                # TODO: Projection of NaMaster covariance via
+                # Eq. 67 of https://arxiv.org/abs/2012.08568
+                pass
 
         return cov
 

tjpcov/covariance_gaussian_fsky.py

@@ -50,10 +50,8 @@ def get_binning_info(self, binning="linear"):
             ell_edges = out[2]
             return ell, ell_eff, ell_edges
         else:
-            warnings.warn(
-                "No bandpower windows found, \
-                           falling back to linear method"
-            )
+            warnings.warn("No bandpower windows found, \
+                           falling back to linear method")
             ell_eff = self.get_ell_eff()
             nbpw = ell_eff.size
 
@@ -160,16 +158,23 @@ def get_covariance_block(
             else 0
         )
 
+        if for_real:
+            cov = np.diag(
+                (cl[13] + SN[13]) * (cl[24] + SN[24])
+                + (cl[14] + SN[14]) * (cl[23] + SN[23])
+                - np.ones_like(cl[13]) * SN[13] * SN[24]
+                - np.ones_like(cl[14]) * SN[14] * SN[23]
+            )
+            # If it is to compute the real space covariance, return the
+            # covariance before binning or normalizing
+            # pure shot/shape noise term will be computed separately
+            return cov, SN[13] * SN[24] + SN[14] * SN[23]
+
         cov = np.diag(
             (cl[13] + SN[13]) * (cl[24] + SN[24])
             + (cl[14] + SN[14]) * (cl[23] + SN[23])
         )
 
-        if for_real:
-            # If it is to compute the real space covariance, return the
-            # covariance before binning or normalizing
-            return cov
-
         norm = (2 * ell + 1) * np.gradient(ell) * self.fsky
         cov /= norm
 
@@ -193,7 +198,7 @@ def get_covariance_block(
 class RealGaussianFsky(CovarianceProjectedReal):
     """Class to compute the Real space Gaussian cov. with the Knox formula.
 
-    It projects the the Fourier space Gaussian covariance into the real space.
+    It projects the Fourier space Gaussian covariance into the real space.
     """
 
     cov_type = "gauss"
@@ -214,6 +219,7 @@ def __init__(self, config):
         # sacc file.
         self.fourier = FourierGaussianFsky(config)
         self.fsky = self.fourier.fsky
+        self.cov_type = "gauss"
 
     def _get_fourier_block(self, tracer_comb1, tracer_comb2):
         """Return the Fourier covariance block for two pair of tracers.
@@ -227,9 +233,9 @@ def _get_fourier_block(self, tracer_comb1, tracer_comb2):
         """
         # For now we just use the EE block which should be dominant over the
         # EB, BE and BB pieces when projecting to real space
-        cov = self.fourier.get_covariance_block(
+        cov, SN = self.fourier.get_covariance_block(
             tracer_comb1, tracer_comb2, for_real=True, lmax=self.lmax
         )
         norm = np.pi * 4 * self.fsky
 
-        return cov / norm
+        return cov / norm, SN / norm