diff --git a/heracles/dices/jackknife.py b/heracles/dices/jackknife.py index 937d28d..5632dd8 100644 --- a/heracles/dices/jackknife.py +++ b/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) diff --git a/heracles/transforms.py b/heracles/transforms.py index e6d44a0..a784f99 100644 --- a/heracles/transforms.py +++ b/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 diff --git a/heracles/unmixing.py b/heracles/unmixing.py index 1ea5551..7557bd5 100644 --- a/heracles/unmixing.py +++ b/heracles/unmixing.py @@ -18,7 +18,7 @@ # License along with Heracles. If not, see . 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 diff --git a/tests/test_dices.py b/tests/test_dices.py index 40f0175..5759b2b 100644 --- a/tests/test_dices.py +++ b/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() diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 6513713..39279f0 100644 --- a/tests/test_transforms.py +++ b/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()