From e7af6dd46c6218db3a53ec8eb6520b295dcd3370 Mon Sep 17 00:00:00 2001
From: paulrogozenski <paulrogozenski@gmail.com>
Date: Wed, 18 Dec 2024 17:04:01 -0500
Subject: [PATCH] change default values for _get_sigma2_B, edited tests
 accordingly

---
 tests/data/conf_covariance_ssc_fsky.yaml  | 11 -----------
 tests/test_covariance_fourier_ssc_fsky.py |  2 ++
 tjpcov/covariance_fourier_ssc.py          | 11 ++++++-----
 tjpcov/covariance_fourier_ssc_fsky.py     |  6 +++---
 4 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/tests/data/conf_covariance_ssc_fsky.yaml b/tests/data/conf_covariance_ssc_fsky.yaml
index b63b1a0d..e49033ac 100644
--- a/tests/data/conf_covariance_ssc_fsky.yaml
+++ b/tests/data/conf_covariance_ssc_fsky.yaml
@@ -5,17 +5,6 @@ tjpcov:
   # 'set' from parameters OR pass CCL cosmology object OR yaml file
   cosmo: 'set'
 
-  # Setting mask OR fsky approximation
-  mask_file:
-    DESgc__0: ./tests/benchmarks/32_DES_tjpcov_bm/catalogs/mask_DESgc__0.fits.gz
-    DESwl__0: ./tests/benchmarks/32_DES_tjpcov_bm/catalogs/DESwlMETACAL_mask_zbin0_ns32.fits.gz
-    DESwl__1: ./tests/benchmarks/32_DES_tjpcov_bm/catalogs/DESwlMETACAL_mask_zbin1_ns32.fits.gz
-
-  mask_names:
-    DESgc__0: mask_DESgc0
-    DESwl__0: mask_DESwl0
-    DESwl__1: mask_DESwl1
-
   outdir: ./tests/tmp/
 
   # Survey params:
diff --git a/tests/test_covariance_fourier_ssc_fsky.py b/tests/test_covariance_fourier_ssc_fsky.py
index 4e143aec..998df10c 100644
--- a/tests/test_covariance_fourier_ssc_fsky.py
+++ b/tests/test_covariance_fourier_ssc_fsky.py
@@ -260,6 +260,8 @@ def test_get_covariance_block_WL_benchmark(cov_fssc):
 
     # Modify tjpcov instance
     cov_fssc.io.sacc_file = s
+    # Create an empty mask dict since there are no input mask files in the yaml
+    cov_fssc.mask_files = {}
     cov_fssc.mask_files["DESwl__0"] = mask
     cov_fssc.cosmo = cosmo
 
diff --git a/tjpcov/covariance_fourier_ssc.py b/tjpcov/covariance_fourier_ssc.py
index 25e01e14..64820363 100644
--- a/tjpcov/covariance_fourier_ssc.py
+++ b/tjpcov/covariance_fourier_ssc.py
@@ -27,7 +27,7 @@ def __init__(self, config):
 
         self.ssc_conf = self.config.get("SSC", {})
 
-    def _get_sigma2_B(self, cosmo, a_arr, masks={}):
+    def _get_sigma2_B(self, cosmo, a_arr, tr=None):
         """Returns the variance of the projected linear density field,
             for the case when a survey mask is provided.
 
@@ -36,8 +36,8 @@ def _get_sigma2_B(self, cosmo, a_arr, masks={}):
             a_arr (:obj:`float`, `array` or :obj:`None`): an array of
                 scale factor values at which to evaluate
                 the projected variance.
-            masks (:obj:`dict`): dictionary containing the survey
-                masks of the relevant tracers.
+            tr (:obj:`dict`): dictionary containing the
+                tracer name combinations.
 
         Returns:
             - (:obj:`float` or `array`): projected variance.
@@ -46,6 +46,8 @@ def _get_sigma2_B(self, cosmo, a_arr, masks={}):
         # TODO: Optimize this, avoid computing the mask_wl for all blocks.
         # Note that this is correct for same footprint cross-correlations. In
         # case of multisurvey analyses this approximation might break.
+        masks = self.get_masks_dict(tr, {})
+
         m12 = masks[1] * masks[2]
         m34 = masks[3] * masks[4]
         area = hp.nside2pixarea(hp.npix2nside(m12.size))
@@ -172,8 +174,7 @@ def get_covariance_block(
             is_number_counts4=isnc[4],
         )
 
-        masks = self.get_masks_dict(tr, {})
-        sigma2_B = self._get_sigma2_B(cosmo, a, masks=masks)
+        sigma2_B = self._get_sigma2_B(cosmo, a, tr=tr)
 
         ell = self.get_ell_eff()
         cov_ssc = ccl.covariances.angular_cl_cov_SSC(
diff --git a/tjpcov/covariance_fourier_ssc_fsky.py b/tjpcov/covariance_fourier_ssc_fsky.py
index 5b0727db..30b74cef 100644
--- a/tjpcov/covariance_fourier_ssc_fsky.py
+++ b/tjpcov/covariance_fourier_ssc_fsky.py
@@ -28,7 +28,7 @@ def __init__(self, config):
                 "You need to set fsky for FourierSSCHaloModelFsky"
             )
 
-    def _get_sigma2_B(self, cosmo, a_arr, masks={}):
+    def _get_sigma2_B(self, cosmo, a_arr, tr=None):
         """Returns the variance of the projected linear density field,
             for the fsky/disk approximation case.
 
@@ -37,8 +37,8 @@ def _get_sigma2_B(self, cosmo, a_arr, masks={}):
             a_arr (:obj:`float`, `array` or :obj:`None`): an array of
                 scale factor values at which to evaluate
                 the projected variance.
-            masks (:obj:`dict`): dictionary containing the survey
-                masks of the relevant tracers.
+            tr (:obj:`dict`): dictionary containing the
+                tracer name combinations.
         Returns:
             - (:obj:`float` or `array`): projected variance.
         """