From 63643cff35835cdb19d8315d76ab35d83d01f1a2 Mon Sep 17 00:00:00 2001
From: Paddy Roddy <patrickjamesroddy@gmail.com>
Date: Tue, 12 Sep 2023 18:19:27 +0100
Subject: [PATCH] Add various types, including `Unknown` placeholder

---
 heracles/catalog/array.py   |  18 ++++--
 heracles/catalog/base.py    |   6 +-
 heracles/catalog/filters.py |  22 +++++---
 heracles/catalog/fits.py    |  22 ++++----
 heracles/core.py            |  15 ++++-
 heracles/covariance.py      |  35 +++++++-----
 heracles/io.py              | 110 ++++++++++++++++++++++++++++--------
 heracles/maps.py            |  18 +++---
 heracles/plot.py            |  26 +++++----
 heracles/twopoint.py        |  65 +++++++++++++++------
 heracles/util.py            |  15 +++--
 11 files changed, 241 insertions(+), 111 deletions(-)

diff --git a/heracles/catalog/array.py b/heracles/catalog/array.py
index 610ad22..497dcdc 100644
--- a/heracles/catalog/array.py
+++ b/heracles/catalog/array.py
@@ -17,6 +17,8 @@
 # You should have received a copy of the GNU Lesser General Public
 # License along with Heracles. If not, see <https://www.gnu.org/licenses/>.
 """Module for array catalogues."""
+from collections.abc import Iterator
+from typing import Optional, TypeVar
 
 from .base import CatalogBase, CatalogPage
 
@@ -28,7 +30,7 @@ class ArrayCatalog(CatalogBase):
         CatalogBase: _description_
     """
 
-    def __init__(self, arr):
+    def __init__(self, arr: TypeVar("Unknown")) -> None:
         """Create a new array catalogue reader.
 
         Args:
@@ -37,7 +39,7 @@ def __init__(self, arr):
         super().__init__()
         self._arr = arr
 
-    def __copy__(self):
+    def __copy__(self) -> "ArrayCatalog":
         """Return a copy of this catalogue.
 
         Returns
@@ -47,7 +49,7 @@ def __copy__(self):
         other._arr = self._arr
         return other
 
-    def _names(self):
+    def _names(self) -> TypeVar("Unknown"):
         """_summary_.
 
         Returns
@@ -55,7 +57,7 @@ def _names(self):
         """
         return self._arr.dtype.names
 
-    def _size(self, selection):
+    def _size(self, selection: Optional[TypeVar("Unknown")]) -> int:
         """_summary_.
 
         Args:
@@ -68,7 +70,11 @@ def _size(self, selection):
             return len(self._arr)
         return len(self._arr[selection])
 
-    def _join(self, first, *other):
+    def _join(
+        self,
+        first: TypeVar("Unknown"),
+        *other: TypeVar("Unknown"),
+    ) -> TypeVar("Unknown"):
         """Join boolean masks.
 
         Args:
@@ -82,7 +88,7 @@ def _join(self, first, *other):
             mask = mask & a
         return mask
 
-    def _pages(self, selection):
+    def _pages(self, selection: Optional[TypeVar("Unknown")]) -> Iterator[CatalogPage]:
         """Iterate the rows of the array in pages.
 
         Args:
diff --git a/heracles/catalog/base.py b/heracles/catalog/base.py
index 668a169..d7644f9 100644
--- a/heracles/catalog/base.py
+++ b/heracles/catalog/base.py
@@ -32,7 +32,7 @@ class CatalogPage:
     Internally holds all column data as a numpy array.
     """
 
-    def _update(self):
+    def _update(self) -> None:
         """Update internal data after dictionary changes.
 
         Raises
@@ -72,7 +72,7 @@ def __getitem__(self, col):
             return tuple(self._data[c] for c in col)
         return self._data[col]
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Number of columns in the page.
 
         Returns
@@ -97,7 +97,7 @@ def __iter__(self):
         yield from self._data
 
     @property
-    def names(self):
+    def names(self) -> list:
         """Column names in the page.
 
         Returns
diff --git a/heracles/catalog/filters.py b/heracles/catalog/filters.py
index 8815292..51d0f62 100644
--- a/heracles/catalog/filters.py
+++ b/heracles/catalog/filters.py
@@ -19,6 +19,7 @@
 """Module for catalogue filters."""
 
 import warnings
+from typing import Optional, TypeVar
 
 import healpy as hp
 import numpy as np
@@ -27,7 +28,12 @@
 class InvalidValueFilter:
     """Filter invalid values from a catalogue."""
 
-    def __init__(self, *columns, weight=None, warn=True):
+    def __init__(
+        self,
+        *columns: TypeVar("Unknown"),
+        weight: Optional[TypeVar("Unknown")] = None,
+        warn: bool = True,
+    ) -> None:
         """Filter invalid values in the given columns.
 
         If ``warn`` is true, invalid values will emit a warning.
@@ -40,7 +46,7 @@ def __init__(self, *columns, weight=None, warn=True):
         self.weight = weight
         self.warn = warn
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """_summary_.
 
         Returns
@@ -52,7 +58,7 @@ def __repr__(self):
         args = ", ".join(args)
         return f"{name}({args})"
 
-    def __call__(self, page):
+    def __call__(self, page: TypeVar("Unknown")) -> None:
         """Filter a catalog page.
 
         Args:
@@ -73,7 +79,7 @@ def __call__(self, page):
 class FootprintFilter:
     """Filter a catalogue using a footprint map."""
 
-    def __init__(self, footprint, lon, lat):
+    def __init__(self, footprint: TypeVar("Unknown"), lon: float, lat: float) -> None:
         """Filter using the given footprint map and position columns.
 
         Args:
@@ -86,7 +92,7 @@ def __init__(self, footprint, lon, lat):
         self._lonlat = (lon, lat)
 
     @property
-    def footprint(self):
+    def footprint(self) -> TypeVar("Unknown"):
         """Footprint for filter.
 
         Returns
@@ -95,7 +101,7 @@ def footprint(self):
         return self._footprint
 
     @property
-    def lonlat(self):
+    def lonlat(self) -> tuple[float, float]:
         """Longitude and latitude columns.
 
         Returns
@@ -103,7 +109,7 @@ def lonlat(self):
         """
         return self._lonlat
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """_summary_.
 
         Returns
@@ -113,7 +119,7 @@ def __repr__(self):
         lon, lat = self.lonlat
         return f"{name}(..., {lon!r}, {lat!r})"
 
-    def __call__(self, page):
+    def __call__(self, page: TypeVar("Unknown")) -> None:
         """Filter catalog page.
 
         Args:
diff --git a/heracles/catalog/fits.py b/heracles/catalog/fits.py
index 773e87c..537171b 100644
--- a/heracles/catalog/fits.py
+++ b/heracles/catalog/fits.py
@@ -18,6 +18,8 @@
 # License along with Heracles. If not, see <https://www.gnu.org/licenses/>.
 """Module for catalogue processing."""
 
+from collections.abc import Iterator
+from typing import Any, Optional, TypeVar
 from weakref import finalize, ref
 
 import fitsio
@@ -25,7 +27,7 @@
 from .base import CatalogBase, CatalogPage
 
 
-def _is_table_hdu(hdu):
+def _is_table_hdu(hdu: fitsio.hdu.TableHDU) -> bool:
     """Return true if HDU is a table with data.
 
     Args:
@@ -37,7 +39,7 @@ def _is_table_hdu(hdu):
     return isinstance(hdu, fitsio.hdu.TableHDU) and hdu.has_data()
 
 
-def rowfilter(array, expr):
+def rowfilter(array: TypeVar("Unknown"), expr: str) -> Any:
     """Filter the rows of a structured array.
 
     Args:
@@ -57,7 +59,7 @@ class FitsCatalog(CatalogBase):
         CatalogBase: _description_
     """
 
-    def __init__(self, filename, *, columns=None, ext=None):
+    def __init__(self, filename, *, columns=None, ext=None) -> None:
         """Create a new FITS catalogue reader.
 
         Neither opens the FITS file nor reads the catalogue immediately.
@@ -72,7 +74,7 @@ def __init__(self, filename, *, columns=None, ext=None):
         self._columns = columns
         self._ext = ext
 
-    def __copy__(self):
+    def __copy__(self) -> "FitsCatalog":
         """Return a copy of this catalog.
 
         Returns
@@ -84,7 +86,7 @@ def __copy__(self):
         other._ext = self._ext
         return other
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """String representation of FitsCatalog.
 
         Returns
@@ -95,7 +97,7 @@ def __repr__(self):
             s = s + f"[{self._ext!r}]"
         return s
 
-    def hdu(self):
+    def hdu(self) -> TypeVar("Unknown"):
         """HDU for catalogue data.
 
         Raises
@@ -142,7 +144,7 @@ def hdu(self):
 
         return hdu
 
-    def _names(self):
+    def _names(self) -> Optional[TypeVar("Unknown")]:
         """Column names in FITS catalogue.
 
         Returns
@@ -153,7 +155,7 @@ def _names(self):
             self._columns = self.hdu().get_colnames()
         return self._columns
 
-    def _size(self, selection):
+    def _size(self, selection: TypeVar("Unknown")) -> int:
         """Size of FITS catalogue; selection is ignored.
 
         Args:
@@ -164,7 +166,7 @@ def _size(self, selection):
         """
         return self.hdu().get_nrows()
 
-    def _join(self, *where):
+    def _join(self, *where: TypeVar("Unknown")) -> Optional[str]:
         """Join rowfilter expressions.
 
         Returns
@@ -174,7 +176,7 @@ def _join(self, *where):
             return None
         return "(" + ") & (".join(map(str, filter(None, where))) + ")"
 
-    def _pages(self, selection):
+    def _pages(self, selection: TypeVar("Unknown")) -> Iterator[CatalogPage]:
         """Iterate pages of rows in FITS file, optionally using the query.
 
         Args:
diff --git a/heracles/core.py b/heracles/core.py
index dd700d2..4e3282c 100644
--- a/heracles/core.py
+++ b/heracles/core.py
@@ -20,9 +20,14 @@
 
 from collections import UserDict
 from collections.abc import Mapping, Sequence
+from typing import Optional, TypeVar, Union
 
 
-def toc_match(key, include=None, exclude=None):
+def toc_match(
+    key: TypeVar("Unknown"),
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> bool:
     """Return whether a tocdict entry matches include/exclude criteria.
 
     Args:
@@ -46,7 +51,11 @@ def toc_match(key, include=None, exclude=None):
     return True
 
 
-def toc_filter(obj, include=None, exclude=None):
+def toc_filter(
+    obj: Union[Sequence, Mapping],
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> Union[dict, list]:
     """Return a filtered toc dict ``d``.
 
     Args:
@@ -77,7 +86,7 @@ class TocDict(UserDict):
         UserDict: _description_
     """
 
-    def __getitem__(self, pattern):
+    def __getitem__(self, pattern: TypeVar("Unknown")) -> TypeVar("Unknown"):
         """Look up one or many keys in dict.
 
         Args:
diff --git a/heracles/covariance.py b/heracles/covariance.py
index eef320a..ab7d01b 100644
--- a/heracles/covariance.py
+++ b/heracles/covariance.py
@@ -22,6 +22,7 @@
 import time
 from datetime import timedelta
 from itertools import combinations_with_replacement
+from typing import Optional, TypeVar
 
 import healpy as hp
 import numpy as np
@@ -38,7 +39,11 @@ class SampleCovariance(np.ndarray):
         np: _description_
     """
 
-    def __new__(cls, nrows, ncols=None):
+    def __new__(
+        cls,
+        nrows: TypeVar("Unknown"),
+        ncols: Optional[TypeVar("Unknown")] = None,
+    ) -> TypeVar("Unknown"):
         """_summary_.
 
         Args:
@@ -56,7 +61,7 @@ def __new__(cls, nrows, ncols=None):
         cov.sample_col_mean = np.zeros(ncols)
         return cov
 
-    def __array_finalize__(self, cov):
+    def __array_finalize__(self, cov: TypeVar("Unknown")) -> None:
         """_summary_.
 
         Args:
@@ -72,7 +77,11 @@ def __array_finalize__(self, cov):
         self.sample_col_mean[:] = getattr(cov, "sample_col_mean", 0.0)
 
 
-def add_sample(cov, x, y=None):
+def add_sample(
+    cov: TypeVar("Unknown"),
+    x: TypeVar("Unknown"),
+    y: Optional[TypeVar("Unknown")] = None,
+) -> None:
     """Add a sample to a sample covariance matrix.
 
     Args:
@@ -101,7 +110,7 @@ def add_sample(cov, x, y=None):
         cov += (np.outer(delta, y - cov.sample_col_mean) - cov) / (cov.sample_count - 1)
 
 
-def update_covariance(cov, sample):
+def update_covariance(cov: TypeVar("Unknown"), sample: TypeVar("Unknown")) -> None:
     """Update a set of sample covariances given a sample.
 
     Args:
@@ -133,14 +142,14 @@ def update_covariance(cov, sample):
 
 
 def jackknife_regions_kmeans(
-    fpmap,
-    n,
+    fpmap: TypeVar("Unknown"),
+    n: TypeVar("Unknown"),
     *,
-    maxrepeat=5,
-    maxiter=1000,
-    tol=1e-5,
-    return_centers=False,
-):
+    maxrepeat: int = 5,
+    maxiter: int = 1_000,
+    tol: float = 1e-5,
+    return_centers: bool = False,
+) -> TypeVar("Unknown"):
     """Partition a footprint map into n regions using k-means.
 
     Args:
@@ -196,9 +205,9 @@ def jackknife_regions_kmeans(
     areas = 60**4 // 100 / np.pi / npix * np.bincount(km.labels)
     area_mean, area_std, area_unit = np.mean(areas), np.std(areas), "deg2"
     if area_mean < 1.0:
-        area_mean, area_std, area_unit = area_mean / 3600, area_std / 3600, "arcmin2"
+        area_mean, area_std, area_unit = area_mean / 3_600, area_std / 3_600, "arcmin2"
     if area_mean < 1.0:
-        area_mean, area_std, area_unit = area_mean / 3600, area_std / 3600, "arcsec2"
+        area_mean, area_std, area_unit = area_mean / 3_600, area_std / 3_600, "arcsec2"
     logger.info("region area is %.3f ± %.3f %s", area_mean, area_std, area_unit)
 
     jkmap = np.zeros(npix, dtype=int)
diff --git a/heracles/io.py b/heracles/io.py
index bab3995..b9a4a9a 100644
--- a/heracles/io.py
+++ b/heracles/io.py
@@ -20,10 +20,12 @@
 
 import logging
 import os
+from typing import Optional, TypeVar
 
 import fitsio
 import healpy as hp
 import numpy as np
+import numpy.typing as npt
 
 from .core import TocDict, toc_match
 
@@ -44,7 +46,7 @@
 }
 
 
-def _write_metadata(hdu, metadata):
+def _write_metadata(hdu: TypeVar("Unknown"), metadata: TypeVar("Unknown")) -> None:
     """Write array metadata to FITS HDU.
 
     Args:
@@ -56,7 +58,7 @@ def _write_metadata(hdu, metadata):
         hdu.write_key("META " + key.upper(), value, _METADATA_COMMENTS.get(key))
 
 
-def _read_metadata(hdu):
+def _read_metadata(hdu: TypeVar("Unknown")) -> dict:
     """Read array metadata from FITS HDU.
 
     Args:
@@ -73,7 +75,12 @@ def _read_metadata(hdu):
     return md
 
 
-def read_mask(mask_name, nside=None, field=0, extra_mask_name=None):
+def read_mask(
+    mask_name: TypeVar("Unknown"),
+    nside: Optional[TypeVar("Unknown")] = None,
+    field: int = 0,
+    extra_mask_name: Optional[TypeVar("Unknown")] = None,
+) -> npt.NDArray:
     """Read visibility map from a HEALPix map file.
 
     Args:
@@ -113,14 +120,14 @@ def read_mask(mask_name, nside=None, field=0, extra_mask_name=None):
 
 
 def write_maps(
-    filename,
-    maps,
+    filename: str,
+    maps: TypeVar("Unknown"),
     *,
-    clobber=False,
-    workdir=".",
-    include=None,
-    exclude=None,
-):
+    clobber: bool = False,
+    workdir: str = ".",
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> None:
     """Write a set of maps to FITS file.
 
     If the output file exists, the new estimates will be appended, unless the
@@ -217,7 +224,13 @@ def write_maps(
     logger.info("done with %d maps", len(maps))
 
 
-def read_maps(filename, workdir=".", *, include=None, exclude=None):
+def read_maps(
+    filename: str,
+    workdir: str = ".",
+    *,
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> TocDict:
     """Read a set of maps from a FITS file.
 
     Args:
@@ -273,14 +286,14 @@ def read_maps(filename, workdir=".", *, include=None, exclude=None):
 
 
 def write_alms(
-    filename,
-    alms,
+    filename: str,
+    alms: TypeVar("Unknown"),
     *,
-    clobber=False,
-    workdir=".",
-    include=None,
-    exclude=None,
-):
+    clobber: bool = False,
+    workdir: str = ".",
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> None:
     """Write a set of alms to FITS file.
 
     If the output file exists, the new estimates will be appended, unless the
@@ -347,7 +360,13 @@ def write_alms(
     logger.info("done with %d alms", len(alms))
 
 
-def read_alms(filename, workdir=".", *, include=None, exclude=None):
+def read_alms(
+    filename: str,
+    workdir: str = ".",
+    *,
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+):
     """Read a set of alms from a FITS file.
 
     Args:
@@ -401,7 +420,15 @@ def read_alms(filename, workdir=".", *, include=None, exclude=None):
     return alms
 
 
-def write_cls(filename, cls, *, clobber=False, workdir=".", include=None, exclude=None):
+def write_cls(
+    filename: str,
+    cls,
+    *,
+    clobber: bool = False,
+    workdir: str = ".",
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> None:
     """Write a set of cls to FITS file.
 
     If the output file exists, the new estimates will be appended, unless the
@@ -487,7 +514,13 @@ def write_cls(filename, cls, *, clobber=False, workdir=".", include=None, exclud
     logger.info("done with %d cls", len(cls))
 
 
-def read_cls(filename, workdir=".", *, include=None, exclude=None):
+def read_cls(
+    filename: str,
+    workdir: str = ".",
+    *,
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> TocDict:
     """Read a set of cls from a FITS file.
 
     Args:
@@ -537,7 +570,15 @@ def read_cls(filename, workdir=".", *, include=None, exclude=None):
     return cls
 
 
-def write_mms(filename, mms, *, clobber=False, workdir=".", include=None, exclude=None):
+def write_mms(
+    filename: str,
+    mms: TypeVar("Unknown"),
+    *,
+    clobber: bool = False,
+    workdir: str = ".",
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> None:
     """Write a set of mixing matrices to FITS file.
 
     If the output file exists, the new mixing matrices will be appended, unless
@@ -613,7 +654,13 @@ def write_mms(filename, mms, *, clobber=False, workdir=".", include=None, exclud
     logger.info("done with %d mm(s)", len(mms))
 
 
-def read_mms(filename, workdir=".", *, include=None, exclude=None):
+def read_mms(
+    filename: str,
+    workdir: str = ".",
+    *,
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> TypeVar("Unknown"):
     """Read a set of mixing matrices from a FITS file.
 
     Args:
@@ -663,7 +710,14 @@ def read_mms(filename, workdir=".", *, include=None, exclude=None):
     return mms
 
 
-def write_cov(filename, cov, clobber=False, workdir=".", include=None, exclude=None):
+def write_cov(
+    filename: str,
+    cov: TypeVar("Unknown"),
+    clobber: bool = False,
+    workdir: str = ".",
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> None:
     """Write a set of covariance matrices to FITS file.
 
     If the output file exists, the new estimates will be appended, unless the
@@ -749,7 +803,13 @@ def write_cov(filename, cov, clobber=False, workdir=".", include=None, exclude=N
     logger.info("done with %d covariance(s)", len(cov))
 
 
-def read_cov(filename, workdir=".", *, include=None, exclude=None):
+def read_cov(
+    filename: str,
+    workdir: str = ".",
+    *,
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> TocDict:
     """Read a set of covariances matrices from a FITS file.
 
     Args:
diff --git a/heracles/maps.py b/heracles/maps.py
index b7fd9c2..b007e7f 100644
--- a/heracles/maps.py
+++ b/heracles/maps.py
@@ -61,7 +61,7 @@ def wrapper(*inputs):
 
 @_nativebyteorder
 @njit(nogil=True, fastmath=True)
-def _map_pos(pos, ipix):
+def _map_pos(pos, ipix) -> None:
     """_summary_.
 
     Args:
@@ -74,7 +74,7 @@ def _map_pos(pos, ipix):
 
 @_nativebyteorder
 @njit(nogil=True, fastmath=True)
-def _map_real(wht, val, ipix, w, v):
+def _map_real(wht, val, ipix, w, v) -> None:
     """_summary_.
 
     Args:
@@ -91,7 +91,7 @@ def _map_real(wht, val, ipix, w, v):
 
 @_nativebyteorder
 @njit(nogil=True, fastmath=True)
-def _map_complex(wht, val, ipix, w, re, im):
+def _map_complex(wht, val, ipix, w, re, im) -> None:
     """_summary_.
 
     Args:
@@ -110,7 +110,7 @@ def _map_complex(wht, val, ipix, w, re, im):
 
 @_nativebyteorder
 @njit(nogil=True, fastmath=True)
-def _map_weight(wht, ipix, w):
+def _map_weight(wht, ipix, w) -> None:
     """_summary_.
 
     Args:
@@ -122,7 +122,7 @@ def _map_weight(wht, ipix, w):
         wht[i] += w_i
 
 
-def update_metadata(array, **metadata):
+def update_metadata(array, **metadata) -> None:
     """Update metadata of an array dtype.
 
     Args:
@@ -211,7 +211,7 @@ class HealpixMap(Map):
         Map: _description_
     """
 
-    def __init__(self, nside: int, **kwargs) -> None:
+    def __init__(self, nside: int, **kwargs: t.Any) -> None:
         """Initialize map with the given nside parameter.
 
         Args:
@@ -852,7 +852,7 @@ def __init__(
         lat: str,
         weight: str,
         *,
-        normalize=True,
+        normalize: bool = True,
     ) -> None:
         """Create a new weight map.
 
@@ -933,7 +933,7 @@ def mapper(page: "CatalogPage") -> None:
 EllipticityMap = Spin2Map
 
 
-def _close_and_return(generator):
+def _close_and_return(generator: t.Generator):
     """_summary_.
 
     Args:
@@ -1088,7 +1088,7 @@ def transform_maps(
     *,
     out: t.MutableMapping[t.Any, t.Any] = None,
     progress: bool = False,
-    **kwargs,
+    **kwargs: t.Any,
 ) -> dict[tuple[t.Any, t.Any], np.ndarray]:
     """Transform a set of maps to alms.
 
diff --git a/heracles/plot.py b/heracles/plot.py
index 1e6f5ed..cc4149f 100644
--- a/heracles/plot.py
+++ b/heracles/plot.py
@@ -21,7 +21,9 @@
 from collections import defaultdict
 from collections.abc import Mapping
 from itertools import chain, count, cycle
+from typing import Optional, TypeVar
 
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 from cycler import cycler
@@ -29,7 +31,7 @@
 DEFAULT_CYCLER = cycler(linestyle=["-", "--", ":", "-."])
 
 
-def _dont_draw_zero_tick(tick):
+def _dont_draw_zero_tick(tick: TypeVar("Unknown")) -> None:
     """Custom draw function for ticks that does not draw zero.
 
     Args:
@@ -49,7 +51,7 @@ def wrap(*args, **kwargs):
     return wrap
 
 
-def _pad_ylim(ymin, ymax):
+def _pad_ylim(ymin: float, ymax: float) -> tuple[float, float]:
     """Pad the y axis range depending on signs.
 
     Args:
@@ -63,17 +65,17 @@ def _pad_ylim(ymin, ymax):
 
 
 def postage_stamps(
-    plot=None,
-    transpose=None,
+    plot: Optional[TypeVar("Unknown")] = None,
+    transpose: Optional[TypeVar("Unknown")] = None,
     *,
-    scale=None,
-    trxshift=0,
-    tryshift=0,
-    stampsize=1.0,
-    hatch_empty=False,
-    linscale=0.01,
-    cycler=None,
-):
+    scale: Optional[TypeVar("Unknown")] = None,
+    trxshift: int = 0,
+    tryshift: int = 0,
+    stampsize: float = 1.0,
+    hatch_empty: bool = False,
+    linscale: float = 0.01,
+    cycler: Optional[TypeVar("Unknown")] = None,
+) -> mpl.figure.Figure:
     """Create a postage stamp plot for cls.
 
     Args:
diff --git a/heracles/twopoint.py b/heracles/twopoint.py
index b5de4fe..942efd6 100644
--- a/heracles/twopoint.py
+++ b/heracles/twopoint.py
@@ -22,9 +22,11 @@
 import time
 from datetime import timedelta
 from itertools import combinations_with_replacement, product
+from typing import Any, Optional, TypeVar, Union
 
 import healpy as hp
 import numpy as np
+import numpy.typing as npt
 from convolvecl import mixmat, mixmat_eb
 
 from .core import TocDict, toc_match
@@ -41,7 +43,14 @@
 logger = logging.getLogger(__name__)
 
 
-def angular_power_spectra(alms, alms2=None, *, lmax=None, include=None, exclude=None):
+def angular_power_spectra(
+    alms: TypeVar("Unknown"),
+    alms2: Optional[TypeVar("Unknown")] = None,
+    *,
+    lmax: Optional[TypeVar("Unknown")] = None,
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+) -> TocDict:
     """Compute angular power spectra from a set of alms.
 
     Args:
@@ -126,7 +135,12 @@ def angular_power_spectra(alms, alms2=None, *, lmax=None, include=None, exclude=
     return cls
 
 
-def debias_cls(cls, bias=None, *, inplace=False):
+def debias_cls(
+    cls,
+    bias: Optional[TypeVar("Unknown")] = None,
+    *,
+    inplace: bool = False,
+) -> Union[TypeVar("Unknown"), TocDict]:
     """Remove bias from cls.
 
     Args:
@@ -184,7 +198,7 @@ def debias_cls(cls, bias=None, *, inplace=False):
     return out
 
 
-def depixelate_cls(cls, *, inplace=False):
+def depixelate_cls(cls, *, inplace: bool = False) -> Union[TypeVar("Unknown"), TocDict]:
     """Remove discretisation kernel from cls.
 
     Args:
@@ -280,7 +294,13 @@ def depixelate_cls(cls, *, inplace=False):
     return out
 
 
-def mixing_matrices(cls, *, l1max=None, l2max=None, l3max=None):
+def mixing_matrices(
+    cls,
+    *,
+    l1max: Optional[TypeVar("Unknown")] = None,
+    l2max: Optional[TypeVar("Unknown")] = None,
+    l3max: Optional[TypeVar("Unknown")] = None,
+) -> TocDict:
     """Compute mixing matrices from a set of cls.
 
     Args:
@@ -349,7 +369,12 @@ def mixing_matrices(cls, *, l1max=None, l2max=None, l3max=None):
     return mms
 
 
-def pixelate_mms_healpix(mms, nside, *, inplace=False):
+def pixelate_mms_healpix(
+    mms: TypeVar("Unknown"),
+    nside: TypeVar("Unknown"),
+    *,
+    inplace: bool = False,
+) -> Union[TypeVar("Unknown"), TocDict]:
     """Apply HEALPix pixel window function to mms.
 
     Args:
@@ -415,7 +440,13 @@ def pixelate_mms_healpix(mms, nside, *, inplace=False):
     return out
 
 
-def binned_cls(cls, bins, *, weights=None, out=None):
+def binned_cls(
+    cls,
+    bins: TypeVar("Unknown"),
+    *,
+    weights: Optional[TypeVar("Unknown")] = None,
+    out: Optional[TypeVar("Unknown")] = None,
+) -> TocDict:
     """Compute binned angular power spectra.
 
     Args:
@@ -424,7 +455,7 @@ def binned_cls(cls, bins, *, weights=None, out=None):
         out: _description_
     """
 
-    def norm(a, b):
+    def norm(a: npt.NDArray, b: npt.NDArray) -> npt.NDArray:
         """Divide a by b if a is nonzero.
 
         Args:
@@ -493,17 +524,17 @@ def norm(a, b):
 
 
 def random_bias(
-    maps,
-    catalogs,
+    maps: TypeVar("Unknown"),
+    catalogs: TypeVar("Unknown"),
     *,
-    repeat=1,
-    full=False,
-    parallel=False,
-    include=None,
-    exclude=None,
-    progress=False,
-    **kwargs,
-):
+    repeat: int = 1,
+    full: bool = False,
+    parallel: bool = False,
+    include: Optional[TypeVar("Unknown")] = None,
+    exclude: Optional[TypeVar("Unknown")] = None,
+    progress: bool = False,
+    **kwargs: Any,
+) -> TocDict:
     """Bias estimate from randomised maps.
 
     The ``include`` and ``exclude`` selection is applied to the maps.
diff --git a/heracles/util.py b/heracles/util.py
index 7304003..5f258ef 100644
--- a/heracles/util.py
+++ b/heracles/util.py
@@ -17,17 +17,18 @@
 # You should have received a copy of the GNU Lesser General Public
 # License along with Heracles. If not, see <https://www.gnu.org/licenses/>.
 """Module for utilities."""
-
+import io
 import os
 import sys
 import time
 from datetime import timedelta
+from typing import Optional, TypeVar
 
 
 class Progress:
     """Simple progress bar for operations."""
 
-    def __init__(self, out=sys.stdout):
+    def __init__(self, out: io.IOBase = sys.stdout) -> None:
         """Create a new progress bar.
 
         Args:
@@ -39,7 +40,11 @@ def __init__(self, out=sys.stdout):
         self.total = 0
         self.title = None
 
-    def start(self, total, title=None):
+    def start(
+        self,
+        total: TypeVar("Unknown"),
+        title: Optional[TypeVar("Unknown")] = None,
+    ) -> None:
         """Start new progress.
 
         Args:
@@ -52,7 +57,7 @@ def start(self, total, title=None):
         self.title = title
         self.update(0)
 
-    def update(self, step=1):
+    def update(self, step: int = 1) -> None:
         """Update progress.
 
         Args:
@@ -75,7 +80,7 @@ def update(self, step=1):
         self.out.write(s)
         self.out.flush()
 
-    def stop(self, complete=True):
+    def stop(self, complete: bool = True) -> None:
         """Stop progress and end line.
 
         Args: