Add type annotations to dices/jackknife.py

heracles/dices/jackknife.py

@@ -16,6 +16,11 @@
 #
 # You should have received a copy of the GNU Lesser General Public
 # License along with DICES. If not, see <https://www.gnu.org/licenses/>.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
 import os
 import numpy as np
 import itertools
@@ -29,25 +34,32 @@
 from ..unmixing import _naturalspice
 from ..transforms import cl2corr, corr2cl
 from ..io import write_alms, read_alms, write, read
-from ..progress import NoProgress
+from ..progress import NoProgress, Progress
 
 try:
     from copy import replace
 except ImportError:
     # Python < 3.13
     from dataclasses import replace
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from numpy.typing import NDArray
+
+    from heracles.fields import Field
+
 
 def jackknife_cls(
     data_maps,
     vis_maps,
-    jk_map,
-    fields,
-    mask_correction="Fast",
-    unmixed=False,
-    nd=1,
-    dir="./dices",
-    progress=None,
+    jk_map: NDArray,
+    fields: Mapping[Any, Field],
+    mask_correction: str = "Fast",
+    unmixed: bool = False,
+    nd: int = 1,
+    dir: str = "./dices",
+    progress: Progress | None = None,
 ):
     """
     Compute the Cls of removing 1 Jackknife.
@@ -143,7 +155,7 @@ def jackknife_cls(
     return cls
 
 
-def _get_region_maps(maps, jk_map, jk):
+def _get_region_maps(maps, jk_map: NDArray, jk: int):
     """
     Returns maps with only the pixels belonging to jackknife region *jk* active.
     All other pixels are set to zero.
@@ -214,7 +226,12 @@ def bias(cls):
     return bias
 
 
-def jackknife_fsky(jk_map, jk=0, jk2=0, ratio=True):
+def jackknife_fsky(
+    jk_map: NDArray,
+    jk: int = 0,
+    jk2: int = 0,
+    ratio: bool = True,
+) -> float:
     """
     Returns the fraction of the sky after deleting two regions.
     inputs:
@@ -235,12 +252,12 @@ def jackknife_fsky(jk_map, jk=0, jk2=0, ratio=True):
     return fskyjk
 
 
-def jackknife_bias(bias, fsky, fields):
+def jackknife_bias(bias, fsky: float, fields: Mapping[Any, Field]):
     """
     Returns the bias for deleting a Jackknife region.
     inputs:
         bias (dict): Dictionary of biases
-        fsky (dict): Dictionary of relative fskys
+        fsky (float): Relative fsky
         fields (dict): Dictionary of fields
     returns:
         bias_jk (dict): Dictionary of biases
@@ -254,7 +271,13 @@ def jackknife_bias(bias, fsky, fields):
     return bias_jk
 
 
-def correct_bias(cls, jk_map, fields, jk=0, jk2=0):
+def correct_bias(
+    cls,
+    jk_map: NDArray,
+    fields: Mapping[Any, Field],
+    jk: int = 0,
+    jk2: int = 0,
+):
     """
     Corrects the bias of the Cls due to taking out a region.
     inputs:
@@ -282,7 +305,13 @@ def correct_bias(cls, jk_map, fields, jk=0, jk2=0):
     return cls
 
 
-def correct_footprint_fsky(cls, jk_map, jk=0, jk2=0, unmixed=False):
+def correct_footprint_fsky(
+    cls,
+    jk_map: NDArray,
+    jk: int = 0,
+    jk2: int = 0,
+    unmixed: bool = False,
+):
     """
     Corrects the Cls for the footprint reduction due to taking out a region.
     inputs:
@@ -303,7 +332,7 @@ def correct_footprint_fsky(cls, jk_map, jk=0, jk2=0, unmixed=False):
     return _cls
 
 
-def _mask_correlation_ratio(mljk, mls0, unmixed=False):
+def _mask_correlation_ratio(mljk, mls0, unmixed: bool = False):
     alphas = {}
     wmls0 = cl2corr(mls0)
     wmljk = cl2corr(mljk)
@@ -317,7 +346,9 @@ def _mask_correlation_ratio(mljk, mls0, unmixed=False):
     return alphas
 
 
-def correct_footprint_naturalspice(cls, cls_mm, mls0, fields, unmixed=False):
+def correct_footprint_naturalspice(
+    cls, cls_mm, mls0, fields: Mapping[Any, Field], unmixed: bool = False
+):
     """
     Corrects the Cls for footprint reduction using the full NaMaster/naturalspice approach.
     inputs:
@@ -341,15 +372,15 @@ def correct_footprint_naturalspice(cls, cls_mm, mls0, fields, unmixed=False):
     return binned(cls, np.arange(0, lmax + 1))
 
 
-def jackknife_covariance(dict, nd=1):
+def jackknife_covariance(dict, nd: int = 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):
+def _jackknife_covariance(samples, nd: int = 1):
     """
     Compute the jackknife covariance matrix from a sequence
     of spectra dictionaries *samples*.
@@ -396,7 +427,7 @@ def _jackknife_covariance(samples, nd=1):
     return cov
 
 
-def sample_covariance(samples, samples2=None):
+def sample_covariance(samples: NDArray, samples2: NDArray | None = None) -> NDArray:
     """
     Returns the sample covariance matrix of *samples*, or the sample
     cross-covariance between *samples* and *samples2* if the latter is