gh-358: add static types and fix (some) mypy

glass/core/algorithm.py

@@ -2,17 +2,19 @@
 
 from __future__ import annotations
 
+import typing
+
 import numpy as np
 import numpy.typing as npt
 
 
 def nnls(
-    a: npt.ArrayLike,
-    b: npt.ArrayLike,
+    a: npt.NDArray[typing.Any],
+    b: npt.NDArray[typing.Any],
     *,
     tol: float = 0.0,
     maxiter: int | None = None,
-) -> npt.ArrayLike:
+) -> npt.NDArray[typing.Any]:
     """
     Compute a non-negative least squares solution.
 

glass/core/array.py

@@ -1,18 +1,26 @@
 """Module for array utilities."""
 
+from __future__ import annotations
+
+import typing
 from functools import partial
 
 import numpy as np
+import numpy.typing as npt
 
 
-def broadcast_first(*arrays):  # type: ignore[no-untyped-def]
+def broadcast_first(
+    *arrays: npt.NDArray[typing.Any],
+) -> tuple[npt.NDArray[typing.Any], ...]:
     """Broadcast arrays, treating the first axis as common."""
     arrays = tuple(np.moveaxis(a, 0, -1) if np.ndim(a) else a for a in arrays)
     arrays = np.broadcast_arrays(*arrays)
     return tuple(np.moveaxis(a, -1, 0) if np.ndim(a) else a for a in arrays)
 
 
-def broadcast_leading_axes(*args):  # type: ignore[no-untyped-def]
+def broadcast_leading_axes(
+    *args: tuple[npt.NDArray[typing.Any], int],
+) -> tuple[tuple[int, ...], ...]:
     """
     Broadcast all but the last N axes.
 
@@ -46,10 +54,18 @@ def broadcast_leading_axes(*args):  # type: ignore[no-untyped-def]
         trails.append(s[i:])
     dims = np.broadcast_shapes(*shapes)
     arrs = (np.broadcast_to(a, dims + t) for (a, _), t in zip(args, trails))
-    return (dims, *arrs)
-
-
-def ndinterp(x, xp, fp, axis=-1, left=None, right=None, period=None):  # type: ignore[no-untyped-def] # noqa: PLR0913
+    return (dims, *arrs)  # type: ignore[arg-type]
+
+
+def ndinterp(  # noqa: PLR0913
+    x: npt.NDArray[typing.Any],
+    xp: npt.NDArray[typing.Any],
+    fp: npt.NDArray[typing.Any],
+    axis: int = -1,
+    left: float | None = None,
+    right: float | None = None,
+    period: float | None = None,
+) -> npt.NDArray[typing.Any]:
     """Interpolate multi-dimensional array over axis."""
     return np.apply_along_axis(
         partial(np.interp, x, xp),
@@ -61,7 +77,11 @@ def ndinterp(x, xp, fp, axis=-1, left=None, right=None, period=None):  # type: i
     )
 
 
-def trapz_product(f, *ff, axis=-1):  # type: ignore[no-untyped-def]
+def trapz_product(
+    f: tuple[npt.NDArray[typing.Any], npt.NDArray[typing.Any]],
+    *ff: tuple[npt.NDArray[typing.Any], npt.NDArray[typing.Any]],
+    axis: int = -1,
+) -> npt.NDArray[typing.Any]:
     """Trapezoidal rule for a product of functions."""
     x, _ = f
     for x_, _ in ff:
@@ -72,10 +92,15 @@ def trapz_product(f, *ff, axis=-1):  # type: ignore[no-untyped-def]
     y = np.interp(x, *f)
     for f_ in ff:
         y *= np.interp(x, *f_)
-    return np.trapz(y, x, axis=axis)  # type: ignore[attr-defined]
+    return np.trapz(y, x, axis=axis)  # type: ignore[attr-defined, no-any-return]
 
 
-def cumtrapz(f, x, dtype=None, out=None):  # type: ignore[no-untyped-def]
+def cumtrapz(
+    f: npt.NDArray[typing.Any],
+    x: npt.NDArray[typing.Any],
+    dtype: np.dtype | None = None,  # type: ignore[type-arg]
+    out: npt.NDArray[typing.Any] | None = None,
+) -> npt.NDArray[typing.Any]:
     """Cumulative trapezoidal rule along last axis."""
     if out is None:
         out = np.empty_like(f, dtype=dtype)

glass/fields.py

@@ -326,7 +326,12 @@ def generate_lognormal(
         yield m
 
 
-def getcl(cls, i, j, lmax=None):  # type: ignore[no-untyped-def]
+def getcl(
+    cls: list[npt.NDArray[typing.Any]],
+    i: int,
+    j: int,
+    lmax: int | None = None,
+) -> npt.NDArray[typing.Any]:
     """
     Return a specific angular power spectrum from an array.
 
@@ -356,8 +361,12 @@ def getcl(cls, i, j, lmax=None):  # type: ignore[no-untyped-def]
     return cl
 
 
-def effective_cls(  # type: ignore[no-untyped-def]
-    cls, weights1, weights2=None, *, lmax=None
+def effective_cls(
+    cls: list[npt.NDArray[typing.Any]],
+    weights1: npt.NDArray[typing.Any],
+    weights2: npt.NDArray[typing.Any] | None = None,
+    *,
+    lmax: int | None = None,
 ) -> npt.NDArray[np.float64]:
     r"""
     Compute effective angular power spectra from weights.
@@ -423,7 +432,7 @@ def effective_cls(  # type: ignore[no-untyped-def]
     for j1, j2 in pairs:
         w1, w2 = weights1[c + j1], weights2[c + j2]
         cl = sum(
-            w1[i1] * w2[i2] * getcl(cls, i1, i2, lmax=lmax)  # type: ignore[no-untyped-call]
+            w1[i1] * w2[i2] * getcl(cls, i1, i2, lmax=lmax)
             for i1, i2 in np.ndindex(n, n)
         )
         out[j1 + j2] = cl

glass/galaxies.py

@@ -33,7 +33,7 @@
 
 
 def redshifts(
-    n: int | npt.ArrayLike,
+    n: int | npt.NDArray[typing.Any],
     w: RadialWindow,
     *,
     rng: np.random.Generator | None = None,
@@ -57,13 +57,13 @@ def redshifts(
         Random number generator. If not given, a default RNG is used.
 
     """
-    return redshifts_from_nz(n, w.za, w.wa, rng=rng, warn=False)
+    return redshifts_from_nz(n, w.za, w.wa, rng=rng, warn=False)  # type: ignore[arg-type]
 
 
 def redshifts_from_nz(
-    count: int | npt.ArrayLike,
-    z: npt.ArrayLike,
-    nz: npt.ArrayLike,
+    count: int | npt.NDArray[typing.Any],
+    z: npt.NDArray[typing.Any],
+    nz: npt.NDArray[typing.Any],
     *,
     rng: np.random.Generator | None = None,
     warn: bool = True,
@@ -111,7 +111,7 @@ def redshifts_from_nz(
         rng = np.random.default_rng()
 
     # bring inputs' leading axes into common shape
-    dims, count, z, nz = broadcast_leading_axes((count, 0), (z, 1), (nz, 1))  # type: ignore[no-untyped-call]
+    dims, count, z, nz = broadcast_leading_axes((count, 0), (z, 1), (nz, 1))  # type: ignore[arg-type, assignment]
 
     # list of results for all dimensions
     redshifts = np.empty(count.sum())  # type: ignore[union-attr]
@@ -122,16 +122,16 @@ def redshifts_from_nz(
     # go through extra dimensions; also works if dims is empty
     for k in np.ndindex(dims):
         # compute the CDF of each galaxy population
-        cdf = cumtrapz(nz[k], z[k], dtype=float)  # type: ignore[call-overload, index, no-untyped-call]
+        cdf = cumtrapz(nz[k], z[k], dtype=float)  # type: ignore[arg-type]
         cdf /= cdf[-1]
 
         # sample redshifts and store result
-        redshifts[total : total + count[k]] = np.interp(  # type: ignore[call-overload, index, misc, operator]
-            rng.uniform(0, 1, size=count[k]),  # type: ignore[arg-type, call-overload, index]
+        redshifts[total : total + count[k]] = np.interp(  # type: ignore[index]
+            rng.uniform(0, 1, size=count[k]),  # type: ignore[index]
             cdf,
-            z[k],  # type: ignore[arg-type, call-overload, index]
+            z[k],
         )
-        total += count[k]  # type: ignore[assignment, call-overload, index, operator]
+        total += count[k]  # type: ignore[index]
 
     assert total == redshifts.size  # noqa: S101
 
@@ -206,11 +206,11 @@ def galaxy_shear(  # noqa: PLR0913
 
 
 def gaussian_phz(
-    z: npt.ArrayLike,
-    sigma_0: float | npt.ArrayLike,
+    z: npt.NDArray[typing.Any],
+    sigma_0: float | npt.NDArray[typing.Any],
     *,
-    lower: npt.ArrayLike | None = None,
-    upper: npt.ArrayLike | None = None,
+    lower: npt.NDArray[typing.Any] | None = None,
+    upper: npt.NDArray[typing.Any] | None = None,
     rng: np.random.Generator | None = None,
 ) -> npt.NDArray[typing.Any]:
     r"""
@@ -264,26 +264,26 @@ def gaussian_phz(
     sigma = np.add(1, z) * sigma_0
     dims = np.shape(sigma)
 
-    zphot = rng.normal(z, sigma)  # type: ignore[arg-type]
+    zphot = rng.normal(z, sigma)
 
     if lower is None:
-        lower = 0.0
+        lower = 0.0  # type: ignore[assignment]
     if upper is None:
-        upper = np.inf
+        upper = np.inf  # type: ignore[assignment]
 
     if not np.all(lower < upper):  # type: ignore[operator]
         msg = "requires lower < upper"
         raise ValueError(msg)
 
     if not dims:
         while zphot < lower or zphot > upper:  # type: ignore[operator]
-            zphot = rng.normal(z, sigma)  # type: ignore[arg-type]
+            zphot = rng.normal(z, sigma)
     else:
         z = np.broadcast_to(z, dims)
         trunc = np.where((zphot < lower) | (zphot > upper))[0]  # type: ignore[operator]
         while trunc.size:
-            znew = rng.normal(z[trunc], sigma[trunc])  # type: ignore[arg-type]
-            zphot[trunc] = znew  # type: ignore[index]
+            znew = rng.normal(z[trunc], sigma[trunc])
+            zphot[trunc] = znew
             trunc = trunc[(znew < lower) | (znew > upper)]  # type: ignore[operator]
 
-    return zphot  # type: ignore[return-value]
+    return zphot

glass/lensing.py

@@ -293,9 +293,9 @@ def add_window(self, delta: npt.NDArray[typing.Any], w: RadialWindow) -> None:
 
         """
         zsrc = w.zeff
-        lens_weight = np.trapz(w.wa, w.za) / np.interp(zsrc, w.za, w.wa)  # type: ignore[attr-defined]
+        lens_weight = np.trapz(w.wa, w.za) / np.interp(zsrc, w.za, w.wa)  # type: ignore[arg-type, attr-defined]
 
-        self.add_plane(delta, zsrc, lens_weight)
+        self.add_plane(delta, zsrc, lens_weight)  # type: ignore[arg-type]
 
     def add_plane(
         self, delta: npt.NDArray[typing.Any], zsrc: float, wlens: float = 1.0
@@ -366,7 +366,7 @@ def wlens(self) -> float:
 def multi_plane_matrix(
     shells: collections.abc.Sequence[RadialWindow],
     cosmo: Cosmology,
-) -> npt.ArrayLike:
+) -> npt.NDArray[typing.Any]:
     """Compute the matrix of lensing contributions from each shell."""
     mpc = MultiPlaneConvergence(cosmo)
     wmat = np.eye(len(shells))
@@ -377,10 +377,10 @@ def multi_plane_matrix(
 
 
 def multi_plane_weights(
-    weights: npt.ArrayLike,
+    weights: npt.NDArray[typing.Any],
     shells: collections.abc.Sequence[RadialWindow],
     cosmo: Cosmology,
-) -> npt.ArrayLike:
+) -> npt.NDArray[typing.Any]:
     """
     Compute effective weights for multi-plane convergence.
 
@@ -413,11 +413,13 @@ def multi_plane_weights(
     weights = weights / np.sum(weights, axis=0)
     # combine weights and the matrix of lensing contributions
     mat = multi_plane_matrix(shells, cosmo)
-    return np.matmul(mat.T, weights)  # type: ignore[no-any-return, union-attr]
+    return np.matmul(mat.T, weights)  # type: ignore[no-any-return]
 
 
 def deflect(
-    lon: npt.ArrayLike, lat: npt.ArrayLike, alpha: npt.ArrayLike
+    lon: npt.NDArray[typing.Any],
+    lat: npt.NDArray[typing.Any],
+    alpha: npt.NDArray[typing.Any],
 ) -> npt.NDArray[typing.Any]:
     r"""
     Apply deflections to positions.

glass/observations.py

@@ -86,10 +86,10 @@ def vmap_galactic_ecliptic(
 
 def gaussian_nz(
     z: npt.NDArray[typing.Any],
-    mean: npt.ArrayLike,
-    sigma: npt.ArrayLike,
+    mean: npt.NDArray[typing.Any],
+    sigma: npt.NDArray[typing.Any],
     *,
-    norm: npt.ArrayLike | None = None,
+    norm: npt.NDArray[typing.Any] | None = None,
 ) -> npt.NDArray[typing.Any]:
     r"""
     Gaussian redshift distribution.
@@ -131,11 +131,11 @@ def gaussian_nz(
 
 def smail_nz(
     z: npt.NDArray[typing.Any],
-    z_mode: npt.ArrayLike,
-    alpha: npt.ArrayLike,
-    beta: npt.ArrayLike,
+    z_mode: npt.NDArray[typing.Any],
+    alpha: npt.NDArray[typing.Any],
+    beta: npt.NDArray[typing.Any],
     *,
-    norm: npt.ArrayLike | None = None,
+    norm: float | npt.NDArray[typing.Any] | None = None,
 ) -> npt.NDArray[typing.Any]:
     r"""
     Redshift distribution following Smail et al. (1994).
@@ -262,7 +262,7 @@ def equal_dens_zbins(
     # first compute the cumulative integral (by trapezoidal rule)
     # then normalise: the first z is at CDF = 0, the last z at CDF = 1
     # interpolate to find the z values at CDF = i/nbins for i = 0, ..., nbins
-    cuml_nz = cumtrapz(nz, z)  # type: ignore[no-untyped-call]
+    cuml_nz = cumtrapz(nz, z)
     cuml_nz /= cuml_nz[[-1]]
     zbinedges = np.interp(np.linspace(0, 1, nbins + 1), cuml_nz, z)