gh-376: Transforms to correlations

heracles/dices/jackknife.py

@@ -26,7 +26,7 @@
 from ..mapping import transform
 from ..twopoint import angular_power_spectra
 from ..unmixing import _natural_unmixing, logistic
-from ..transforms import cl2corr
+from ..transforms import _cl2corr
 
 try:
     from copy import replace
@@ -265,13 +265,13 @@ def get_mask_correlation_ratio(Mljk, Mls0, unmixed=False):
         mljk = Mljk[key]
         mls0 = Mls0[key]
         # Transform to real space
-        wmljk = cl2corr(mljk)
+        wmljk = _cl2corr(mljk)
         wmljk = wmljk.T[0]
         wmljk *= logistic(np.log10(abs(wmljk)))
         # Compute alpha
         alpha = wmljk
         if not unmixed:
-            wmls0 = cl2corr(mls0)
+            wmls0 = _cl2corr(mls0)
             wmls0 = wmls0.T[0]
             alpha /= wmls0
         alphas[key] = replace(Mls0[key], array=alpha)

heracles/transforms.py

@@ -1,6 +1,12 @@
 from scipy.special import lpn as legendrep
 import numpy as np
 
+try:
+    from copy import replace
+except ImportError:
+    # Python < 3.13
+    from dataclasses import replace
+
 gauss_legendre = None
 _gauss_legendre_cache = {}
 
@@ -91,7 +97,7 @@ def legendre_funcs(lmax, x, m=(0, 2), lfacs=None, lfacs2=None, lrootfacs=None):
     return res
 
 
-def cl2corr(cls, lmax=None, sampling_factor=1):
+def _cl2corr(cls, lmax=None, sampling_factor=1):
     """
     Get the correlation function from the power spectra, evaluated at points cos(theta) = xvals.
     Use roots of Legendre polynomials (np.polynomial.legendre.leggauss) for accurate back integration with corr2cl.
@@ -137,11 +143,11 @@ def cl2corr(cls, lmax=None, sampling_factor=1):
         corrs[i, 1] = np.dot(cp, d22)  # Q+U
         corrs[i, 2] = np.dot(cm, d2m2)  # Q-U
         corrs[i, 3] = np.dot(cc, d20)  # cross
-    corrs[:, 0] += corrs[:, 0][0] / (4 * np.pi)
+    # corrs[:, 0] += corrs[:, 0][0] / (4 * np.pi)
     return corrs
 
 
-def corr2cl(corrs, lmax=None, sampling_factor=1):
+def _corr2cl(corrs, lmax=None, sampling_factor=1):
     """
     Transform from correlation functions to power spectra.
     Note that using cl2corr followed by corr2cl is generally very accurate (< 1e-5 relative error) if
@@ -155,6 +161,11 @@ def corr2cl(corrs, lmax=None, sampling_factor=1):
       They include :math:`\ell(\ell+1)/2\pi` factors.
     """
 
+    if corrs.ndim == 1:
+        corrs = np.array(
+            [corrs, np.zeros_like(corrs), np.zeros_like(corrs), np.zeros_like(corrs)]
+        ).T
+
     if lmax is None:
         lmax = corrs.shape[0] - 1
 
@@ -177,6 +188,157 @@ def corr2cl(corrs, lmax=None, sampling_factor=1):
         cls[2:, 2] += T2 - T4
         cls[2:, 3] += (weight * corrs[i, 3]) * d20
 
-    cls[1, :] *= 2
-    cls[2:, :] = cls[2:, :]
+    # cls[1, :] *= 2
+    # cls[2:, :] = cls[2:, :]
     return 2 * np.pi * cls
+
+
+def cl2corr(cls):
+    """
+    Transforms cls to correlation functions
+    Args:
+        cls: Data Cl
+    Returns:
+        corr: correlation function
+    """
+    wds = {}
+    for key in cls.keys():
+        cl = cls[key]
+        s1, s2 = cl.spin
+        # Grab metadata
+        dtype = cl.array.dtype
+        # Initialize wd
+        wd = np.zeros_like(cl)
+        if (s1 != 0) and (s2 != 0):
+            _cl = np.array(
+                [
+                    np.zeros_like(cl[0, 0]),
+                    cl[0, 0],  # EE like spin-2
+                    cl[1, 1],  # BB like spin-2
+                    np.zeros_like(cl[0, 0]),
+                ]
+            )
+            _icl = np.array(
+                [
+                    np.zeros_like(cl[0, 0]),
+                    -cl[0, 1],  # EB like spin-0
+                    cl[1, 0],  # EB like spin-0
+                    np.zeros_like(cl[0, 0]),
+                ]
+            )
+            # transform to corrs
+            _wd = _cl2corr(_cl.T).T + 1j * _cl2corr(_icl.T).T
+            _rwd = _wd.real
+            _iwd = _wd.imag
+            # reorder
+            wd[0, 0] = _rwd[1]  # EE like spin-2
+            wd[1, 1] = _rwd[2]  # BB like spin-2
+            wd[0, 1] = _iwd[1]  # EB like spin-0
+            wd[1, 0] = _iwd[2]  # EB like spin-0
+        elif (s1 != 0) or (s2 != 0):
+            _clp = np.array(
+                [
+                    np.zeros_like(cl[0]),
+                    np.zeros_like(cl[0]),
+                    np.zeros_like(cl[0]),
+                    cl[0] + cl[1],  # TE like spin-2
+                ]
+            )
+            _clm = np.array(
+                [
+                    np.zeros_like(cl[0]),
+                    np.zeros_like(cl[0]),
+                    np.zeros_like(cl[0]),
+                    cl[0] - cl[1],  # TE like spin-2
+                ]
+            )
+            # trnsform to corrs
+            wd[0] = _cl2corr(_clp.T).T[3]
+            wd[1] = _cl2corr(_clm.T).T[3]
+        elif (s1 == 0) and (s2 == 0):
+            wd = _cl2corr(cl).T[0]
+        else:
+            raise ValueError("Invalid spin combination")
+        # Add metadata back
+        wd = np.array(list(wd), dtype=dtype)
+        wds[key] = replace(
+            cls[key],
+            array=wd,
+        )
+    return wds
+
+
+def corr2cl(wds):
+    """
+    Transforms correlation functions to cls
+    Args:
+        corrs: data corrs
+    Returns:
+        corr: correlation function
+    """
+    cls = {}
+    for key in wds.keys():
+        wd = wds[key]
+        s1, s2 = wd.spin
+        # Grab metadata
+        dtype = wd.array.dtype
+        # initialize cl
+        cl = np.zeros_like(wd)
+        if (s1 != 0) and (s2 != 0):
+            _rwd = np.array(
+                [
+                    np.zeros_like(wd[0, 0]),
+                    wd[0, 0],  # EE like spin-2
+                    wd[1, 1],  # BB like spin-2
+                    np.zeros_like(wd[0, 0]),
+                ]
+            )
+            _iwd = np.array(
+                [
+                    np.zeros_like(wd[0, 0]),
+                    wd[0, 1],  # EB like spin-0
+                    wd[1, 0],  # EB like spin-0
+                    np.zeros_like(wd[0, 0]),
+                ]
+            )
+            # transform back to Cl
+            _rcl = _corr2cl(_rwd.T).T
+            _icl = _corr2cl(_iwd.T).T
+            # reorder
+            cl[0, 0] = _rcl[1]  # EE like spin-2
+            cl[1, 1] = _rcl[2]  # BB like spin-2
+            cl[0, 1] = -_icl[1]  # EB like spin-0
+            cl[1, 0] = _icl[2]  # EB like spin-0
+        elif (s1 != 0) or (s2 != 0):
+            _wp = np.array(
+                [
+                    np.zeros_like(wd[0]),
+                    np.zeros_like(wd[0]),
+                    np.zeros_like(wd[0]),
+                    wd[0],  # TE like spin-2
+                ]
+            )
+            _wm = np.array(
+                [
+                    np.zeros_like(wd[0]),
+                    np.zeros_like(wd[0]),
+                    np.zeros_like(wd[0]),
+                    wd[1],  # TE like spin-2
+                ]
+            )
+            _clp = _corr2cl(_wp.T).T[3]
+            _clm = _corr2cl(_wm.T).T[3]
+            cl[0] = (_clp + _clm) / 2
+            cl[1] = (_clp - _clm) / 2
+        elif (s1 == 0) and (s2 == 0):
+            # Treat everything as spin-0 and preserve 1D shape.
+            cl = _corr2cl(wd).T[0]
+        else:
+            raise ValueError("Invalid spin combination")
+        # Add metadata back
+        cl = np.array(list(cl), dtype=dtype)
+        cls[key] = replace(
+            wds[key],
+            array=cl,
+        )
+    return cls

heracles/unmixing.py

@@ -18,7 +18,7 @@
 # License along with Heracles. If not, see <https://www.gnu.org/licenses/>.
 import numpy as np
 from .result import binned
-from .transforms import cl2corr, corr2cl
+from .transforms import _cl2corr, _corr2cl
 from .utils import get_cl
 
 try:
@@ -43,7 +43,7 @@ def natural_unmixing(d, m, fields, x0=-2, k=50, patch_hole=True, lmax=None):
     m_keys = list(m.keys())
     for m_key in m_keys:
         _m = m[m_key].array
-        _wm = cl2corr(_m).T[0]
+        _wm = _cl2corr(_m).T[0]
         if patch_hole:
             _wm *= logistic(np.log10(abs(_wm)), x0=x0, k=k)
         wm[m_key] = replace(m[m_key], array=_wm)
@@ -98,12 +98,12 @@ def _natural_unmixing(d, wm, fields, lmax=None):
                 ]
             )
             # Correct by alpha
-            wd = cl2corr(__d.T).T + 1j * cl2corr(__id.T).T
+            wd = _cl2corr(__d.T).T + 1j * _cl2corr(__id.T).T
             corr_wd = (wd / _wm).real
             icorr_wd = (wd / _wm).imag
             # Transform back to Cl
-            __corr_d = corr2cl(corr_wd.T).T
-            __icorr_d = corr2cl(icorr_wd.T).T
+            __corr_d = _corr2cl(corr_wd.T).T
+            __icorr_d = _corr2cl(icorr_wd.T).T
             # reorder
             _corr_d = np.zeros_like(_d)
             _corr_d[0, 0] = __corr_d[1]  # EE like spin-2
@@ -128,23 +128,23 @@ def _natural_unmixing(d, wm, fields, lmax=None):
                 ]
             )
             # Correct by alpha
-            wplus = cl2corr(__dp.T).T
-            wminus = cl2corr(__dm.T).T
+            wplus = _cl2corr(__dp.T).T
+            wminus = _cl2corr(__dm.T).T
             corr_wplus = wplus / _wm
             corr_wminus = wminus / _wm
             # Transform back to Cl
-            corr_dp = corr2cl(corr_wplus.T).T
-            corr_dm = corr2cl(corr_wminus.T).T
+            corr_dp = _corr2cl(corr_wplus.T).T
+            corr_dm = _corr2cl(corr_wminus.T).T
             # reorder
             _corr_d = np.zeros_like(_d)
             _corr_d[0] = 0.5 * (corr_dp[3] + corr_dm[3])  # TE
             _corr_d[1] = 0.5 * (corr_dp[3] - corr_dm[3])  # TB
         elif (s1 == 0) and (s2 == 0):
             # Treat everything as spin-0
-            wd = cl2corr(_d).T
+            wd = _cl2corr(_d).T
             corr_wd = wd / _wm
             # Transform back to Cl
-            _corr_d = corr2cl(corr_wd.T).T[0]
+            _corr_d = _corr2cl(corr_wd.T).T[0]
         else:
             raise ValueError(f"Invalid spin combination: {s1}, {s2}")
         # Add metadata back

tests/test_dices.py

@@ -111,7 +111,7 @@ def test_full_mask_correction(cls0, mls0, fields):
 
     _alphas = dices.get_mask_correlation_ratio(mls0, mls0, unmixed=True)
     for key in list(_alphas.keys()):
-        wmls0 = heracles.transforms.cl2corr(mls0[key]).T[0]
+        wmls0 = heracles.transforms._cl2corr(mls0[key]).T[0]
         alpha = alphas[key].array
         _alpha = _alphas[key].array / wmls0
         assert np.isclose(alpha, _alpha).all()
@@ -136,8 +136,8 @@ def test_polspice(cls0):
             get_cl(("POS", "SHE", 1, 1), cls0)[0],
         ]
     ).T
-    corrs = heracles.cl2corr(cls)
-    _cls = heracles.corr2cl(corrs)
+    corrs = heracles.transforms._cl2corr(cls)
+    _cls = heracles.transforms._corr2cl(corrs)
     for cl, _cl in zip(cls.T, _cls.T):
         assert np.isclose(cl[2:], _cl[2:]).all()
 

tests/test_transforms.py

@@ -2,29 +2,14 @@
 import heracles
 
 
-def test_cl2corr():
-    # Is there something more clever we can do here?
-    # Like transforming the legendre nodes and return ones?
-    cl = np.array(
-        [
-            np.ones(10),
-            np.ones(10),
-            np.ones(10),
-            np.zeros(10),
-        ]
-    )
-    corr = heracles.cl2corr(cl.T).T
-    assert corr.shape == cl.shape
-
-
-def test_corr2cl():
-    corr = np.array(
-        [
-            np.ones(10),
-            np.zeros(10),
-            np.zeros(10),
-            np.zeros(10),
-        ]
-    )
-    cl = heracles.corr2cl(corr.T).T
-    assert corr.shape == cl.shape
+def test_transforms(cls0):
+    corr = heracles.cl2corr(cls0)
+    for key in list(cls0.keys()):
+        print(key, cls0[key].shape, corr[key].shape)
+        assert corr[key].shape == cls0[key].shape
+
+    _cls = heracles.corr2cl(corr)
+    for key in list(cls0.keys()):
+        print(key, cls0[key].shape, _cls[key].shape)
+        assert _cls[key].shape == cls0[key].shape
+        assert np.isclose(cls0[key].array, _cls[key].array).all()