comps2fields working for covs

Author: JaimeRZP

Diff for: examples/jackknife-covariance.ipynb

@@ -55,12 +55,10 @@ def Fields2Components(results):
             comps2 = _split_comps(key2)
             for i, comp1 in enumerate(comps1):
                 for j, comp2 in enumerate(comps2):
-                    #if i <= j:
-                    # Only save the upper triangle
                     _a1, _b1, _i1, _j1 = comp1
                     _a2, _b2, _i2, _j2 = comp2
-                    covkey = (_a1, _b1, _a2, _b2, _i1, _j1, _i2, _j2)
                     _r = r.array
+                    covkey = (_a1, _b1, _a2, _b2, _i1, _j1, _i2, _j2)
                     _results[covkey] = Result(_r[..., i, j, :, :], ell)
         else:
             raise ValueError(
@@ -69,6 +67,64 @@ def Fields2Components(results):
     return _results
 
 
+def Components2Fields(results):
+    _results = {}
+    keys = list(results.keys())
+    naxis = np.unique([len(result.axis) for result in results.values()])
+    if len(naxis) != 1:
+        raise ValueError("Different types of results in the same dictionary.")
+    naxis = naxis[0]
+    if naxis == 1:
+        # We are dealing with Cls
+        # find unique fields in comps
+        keys = set([_unsplit_comps(key) for key in keys])
+        keys = sorted(keys)
+        for key in keys:
+            # get comps of each unique field
+            comps = _split_comps(key)
+            cls = np.array([results[format_key(comp)] for comp in comps])
+            _results[key] = cls
+    elif naxis == 2:
+        # We are dealing with Covariance matrices
+        keys = list(set([(k[0], k[1], k[4], k[5]) for k in keys]))
+        keys = set([_unsplit_comps(key) for key in keys])
+        keys = sorted(keys)
+        for key1, key2 in itertools.combinations_with_replacement(keys, 2):
+            a1, b1, i1, j1 = key1
+            a2, b2, i2, j2 = key2
+            field_covkey = (a1, b1, a2, b2, i1, j1, i2, j2)
+            comps1 = _split_comps(key1)
+            comps2 = _split_comps(key2)
+            fields_cov = []
+            for _key1, _key2 in itertools.product(comps1, comps2):
+                _a1, _b1, _i1, _j1 = _key1
+                _a2, _b2, _i2, _j2 = _key2
+                comp_covkey = (_a1, _b1, _a2, _b2, _i1, _j1, _i2, _j2)
+                if comp_covkey in results.keys():
+                    comp_cov = results[comp_covkey]
+                    fields_cov.append(comp_cov)
+            if len(fields_cov) == 0:
+                # Happens if the BBAA exists but not AABB
+                comps1, comps2 = comps2, comps1
+                field_covkey = (a2, b2, a1, b1, i2, j2, i1, j1)
+                for _key1, _key2 in itertools.product(comps1, comps2):
+                    _a1, _b1, _i1, _j1 = _key1
+                    _a2, _b2, _i2, _j2 = _key2
+                    comp_covkey = (_a1, _b1, _a2, _b2, _i1, _j1, _i2, _j2)
+                    if comp_covkey in results.keys():
+                        comp_cov = results[comp_covkey]
+                        fields_cov.append(comp_cov)
+            # format the covariance
+            first, *rest = fields_cov
+            ell = first.ell
+            axis = first.axis
+            nell1, nell2 = first.shape
+            n, m = len(comps1), len(comps2)
+            fields_cov = np.reshape(fields_cov, (n, m, nell1, nell2))
+            _results[field_covkey] = Result(fields_cov, ell, axis=axis)
+    return _results
+
+
 def Components2Data(results):
     keys = list(results.keys())
     ells = [results[key].ell for key in keys]
@@ -138,66 +194,6 @@ def Data2Components(cov):
     return Cl_cov_dict
 
 
-def Components2Fields(results):
-    _results = {}
-    keys = list(results.keys())
-    naxis = np.unique([len(result.axis) for result in results.values()])
-    if len(naxis) != 1:
-        raise ValueError("Different types of results in the same dictionary.")
-    naxis = naxis[0]
-    if naxis == 1:
-        # We are dealing with Cls
-        # find unique fields in comps
-        keys = set([_unsplit_comps(key) for key in keys])
-        keys = sorted(keys)
-        for key in keys:
-            # get comps of each unique field
-            comps = _split_comps(key)
-            cls = np.array([results[format_key(comp)] for comp in comps])
-            _results[key] = cls
-    elif naxis == 2:
-        # We are dealing with Covariance matrices
-        keys = [(k[0], k[1], k[4], k[5]) for k in keys]
-        keys = list(set(keys))
-        for key1, key2 in itertools.combinations_with_replacement(keys, 2):
-            a1, b1, i1, j1 = key1
-            a2, b2, i2, j2 = key2
-            covkey = (a1, b1, a2, b2, i1, j1, i2, j2)
-            ells1, ells2 = results[covkey].ell
-            nells1, nells2 = len(ells1), len(ells2)
-            # Writes the covkeys of the spin components associated with covkey
-            # it also returns what fields go in what axis
-            #  comps1/comp2  (E, E) (B, B) (E,B)
-            # (POS, E)
-            # (POS, B)
-            comps1 = _split_comps(key1)
-            comps2 = _split_comps(key2)
-            # Save comps in dtype metadata
-            dt = np.dtype(
-                float,
-                metadata={
-                    "fields1": comps1,
-                    "fields2": comps2,
-                },
-            )
-            _cov = np.zeros((len(comps1), len(comps2), nells1, nells2), dtype=dt)
-            for i in range(len(comps1)):
-                for j in range(len(comps2)):
-                    comp1 = comps1[i]
-                    comp2 = comps2[j]
-                    _a1, _b1, _, _ = comp1
-                    _a2, _b2, _, _ = comp2
-                    _covkey = _a1, _b1, _a2, _b2, i1, j1, i2, j2
-                    if _covkey not in results.keys():
-                        # This triggers if the element doesn't exist
-                        # but the symmetrical term does
-                        _cov[i, j, :, :] = np.zeros((nells2, nells1))
-                    else:
-                        _cov[i, j, :, :] = results[_covkey]
-            _results[covkey] = Result(_cov, ell=(ells1, ells2))
-    return _results
-
-
 def _split_comps(key):
     _key = copy.deepcopy(key)
     a, b, i, j = _key

Diff for: heracles/dices/io.py

@@ -187,7 +187,15 @@ def correct_bias(cls, jkmaps, fields, jk=0, jk2=0):
     return cls
 
 
-def jackknife_covariance(samples, nd=1):
+def jackknife_covariance(dict, nd=1):
+    """
+    Compute the jackknife covariance matrix from a sequence
+    of spectra dictionaries *dict*.
+    """
+    return _jackknife_covariance(dict.values(), nd=nd)
+
+
+def _jackknife_covariance(samples, nd=1):
     """
     Compute the jackknife covariance matrix from a sequence
     of spectra dictionaries *samples*.
@@ -283,7 +291,7 @@ def delete2_correction(Cls0, Cls1, Cls2):
             _qii = Result(_qii, ell=Cls0[key].ell)
             qii[key] = _qii
             Q_ii.append(qii)
-    Q = jackknife_covariance(Q_ii, nd=2)
+    Q = _jackknife_covariance(Q_ii, nd=2)
     return Q
 
 

Diff for: heracles/dices/jackknife.py

@@ -247,7 +247,7 @@ def test_dices(data_path):
             assert nells == len(lgrid)
 
     # Delete1
-    cov_jk = dices.jackknife_covariance(cqs1.values())
+    cov_jk = dices.jackknife_covariance(cqs1)
     #  Shrinkage
     # target_cov = dices.gaussian_covariance(cqs0)
     # shrinkage = dices.shrinkage_factor(cqs0, cqs1, target_cov)
@@ -315,7 +315,7 @@ def test_dices(data_path):
     #     assert np.all(d == _d)
 
 
-def test_Fields2Components(data_path):
+def test_cls_io(data_path):
     data_maps = make_data_maps()
     fields = get_fields()
     jkmaps = make_jkmaps(data_path)
@@ -324,6 +324,26 @@ def test_Fields2Components(data_path):
     __cls = dices.Components2Fields(_cls)
     assert sorted(list(cls.keys())) == list(__cls.keys())
     for key in list(cls.keys()):
-        cl = cls[key]
-        __cl = __cls[key]
-        assert np.all(cl == __cl)
+        cl = cls[key].array
+        __cl = __cls[key].array
+        assert (cl == __cl).all()
+
+
+def test_cov_io(data_path):
+    nside = 128
+    data_maps = make_data_maps()
+    vis_maps = make_vis_maps()
+    fields = get_fields()
+    jkmaps = make_jkmaps(data_path)
+    cls1 = dices.jackknife_cls(data_maps, vis_maps, jkmaps, fields, nd=1)
+    lbins = 5
+    ledges = np.logspace(np.log10(10), np.log10(nside), lbins + 1)
+    cqs1 = heracles.binned(cls1, ledges)
+    cov_jk = dices.jackknife_covariance(cqs1)
+    _cov_jk = dices.Fields2Components(cov_jk)
+    __cov_jk = dices.Components2Fields(_cov_jk)
+    assert sorted(list(cov_jk.keys())) == sorted(list(__cov_jk.keys()))
+    for key in list(cov_jk.keys()):
+        cov = cov_jk[key].array
+        __cov = __cov_jk[key].array
+        assert (cov == __cov).all()