Getting rid of components

heracles/dices/__init__.py

@@ -39,6 +39,7 @@
     "flatten",
     # utils
     "impose_correlation",
+    "get_cl",
 ]
 
 from .jackknife import (

heracles/dices/io.py

@@ -51,25 +51,31 @@ def _fields2components(results):
         if len(axis) == 1:
             # We are dealing with Cls
             a1, b1, i1, j1 = key
-            _a1, idx1 = _split_key(a1)
-            _b1, idx2 = _split_key(b1)
+            sa1, sb1 = r.spin
+            _a1, idx1 = _split_key(a1, sa1)
+            _b1, idx2 = _split_key(b1, sb1)
             for k, idx in zip(
                 itertools.product(_a1, _b1), itertools.product(idx1, idx2)
             ):
                 __a1, __b1 = k
                 _key = (__a1, __b1, i1, j1)
                 _r = r[idx]
                 _r = np.squeeze(_r)
-                duplicate_cond = __a1 == "G_B" and __b1 == "G_E" and i1 == j1
+                if sa1 != 0 and sb1 != 0:
+                    mode1, mode2 = __a1[-1], __b1[-1]
+                    duplicate_cond = mode1 == "B" and mode2 == "E" and i1 == j1
+                else:
+                    duplicate_cond = False
                 if not duplicate_cond:
-                    _results[_key] = Result(_r, ell=ell)
+                    _results[_key] = Result(_r, spin=(0, 0), axis=(0,), ell=ell)
         elif len(axis) == 2:
             # We are dealing with Covariance matrices
             a1, b1, a2, b2, i1, j1, i2, j2 = key
-            _a1, idx1 = _split_key(a1)
-            _b1, idx2 = _split_key(b1)
-            _a2, idx3 = _split_key(a2)
-            _b2, idx4 = _split_key(b2)
+            sa1, sb1, sa2, sb2 = r.spin
+            _a1, idx1 = _split_key(a1, sa1)
+            _b1, idx2 = _split_key(b1, sb1)
+            _a2, idx3 = _split_key(a2, sa2)
+            _b2, idx4 = _split_key(b2, sb2)
             for k, idx in zip(
                 itertools.product(_a1, _b1, _a2, _b2),
                 itertools.product(idx1, idx2, idx3, idx4),
@@ -78,7 +84,7 @@ def _fields2components(results):
                 _key = (__a1, __b1, __a2, __b2, i1, j1, i2, j2)
                 _r = r[idx]
                 _r = np.squeeze(_r)
-                _results[_key] = Result(_r, ell=ell)
+                _results[_key] = Result(_r, spin=(0, 0, 0, 0), axis=(0, 1), ell=ell)
         else:
             raise ValueError(
                 "Results with more than 3 axes are not supported at the moment."
@@ -109,11 +115,18 @@ def _components2data(results, order=None):
             order = []
             nells = []
             for key in list(results.keys()):
+                # The order only depends on the unique fields
+                # So we only need the first two entries of the key
+                s1, s2, _, _ = results[key].spin
                 ell = results[key].ell
                 nell = len(ell[0])
                 _key = (key[0], key[1], key[4], key[5])
                 a, b, i, j = _key
-                duplicate_cond = a == "G_B" and b == "G_E" and i == j
+                if s1 != 0 and s2 != 0:
+                    mode1, mode2 = _key[-1], _key[-1]
+                    duplicate_cond = mode1 == "B" and mode2 == "E" and i == j
+                else:
+                    duplicate_cond = False
                 if _key not in order and not duplicate_cond:
                     order.append(_key)
                     nells.append(nell)
@@ -159,26 +172,8 @@ def _components2data(results, order=None):
     return data
 
 
-def _split_key(f, pos=None):
-    if f == "POS":
-        return ["POS"], [pos]
-    if f == "SHE":
-        return ["G_E", "G_B"], [0, 1]
-
-
-def format_key(key):
-    """
-    Produces a Cl key for data maps.
-    input:
-        key: Cl key
-    returns:
-        Clkey: Cl key
-    """
-    _key = copy.deepcopy(key)
-    a, b, i, j = _key
-    if i > j:
-        i, j = j, i
-        a, b = b, a
-    if (b == "POS") or (b == "G_E" and a == "G_B"):
-        a, b = b, a
-    return (a, b, i, j)
+def _split_key(f, spin, pos=None):
+    if spin == 0:
+        return [f], [pos]
+    if spin != 0:
+        return [f + "_E", f + "_B"], [0, 1]

heracles/dices/jackknife.py

@@ -28,6 +28,12 @@
 from ..unmixing import _natural_unmixing, logistic
 from ..transforms import cl2corr
 
+try:
+    from copy import replace
+except ImportError:
+    # Python < 3.13
+    from dataclasses import replace
+
 
 def jackknife_cls(data_maps, vis_maps, jk_maps, fields, nd=1):
     """
@@ -97,7 +103,7 @@ def jackknife_maps(maps, jkmaps, jk=0, jk2=0):
         _map = _maps[key_data]
         _jkmap = jkmaps[key_mask]
         _mask = np.copy(_jkmap)
-        _mask[_mask != 0] = _mask[_mask != 0] / _mask[_mask != 0]
+        _mask = (_mask > 0).astype(int)
         # Remove jk 2 regions
         cond = np.where((_jkmap == float(jk)) | (_jkmap == float(jk2)))[0]
         _mask[cond] = 0.0
@@ -135,7 +141,7 @@ def jackknife_fsky(jkmaps, jk=0, jk2=0):
     for key in jkmaps.keys():
         jkmap = jkmaps[key]
         mask = np.copy(jkmap)
-        mask[mask != 0] = mask[mask != 0] / mask[mask != 0]
+        mask = (mask > 0).astype(int)
         fsky = sum(mask) / len(mask)
         cond = np.where((mask == 1.0) & (jkmap != jk) & (jkmap != jk2))[0]
         rel_fskys[key] = (len(cond) / len(mask)) / fsky
@@ -191,7 +197,7 @@ def correct_bias(cls, jkmaps, fields, jk=0, jk2=0):
     for key in cls.keys():
         cl = cls[key].array
         update_metadata(cl, bias=b_jk[key])
-        cls[key] = Result(cl)
+        cls[key] = replace(cls[key], array=cl)
     return cls
 
 
@@ -241,6 +247,8 @@ def _jackknife_covariance(samples, nd=1):
         # get reference results
         result1 = first[key1]
         result2 = first[key2]
+        sa1, sb1 = result1.spin
+        sa2, sb2 = result2.spin
         # gather samples for this key combination
         samples1 = np.stack([result1] + [spectra[key1] for spectra in rest])
         samples2 = np.stack([result2] + [spectra[key2] for spectra in rest])
@@ -267,7 +275,7 @@ def _jackknife_covariance(samples, nd=1):
             # add extra axis if needed
             a1, b1, i1, j1 = key1
             a2, b2, i2, j2 = key2
-            result = Result(a, axis=axis, ell=ell)
+            result = Result(a, axis=axis, spin=(sa1, sb1, sa2, sb2), ell=ell)
             # store result
             cov[a1, b1, a2, b2, i1, j1, i2, j2] = result
     return cov
@@ -321,7 +329,7 @@ def delete2_correction(cls0, cls1, cls2):
             _qii -= (Njk - 1) * cls1[(k1,)][key].array
             _qii -= (Njk - 1) * cls1[(k2,)][key].array
             _qii += (Njk - 2) * cls2[kk][key].array
-            _qii = Result(_qii)
+            _qii = replace(cls0[key], array=_qii)
             qii[key] = _qii
         Q_ii.append(qii)
     # Compute the correction from the ensemble
@@ -334,7 +342,7 @@ def delete2_correction(cls0, cls1, cls2):
         q_diag_exp = np.zeros_like(q)
         diag_indices = np.arange(length)  # Indices for the diagonal
         q_diag_exp[..., diag_indices, diag_indices] = q_diag
-        Q[key] = Result(q_diag_exp, axis=q.axis, ell=q.ell)
+        Q[key] = replace(q, array=q_diag_exp)
     return Q
 
 
@@ -361,9 +369,5 @@ def _debias_covariance(cov_jk, Q):
     debiased_cov = {}
     for key in list(cov_jk.keys()):
         c = cov_jk[key].array - Q[key].array
-        debiased_cov[key] = Result(
-            c,
-            ell=cov_jk[key].ell,
-            axis=cov_jk[key].axis,
-        )
+        debiased_cov[key] = replace(cov_jk[key], array=c)
     return debiased_cov

heracles/dices/shrinkage.py

@@ -28,14 +28,20 @@
 from .utils import (
     add_to_Cls,
     impose_correlation,
+    get_cl,
 )
 from .io import (
     _fields2components,
     flatten,
-    format_key,
     _split_key,
 )
 
+try:
+    from copy import replace
+except ImportError:
+    # Python < 3.13
+    from dataclasses import replace
+
 
 def shrink(cov, target, shrinkage_factor):
     """
@@ -53,7 +59,7 @@ def shrink(cov, target, shrinkage_factor):
         c = cov[key].array
         tc = correlated_target[key].array
         sc = shrinkage_factor * tc + (1 - shrinkage_factor) * c
-        shrunk_cov[key] = Result(sc, axis=cov[key].axis, ell=cov[key].ell)
+        shrunk_cov[key] = replace(cov[key], array=sc)
     return shrunk_cov
 
 
@@ -120,11 +126,13 @@ def gaussian_covariance(Cls):
         # get reference results
         cl1 = Cls[key1]
         cl2 = Cls[key2]
+        sa1, sb1 = cl1.spin
+        sa2, sb2 = cl2.spin
         # get components
-        _a1, idx1 = _split_key(a1, pos=0)
-        _b1, idx2 = _split_key(b1, pos=0)
-        _a2, idx3 = _split_key(a2, pos=0)
-        _b2, idx4 = _split_key(b2, pos=0)
+        _a1, idx1 = _split_key(a1, sa1, pos=0)
+        _b1, idx2 = _split_key(b1, sb1, pos=0)
+        _a2, idx3 = _split_key(a2, sa2, pos=0)
+        _b2, idx4 = _split_key(b2, sb2, pos=0)
         # get attributes of result
         ell1 = get_result_array(cl1, "ell")
         ell2 = get_result_array(cl2, "ell")
@@ -145,7 +153,7 @@ def gaussian_covariance(Cls):
         # Remove the extra dimensions
         r = np.squeeze(r)
         # Make Result
-        result = Result(r, ell=ell)
+        result = Result(r, spin=(sa1, sb1, sa2, sb2), ell=ell)
         cov[covkey] = result
     return cov
 
@@ -160,15 +168,10 @@ def _gaussian_covariance(cls, key):
         cov: covariance matrix
     """
     a1, b1, a2, b2, i1, j1, i2, j2 = key
-    clkey1 = format_key((a1, a2, i1, i2))
-    clkey2 = format_key((b1, b2, j1, j2))
-    clkey3 = format_key((a1, b2, i1, j2))
-    clkey4 = format_key((b1, a2, j1, i2))
-    cl1 = cls[clkey1].array
-    cl2 = cls[clkey2].array
-    cl3 = cls[clkey3].array
-    cl4 = cls[clkey4].array
-    # Compute the Gaussian covariance
+    cl1 = get_cl((a1, a2, i1, i2), cls)
+    cl2 = get_cl((b1, b2, j1, j2), cls)
+    cl3 = get_cl((a1, b2, i1, j2), cls)
+    cl4 = get_cl((b1, a2, j1, i2), cls)
     cov = cl1 * cl2 + cl3 * cl4
     return cov
 

heracles/dices/utils.py

@@ -17,10 +17,43 @@
 # You should have received a copy of the GNU Lesser General Public
 # License along with DICES. If not, see <https://www.gnu.org/licenses/>.
 import numpy as np
-from ..result import Result
 
+try:
+    from copy import replace
+except ImportError:
+    # Python < 3.13
+    from dataclasses import replace
 
-def add_to_Cls(Cls, x):
+
+def get_cl(key, cls):
+    """
+    Internal method to get a Cl from a dictionary of Cls.
+    Check if the key exists if not tries to find the symmetric key.
+    input:
+        key: key of the Cl
+        cls: dictionary of Cls
+    returns:
+        cl: Cl
+    """
+    if key in cls:
+        return cls[key].array
+    else:
+        a, b, i, j = key
+        key_sym = (b, a, j, i)
+        if key_sym in cls:
+            arr = cls[key_sym].array
+            s1, s2 = cls[key_sym].spin
+            if s1 != 0 and s2 != 0:
+                print("dims of arr:", key_sym, arr.shape)
+                return np.transpose(arr, axes=(1, 0, 2))
+            else:
+                return arr
+
+        else:
+            raise KeyError(f"Key {key} not found in Cls.")
+
+
+def add_to_Cls(cls, x):
     """
     Adds a dictionary of Cl values to another.
     input:
@@ -29,14 +62,14 @@ def add_to_Cls(Cls, x):
     returns:
         Cls: updated dictionary of Cl values
     """
-    _Cls = {}
-    for key in Cls.keys():
-        ell = Cls[key].ell
-        _Cls[key] = Result(Cls[key].array + x[key], ell)
-    return _Cls
+    _cls = {}
+    for key in cls.keys():
+        arr = cls[key].array + x[key]
+        _cls[key] = replace(cls[key], array=arr)
+    return _cls
 
 
-def sub_to_Cls(Cls, x):
+def sub_to_Cls(cls, x):
     """
     Substracts a dictionary of Cl values to another.
     input:
@@ -45,11 +78,11 @@ def sub_to_Cls(Cls, x):
     returns:
         Cls: updated dictionary of Cl values
     """
-    _Cls = {}
-    for key in Cls.keys():
-        ell = Cls[key].ell
-        _Cls[key] = Result(Cls[key].array - x[key], ell)
-    return _Cls
+    _cls = {}
+    for key in cls.keys():
+        arr = cls[key].array - x[key]
+        _cls[key] = replace(cls[key], array=arr)
+    return _cls
 
 
 def impose_correlation(cov_a, cov_b):
@@ -71,5 +104,5 @@ def impose_correlation(cov_a, cov_b):
         b_std = np.sqrt(b_v[..., None, :])
         c = a * (b_std * np.swapaxes(b_std, -1, -2))
         c /= a_std * np.swapaxes(a_std, -1, -2)
-        cov_c[key] = Result(c, axis=a.axis, ell=a.ell)
+        cov_c[key] = replace(a, array=c)
     return cov_c