gh-395: Theta max (#394)

Author: JaimeRZP

examples/unmixing.ipynb

@@ -7,6 +7,8 @@
     # Python < 3.13
     from dataclasses import replace
 
+from .result import get_result_array
+
 gauss_legendre = None
 _gauss_legendre_cache = {}
 
@@ -203,6 +205,9 @@ def cl2corr(cls):
         s1, s2 = cl.spin
         # Grab metadata
         dtype = cl.array.dtype
+        # Determine lmax from ell field or shape along ell axis
+        lmax = len(get_result_array(cl, "ell")[0]) - 1
+        xvals, _ = _cached_gauss_legendre(lmax + 1)
         # Initialize wd
         wd = np.zeros_like(cl)
         if (s1 != 0) and (s2 != 0):
@@ -259,6 +264,7 @@ def cl2corr(cls):
         wd = np.array(list(wd), dtype=dtype)
         wds[key] = replace(
             cls[key],
+            ell=xvals,
             array=wd,
         )
     return wds
@@ -278,6 +284,9 @@ def corr2cl(wds):
         s1, s2 = wd.spin
         # Grab metadata
         dtype = wd.array.dtype
+        # Derive lmax from xvals stored in the correlation's ell field
+        xvals = get_result_array(wd, "ell")[0]
+        lmax = len(xvals) - 1
         # initialize cl
         cl = np.zeros_like(wd)
         if (s1 != 0) and (s2 != 0):
@@ -335,6 +344,7 @@ def corr2cl(wds):
         cl = np.array(list(cl), dtype=dtype)
         cls[key] = replace(
             wds[key],
+            ell=np.arange(lmax + 1),
             array=cl,
         )
     return cls

heracles/transforms.py

@@ -17,10 +17,10 @@
 # You should have received a copy of the GNU Lesser General Public
 # License along with Heracles. If not, see <https://www.gnu.org/licenses/>.
 import numpy as np
-from collections.abc import Mapping
 from .result import binned
 from .transforms import cl2corr, corr2cl
 from .utils import get_cl
+from .transforms import _cached_gauss_legendre
 
 try:
     from copy import replace
@@ -29,38 +29,36 @@
     from dataclasses import replace
 
 
-def naturalspice(d, m, fields, rcond=0.01):
+def naturalspice(d, m, fields, theta_max=None):
     """
     Natural unmixing of the data Cl.
     Args:
         d: Data Cl
         m: mask Cl
         fields: list of fields
-        patch_hole: If True, apply the patch hole correction
+        theta_max: maximum angle to use for the unmixing, in degrees. If None, use all angles.
     Returns:
         corr_d: Corrected Cl
     """
     first_wd = list(d.values())[0]
     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)
-    for m_key in list(wm.keys()):
-        if isinstance(rcond, Mapping):
-            if m_key not in rcond:
-                raise KeyError(f"Missing rcond value for wm key: {m_key}")
-            _rcond = rcond[m_key]
-        else:
-            _rcond = rcond
-        _wm = wm[m_key].array
-        _wm = _wm * logistic(np.log10(abs(_wm)), x0=np.log10(_rcond * np.max(_wm)))
-        wm[m_key] = replace(wm[m_key], array=_wm)
-
+    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)
 
     # trnasform back to Cl
@@ -99,5 +97,5 @@ def _naturalspice(wd, wm, fields):
     return corr_wds
 
 
-def logistic(x, x0=-5, k=50):
+def logistic(x, x0=-2, k=50):
     return 1.0 + np.exp(-k * (x - x0))

heracles/unmixing.py

@@ -117,23 +117,12 @@ def test_full_mask_correction(cls0, mls0, fields):
         assert np.isclose(cl[2:], _cl[2:]).all()
 
     cls_alphas = heracles.corr2cl(alphas)
-    __cls = heracles.unmixing.naturalspice(cls0, cls_alphas, fields, rcond=1.0)
+    __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()
 
-    wm_keys = heracles.transforms.cl2corr(cls_alphas).keys()
-    rcond_by_key = {key: 1.0 for key in wm_keys}
-    ___cls = heracles.unmixing.naturalspice(
-        cls0,
-        cls_alphas,
-        fields,
-        rcond=rcond_by_key,
-    )
-    for key in list(cls0.keys()):
-        np.testing.assert_allclose(__cls[key].array, ___cls[key].array)
-
     _alphas = dices.get_mask_correlation_ratio(mls0, mls0, unmixed=True)
     for key in list(_alphas.keys()):
         wmls0 = heracles.transforms._cl2corr(mls0[key]).T[0]
@@ -142,17 +131,6 @@ def test_full_mask_correction(cls0, mls0, fields):
         assert np.isclose(alpha, _alpha).all()
 
 
-def test_full_mask_correction_rcond_missing_key(cls0, mls0, fields):
-    cls_alphas = heracles.corr2cl(
-        dices.get_mask_correlation_ratio(mls0, mls0, unmixed=False)
-    )
-    wm_keys = list(heracles.transforms.cl2corr(cls_alphas).keys())
-    rcond_by_key = {key: 1.0 for key in wm_keys[1:]}
-
-    with pytest.raises(KeyError, match="Missing rcond value for wm key"):
-        heracles.unmixing.naturalspice(cls0, cls_alphas, fields, rcond=rcond_by_key)
-
-
 def test_fast_mask_correction(cls0, fields, jk_maps):
     _cls0 = dices.correct_footprint_reduction(cls0, jk_maps, fields, 0, 0)
     for key in list(cls0.keys()):