gh-397: Dices lmax (#398)

Author: JaimeRZP

heracles/dices/__init__.py

@@ -26,8 +26,8 @@
     "jackknife_fsky",
     "jackknife_bias",
     "correct_bias",
-    "get_mask_correlation_ratio",
-    "correct_footprint_reduction",
+    "correct_footprint_naturalspice",
+    "correct_footprint_fsky",
     "jackknife_covariance",
     "debias_covariance",
     "delete2_correction",
@@ -46,8 +46,8 @@
     jackknife_fsky,
     jackknife_bias,
     correct_bias,
-    get_mask_correlation_ratio,
-    correct_footprint_reduction,
+    correct_footprint_naturalspice,
+    correct_footprint_fsky,
     jackknife_covariance,
     debias_covariance,
     delete2_correction,

heracles/dices/jackknife.py

@@ -22,11 +22,11 @@
 from itertools import combinations
 from ..utils import add_to_Cls, sub_to_Cls
 from ..core import update_metadata
-from ..result import Result, get_result_array
+from ..result import Result, get_result_array, binned
 from ..mapping import transform
 from ..twopoint import angular_power_spectra
-from ..unmixing import _naturalspice, logistic
-from ..transforms import _cl2corr, cl2corr, corr2cl
+from ..unmixing import _naturalspice
+from ..transforms import cl2corr, corr2cl
 
 try:
     from copy import replace
@@ -76,12 +76,11 @@ def jackknife_cls(
         if mask_correction == "Full":
             vis_alms_jk = _sum_alms_except(vis_alms_regions, regions)
             _cls_mm = angular_power_spectra(vis_alms_jk)
-            alphas = get_mask_correlation_ratio(_cls_mm, mls0, unmixed=unmixed)
-            _wcls = cl2corr(_cls)
-            _wcls = _naturalspice(_wcls, alphas, fields)
-            _cls = corr2cl(_wcls)
+            _cls = correct_footprint_naturalspice(
+                _cls, _cls_mm, mls0, fields, unmixed=unmixed
+            )
         elif mask_correction == "Fast":
-            _cls = correct_footprint_reduction(
+            _cls = correct_footprint_fsky(
                 _cls, jk_maps, fields, *regions, unmixed=unmixed
             )
         else:
@@ -220,7 +219,7 @@ def correct_bias(cls, jkmaps, fields, jk=0, jk2=0):
     return cls
 
 
-def correct_footprint_reduction(cls, jkmaps, fields, jk=0, jk2=0, unmixed=False):
+def correct_footprint_fsky(cls, jkmaps, fields, jk=0, jk2=0, unmixed=False):
     """
     Corrects the Cls for the footprint reduction due to taking out a region.
     inputs:
@@ -249,33 +248,42 @@ def correct_footprint_reduction(cls, jkmaps, fields, jk=0, jk2=0, unmixed=False)
     return _cls
 
 
-def get_mask_correlation_ratio(Mljk, Mls0, unmixed=False):
+def _mask_correlation_ratio(mljk, mls0, unmixed=False):
+    alphas = {}
+    wmls0 = cl2corr(mls0)
+    wmljk = cl2corr(mljk)
+    for key in list(wmljk.keys()):
+        _wmljk = wmljk[key].array
+        _wmls0 = wmls0[key].array
+        alpha = _wmljk
+        if not unmixed:
+            alpha = alpha / _wmls0
+        alphas[key] = replace(mls0[key], array=alpha)
+    return alphas
+
+
+def correct_footprint_naturalspice(cls, cls_mm, mls0, fields, unmixed=False):
     """
-    Computes the ratio of the correlation
-    functions of the masks Cls.
-    input:
-        Mljk (np.array): mask of delete1 Cls
-        Mls0 (np.array): mask Cls
+    Corrects the Cls for footprint reduction using the full NaMaster/naturalspice approach.
+    inputs:
+        cls (dict): Dictionary of data Cls
+        cls_mm (dict): Dictionary of jackknife mask Cls
+        mls0 (dict): Dictionary of full mask Cls
+        fields (dict): Dictionary of fields
         unmixed (bool): unmix the Cls
     returns:
-        alpha (Float64): Mask correction factor
+        cls (dict): Corrected Cls
     """
-    alphas = {}
-    for key in list(Mljk.keys()):
-        mljk = Mljk[key]
-        mls0 = Mls0[key]
-        # Transform to real space
-        wmljk = _cl2corr(mljk)
-        wmljk = wmljk.T[0]
-        wmljk *= logistic(np.log10(abs(wmljk)))
-        # Compute alpha
-        alpha = wmljk
-        if not unmixed:
-            wmls0 = _cl2corr(mls0)
-            wmls0 = wmls0.T[0]
-            alpha /= wmls0
-        alphas[key] = replace(Mls0[key], array=alpha)
-    return alphas
+    alphas = _mask_correlation_ratio(cls_mm, mls0, unmixed=unmixed)
+    first_cls = list(cls.values())[0]
+    first_mls = list(mls0.values())[0]
+    lmax = first_cls.shape[first_cls.axis[0]]
+    lmax_mask = first_mls.shape[first_mls.axis[0]]
+    cls = binned(cls, np.arange(0, lmax_mask + 1))
+    wcls = cl2corr(cls)
+    wcls = _naturalspice(wcls, alphas, fields)
+    cls = corr2cl(wcls)
+    return binned(cls, np.arange(0, lmax + 1))
 
 
 def jackknife_covariance(dict, nd=1):

heracles/unmixing.py

@@ -29,6 +29,10 @@
     from dataclasses import replace
 
 
+def logistic(x, x0=-2, k=50):
+    return 1.0 + np.exp(-k * (x - x0))
+
+
 def naturalspice(d, m, fields, theta_max=None):
     """
     Natural unmixing of the data Cl.
@@ -44,22 +48,13 @@ def naturalspice(d, m, fields, theta_max=None):
     first_wm = list(m.values())[0]
     lmax = first_wd.shape[first_wd.axis[0]]
     lmax_mask = first_wm.shape[first_wm.axis[0]]
-    xvals, _ = _cached_gauss_legendre(lmax_mask)
-    theta = np.arccos(xvals) * 180 / np.pi
 
     # pad correlation functions to lmax_mask
     d = binned(d, np.arange(0, lmax_mask + 1))
 
     wd = cl2corr(d)
     wm = cl2corr(m)
-    if theta_max is not None:
-        for m_key in list(wm.keys()):
-            _wm = wm[m_key].array
-            i_theta_max = np.abs(theta - theta_max).argmin()
-            wm_at_theta_max = _wm[i_theta_max]
-            _wm = _wm * logistic(np.log10(abs(_wm)), x0=np.log10(abs(wm_at_theta_max)))
-            wm[m_key] = replace(wm[m_key], array=_wm)
-    corr_wds = _naturalspice(wd, wm, fields)
+    corr_wds = _naturalspice(wd, wm, fields, theta_max=theta_max)
 
     # trnasform back to Cl
     corr_d = corr2cl(corr_wds)
@@ -69,14 +64,14 @@ def naturalspice(d, m, fields, theta_max=None):
     return corr_d
 
 
-def _naturalspice(wd, wm, fields):
+def _naturalspice(wd, wm, fields, theta_max=None):
     """
     Natural unmixing of the data correlation function.
     Args:
         wd: data correlation function
         wm: mask correlation function
         fields: list of fields
-        patch_hole: If True, apply the patch hole correction
+        theta_max: maximum angle in degrees for the logistic cutoff. If None, uses default x0=-2.
     Returns:
         corr_d: Corrected Cl
     """
@@ -85,17 +80,22 @@ def _naturalspice(wd, wm, fields):
         if field.mask is not None:
             masks[key] = field.mask
 
+    if theta_max is not None:
+        first_wm = list(wm.values())[0]
+        lmax_mask = first_wm.shape[first_wm.axis[0]]
+        xvals, _ = _cached_gauss_legendre(lmax_mask)
+        theta = np.arccos(xvals) * 180 / np.pi
+        i_theta_max = np.abs(theta - theta_max).argmin()
+
     corr_wds = {}
     for key in wd.keys():
         a, b, i, j = key
         m_key = (masks[a], masks[b], i, j)
-        _wm = get_cl(m_key, wm)
-        _wd = wd[key]
-        # divide by the mask correlation function
-        corr_wds[key] = replace(wd[key], array=(_wd.array / _wm.array))
+        _wm = get_cl(m_key, wm).array
+        _wd = wd[key].array
+        if theta_max is not None:
+            x0 = np.log10(abs(_wm[i_theta_max]))
+            _wm *= logistic(np.log10(abs(_wm)), x0=x0)
+        corr_wds[key] = replace(wd[key], array=(_wd / _wm))
 
     return corr_wds
-
-
-def logistic(x, x0=-2, k=50):
-    return 1.0 + np.exp(-k * (x - x0))

tests/test_dices.py

@@ -107,23 +107,24 @@ def test_get_delete2_fsky(jk_maps, njk):
 
 
 def test_full_mask_correction(cls0, mls0, fields):
-    alphas = dices.get_mask_correlation_ratio(mls0, mls0, unmixed=False)
-    wcls0 = heracles.transforms.cl2corr(cls0)
-    _wcls = heracles.unmixing._naturalspice(wcls0, alphas, fields)
-    _cls = heracles.transforms.corr2cl(_wcls)
+    from heracles.dices.jackknife import _mask_correlation_ratio
+
+    # When mljk == mls0, correct_footprint_naturalspice should recover the original cls
+    _cls = dices.correct_footprint_naturalspice(cls0, mls0, mls0, fields, unmixed=False)
     for key in list(cls0.keys()):
         cl = cls0[key].array
         _cl = _cls[key].array
         assert np.isclose(cl[2:], _cl[2:]).all()
 
+    alphas = _mask_correlation_ratio(mls0, mls0, unmixed=False)
     cls_alphas = heracles.corr2cl(alphas)
     __cls = heracles.unmixing.naturalspice(cls0, cls_alphas, fields, theta_max=180)
     for key in list(cls0.keys()):
         cl = cls0[key].array
         _cl = __cls[key].array
         assert np.isclose(cl[2:], 2 * _cl[2:]).all()
 
-    _alphas = dices.get_mask_correlation_ratio(mls0, mls0, unmixed=True)
+    _alphas = _mask_correlation_ratio(mls0, mls0, unmixed=True)
     for key in list(_alphas.keys()):
         wmls0 = heracles.transforms._cl2corr(mls0[key]).T[0]
         alpha = alphas[key].array
@@ -132,7 +133,7 @@ def test_full_mask_correction(cls0, mls0, fields):
 
 
 def test_fast_mask_correction(cls0, fields, jk_maps):
-    _cls0 = dices.correct_footprint_reduction(cls0, jk_maps, fields, 0, 0)
+    _cls0 = dices.correct_footprint_fsky(cls0, jk_maps, fields, 0, 0)
     for key in list(cls0.keys()):
         cl = cls0[key].array
         _cl = _cls0[key].array