Update from main

heracles/dices/jackknife.py

@@ -41,7 +41,7 @@
 def jackknife_cls(
     data_maps,
     vis_maps,
-    jk_maps,
+    jk_map,
     fields,
     mask_correction="Fast",
     unmixed=False,
@@ -54,7 +54,7 @@ def jackknife_cls(
     inputs:
         data_maps (dict): Dictionary of data maps
         vis_maps (dict): Dictionary of visibility maps
-        jk_maps (dict): Dictionary of mask maps
+        jk_map (array): Jackknife mask map
         fields (dict): Dictionary of fields
         mask_correction (str): Type of mask correction to apply ("Fast" or "Full")
         nd (int): Number of Jackknife regions
@@ -63,99 +63,200 @@ def jackknife_cls(
     returns:
         cls (dict): Dictionary of data Cls
     """
-    if nd < 0 or nd > 2:
-        raise ValueError("number of deletions must be 0, 1, or 2")
+    """Alms calculated and save then cls calculated from saved alms."""
+
+    if progress is None:
+        progress = NoProgress()
+
+    # calc save alms if don't exist
+    compute_jk_alms(
+        data_maps,
+        vis_maps,
+        jk_map,
+        fields,
+        dir=dir,
+        progress=progress,
+    )
+
+    # calculate cls from saved alms
+    return compute_jk_cls_from_alms(
+        jk_map,
+        fields,
+        mask_correction=mask_correction,
+        unmixed=unmixed,
+        nd=nd,
+        dir=dir,
+        progress=progress,
+    )
+
+
+def compute_jk_alms(
+    data_maps,
+    vis_maps,
+    jk_map,
+    fields,
+    dir="./dices",
+    progress=None,
+):
+    """Compute and save ALMs each JK region."""
 
     if progress is None:
         progress = NoProgress()
 
-    cls = {}
-    jkmap = jk_maps[list(jk_maps.keys())[0]]
-    njk = len(np.unique(jkmap)[np.unique(jkmap) != 0])
     os.makedirs(dir, exist_ok=True)
 
-    all_regions = list(combinations(range(1, njk + 1), nd))
-    total = (njk + 1) + len(all_regions)
+    njk = len(np.unique(jk_map)[np.unique(jk_map) != 0])
+
+    total = njk + 1
     current = 0
     progress.update(current, total)
 
-    # Compute ALMs
     for k in range(0, njk + 1):
-        data_path = os.path.join(dir, f"data_alms_{k}.fits")
-        vis_path = os.path.join(dir, f"vis_alms_{k}.fits")
         with progress.task(f"ALMs {k}"):
-            if not (os.path.exists(data_path) and os.path.exists(vis_path)):
-                if k == 0:
-                    data_alms_k = transform(fields, data_maps)
-                    vis_alms_k = transform(fields, vis_maps)
-                else:
-                    data_alms_k = transform(
-                        fields, _get_region_maps(data_maps, jk_maps, k)
-                    )
-                    vis_alms_k = transform(
-                        fields, _get_region_maps(vis_maps, jk_maps, k)
-                    )
-                write_alms(data_path, data_alms_k, clobber=True)
-                write_alms(vis_path, vis_alms_k, clobber=True)
+            _compute_single_jk_alm(
+                k,
+                data_maps,
+                vis_maps,
+                jk_map,
+                fields,
+                dir,
+            )
+
         current += 1
         progress.update(current, total)
 
-    data_alms_full = read_alms(os.path.join(dir, "data_alms_0.fits"))
-    vis_alms_full = read_alms(os.path.join(dir, "vis_alms_0.fits"))
-    mls0 = angular_power_spectra(vis_alms_full)
 
-    # Compute Cls
+def _compute_single_jk_alm(
+    k,
+    data_maps,
+    vis_maps,
+    jk_map,
+    fields,
+    dir="./dices",
+):
+    data_path = os.path.join(dir, f"data_alms_{k}.fits")
+    vis_path = os.path.join(dir, f"vis_alms_{k}.fits")
+
+    if os.path.exists(data_path) and os.path.exists(vis_path):
+        return k, False  # nothing done
+
+    if k == 0:
+        data_alms_k = transform(fields, data_maps)
+        vis_alms_k = transform(fields, vis_maps)
+    else:
+        data_alms_k = transform(fields, _get_region_maps(data_maps, jk_map, k))
+        vis_alms_k = transform(fields, _get_region_maps(vis_maps, jk_map, k))
+
+    write_alms(data_path, data_alms_k, clobber=True)
+    write_alms(vis_path, vis_alms_k, clobber=True)
+
+    return k, True  # processed
+
+
+def compute_jk_cls_from_alms(
+    jk_map,
+    fields,
+    mask_correction="Fast",
+    unmixed=False,
+    nd=1,
+    dir="./dices",
+    progress=None,
+):
+    if nd == 0:
+        data_alms_full = read_alms(os.path.join(dir, "data_alms_0.fits"))
+        cls0 = angular_power_spectra(data_alms_full)
+        return {(): cls0}
+
+    if nd < 1 or nd > 2:
+        raise ValueError("number of deletions must be 1 or 2")
+
+    if progress is None:
+        progress = NoProgress()
+
+    cls = {}
+
+    njk = len(np.unique(jk_map)[np.unique(jk_map) != 0])
+    all_regions = list(combinations(range(1, njk + 1), nd))
+
+    total = len(all_regions)
+    current = 0
+    progress.update(current, total)
+
     for regions in all_regions:
-        regions_tag = "_".join(map(str, regions))
-        cls_path = os.path.join(dir, f"cls_{regions_tag}_unmixed_{unmixed}.fits")
         with progress.task(f"Cls {regions}"):
-            if os.path.exists(cls_path):
-                cls[regions] = read(cls_path)
-            else:
-                alms_jk = _subtract_alms(
-                    data_alms_full,
-                    _accumulate_alms(
-                        os.path.join(dir, f"data_alms_{r}.fits") for r in regions
-                    ),
-                )
-                _cls = angular_power_spectra(alms_jk)
-                _cls = correct_bias(_cls, jk_maps, fields, *regions)
-                if mask_correction == "Full":
-                    vis_alms_jk = _subtract_alms(
-                        vis_alms_full,
-                        _accumulate_alms(
-                            os.path.join(dir, f"vis_alms_{r}.fits") for r in regions
-                        ),
-                    )
-                    _cls_mm = angular_power_spectra(vis_alms_jk)
-                    _cls = correct_footprint_naturalspice(
-                        _cls, _cls_mm, mls0, fields, unmixed=unmixed
-                    )
-                elif mask_correction == "Fast":
-                    _cls = correct_footprint_fsky(
-                        _cls, jk_maps, fields, *regions, unmixed=unmixed
-                    )
-                else:
-                    raise ValueError("mask_correction must be 'Fast' or 'Full'")
-                write(cls_path, _cls, clobber=True)
-                cls[regions] = _cls
+            cls[regions] = _compute_single_jk_cls(
+                regions,
+                jk_map,
+                fields,
+                mask_correction,
+                unmixed,
+                dir,
+            )
+
         current += 1
         progress.update(current, total)
+
     return cls
 
 
-def _get_region_maps(maps, jkmaps, jk):
+def _compute_single_jk_cls(
+    regions,
+    jk_map,
+    fields,
+    mask_correction="Fast",
+    unmixed=False,
+    dir="./dices",
+):
+    """Compute Cls for a single jackknife region combination."""
+
+    regions_tag = "_".join(map(str, regions))
+    cls_path = os.path.join(dir, f"cls_{regions_tag}_unmixed_{unmixed}.fits")
+
+    if os.path.exists(cls_path):
+        return read(cls_path)
+
+    data_alms_full = read_alms(os.path.join(dir, "data_alms_0.fits"))
+    vis_alms_full = read_alms(os.path.join(dir, "vis_alms_0.fits"))
+    mls0 = angular_power_spectra(vis_alms_full)
+
+    alms_jk = _subtract_alms(
+        data_alms_full,
+        _accumulate_alms(os.path.join(dir, f"data_alms_{r}.fits") for r in regions),
+    )
+
+    _cls = angular_power_spectra(alms_jk)
+    _cls = correct_bias(_cls, jk_map, fields, *regions)
+
+    if mask_correction == "Full":
+        vis_alms_jk = _subtract_alms(
+            vis_alms_full,
+            _accumulate_alms(os.path.join(dir, f"vis_alms_{r}.fits") for r in regions),
+        )
+        _cls_mm = angular_power_spectra(vis_alms_jk)
+        _cls = correct_footprint_naturalspice(
+            _cls, _cls_mm, mls0, fields, unmixed=unmixed
+        )
+
+    elif mask_correction == "Fast":
+        _cls = correct_footprint_fsky(_cls, jk_map, *regions, unmixed=unmixed)
+
+    else:
+        raise ValueError("mask_correction must be 'Fast' or 'Full'")
+
+    write(cls_path, _cls, clobber=True)
+
+    return _cls
+
+
+def _get_region_maps(maps, jk_map, jk):
     """
     Returns maps with only the pixels belonging to jackknife region *jk* active.
     All other pixels are set to zero.
     """
     _maps = deepcopy(maps)
-    for key_data, key_mask in zip(maps.keys(), jkmaps.keys()):
+    for key_data in maps.keys():
         _map = _maps[key_data]
-        _jkmap = jkmaps[key_mask]
-        if _jkmap is None:
-            continue
-        _mask = (_jkmap == float(jk)).astype(int)
+        _mask = (jk_map == float(jk)).astype(int)
         _map *= _mask
     return _maps
 
@@ -218,30 +319,25 @@ def bias(cls):
     return bias
 
 
-def jackknife_fsky(jkmaps, jk=0, jk2=0, ratio=True):
+def jackknife_fsky(jk_map, jk=0, jk2=0, ratio=True):
     """
     Returns the fraction of the sky after deleting two regions.
     inputs:
-        jkmaps (dict): Dictionary of Jackknife maps
+        jk_map (array): Jackknife mask map
         jk (int): Jackknife region to remove
         jk2 (int): Jackknife region to remove
         ratio (bool): Return the ratio of fskyjk to fsky
     returns:
         fskyjk2 (np.array): Fraction of the sky after deleting two regions.
     """
-    fskysjk = {}
-    for key in jkmaps.keys():
-        jkmap = jkmaps[key]
-        mask = np.copy(jkmap)
-        mask = (mask > 0).astype(int)
-        fsky = sum(mask) / len(mask)
-        cond = np.where((mask == 1.0) & (jkmap != jk) & (jkmap != jk2))[0]
-        fskyjk = len(cond) / len(mask)
-        if ratio:
-            fskysjk[key] = fskyjk / fsky
-        else:
-            fskysjk[key] = fskyjk
-    return fskysjk
+    mask = np.copy(jk_map)
+    mask = (mask > 0).astype(int)
+    fsky = sum(mask) / len(mask)
+    cond = np.where((mask == 1.0) & (jk_map != jk) & (jk_map != jk2))[0]
+    fskyjk = len(cond) / len(mask)
+    if ratio:
+        fskyjk = fskyjk / fsky
+    return fskyjk
 
 
 def jackknife_bias(bias, fsky, fields):
@@ -258,23 +354,17 @@ def jackknife_bias(bias, fsky, fields):
     for key in list(bias.keys()):
         f1, f2, b1, b2 = key
         b = bias[key]
-        if (f1, b1) == (f2, b2):
-            field = fields[f1]
-            m_f = field.mask
-            fskyjk = fsky[(m_f, b1)]
-        else:
-            fskyjk = 0.0
-        b_jk = b * fskyjk
+        b_jk = b * fsky
         bias_jk[key] = b_jk
-    return bias
+    return bias_jk
 
 
-def correct_bias(cls, jkmaps, fields, jk=0, jk2=0):
+def correct_bias(cls, jk_map, fields, jk=0, jk2=0):
     """
     Corrects the bias of the Cls due to taking out a region.
     inputs:
         cls (dict): Dictionary of Cls
-        jkmaps (dict): Dictionary of Jackknife maps
+        jk_map (array): Jackknife mask map
         fields (dict): Dictionary of fields
         jk (int): Jackknife region to remove
         jk2 (int): Jackknife region to remove
@@ -284,7 +374,7 @@ def correct_bias(cls, jkmaps, fields, jk=0, jk2=0):
     """
     # Bias correction
     b = bias(cls)
-    fskyjk = jackknife_fsky(jkmaps, jk=jk, jk2=jk2)
+    fskyjk = jackknife_fsky(jk_map, jk=jk, jk2=jk2)
     b_jk = jackknife_bias(b, fskyjk, fields)
     # Correct Cls
     cls = add_to_Cls(cls, b)
@@ -297,31 +387,23 @@ def correct_bias(cls, jkmaps, fields, jk=0, jk2=0):
     return cls
 
 
-def correct_footprint_fsky(cls, jkmaps, fields, jk=0, jk2=0, unmixed=False):
+def correct_footprint_fsky(cls, jk_map, jk=0, jk2=0, unmixed=False):
     """
     Corrects the Cls for the footprint reduction due to taking out a region.
     inputs:
         cls (dict): Dictionary of Cls
-        jkmaps (dict): Dictionary of Jackknife maps
-        fields (dict): Dictionary of fields
+        jk_map (array): Jackknife mask map
         jk (int): Jackknife region to remove
         jk2 (int): Jackknife region to remove
         unmixed (bool): unmix the Cls
     returns:
         cls_cf (dict): Corrected Cls
     """
     ratio = not unmixed
-    fskyjk = jackknife_fsky(jkmaps, jk=jk, jk2=jk2, ratio=ratio)
+    fskyjk = jackknife_fsky(jk_map, jk=jk, jk2=jk2, ratio=ratio)
     _cls = {}
     for key in cls.keys():
-        a, b, i, j = key
-        f_a = fields[a]
-        f_b = fields[b]
-        m_a = f_a.mask
-        m_b = f_b.mask
-        fsky_a = fskyjk[(m_a, i)]
-        fsky_b = fskyjk[(m_b, j)]
-        _cl = cls[key].array / np.sqrt(fsky_a * fsky_b)
+        _cl = cls[key].array / fskyjk
         _cls[key] = replace(cls[key], array=_cl)
     return _cls