Merge branch 'paddy/issue-358' into paddy/issue-386

glass-dev · Nov 6, 2024 · 11d8af1 · 11d8af1
2 parents cb45af1 + 113d4aa
commit 11d8af1
Show file tree

Hide file tree

Showing 7 changed files with 53 additions and 72 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -69,5 +69,6 @@ repos:
     rev: v1.13.0
     hooks:
       - id: mypy
+        files: ^glass/
         additional_dependencies:
           - numpy
diff --git a/glass/core/array.py b/glass/core/array.py
@@ -8,6 +8,9 @@
 import numpy as np
 import numpy.typing as npt
 
+if typing.TYPE_CHECKING:
+    import collections.abc
+
 
 def broadcast_first(
     *arrays: npt.NDArray[np.float64],
@@ -65,8 +68,8 @@ def broadcast_leading_axes(
 
 def ndinterp(  # noqa: PLR0913
     x: float | npt.NDArray[np.float64],
-    xp: npt.NDArray[np.float64],
-    fp: npt.NDArray[np.float64],
+    xp: collections.abc.Sequence[float] | npt.NDArray[np.float64],
+    fp: collections.abc.Sequence[float] | npt.NDArray[np.float64],
     axis: int = -1,
     left: float | None = None,
     right: float | None = None,

diff --git a/glass/fields.py b/glass/fields.py
@@ -26,6 +26,7 @@
 
 from __future__ import annotations
 
+import collections.abc
 import typing
 import warnings
 
@@ -34,8 +35,9 @@
 import numpy.typing as npt
 from gaussiancl import gaussiancl
 
-if typing.TYPE_CHECKING:
-    import collections.abc
+Cls = collections.abc.Sequence[
+    typing.Union[npt.NDArray[np.float64], collections.abc.Sequence[float]]
+]
 
 
 def iternorm(
@@ -97,9 +99,7 @@ def iternorm(
 
 
 def cls2cov(
-    cls: collections.abc.Sequence[
-        npt.NDArray[np.float64] | collections.abc.Sequence[float]
-    ],
+    cls: Cls,
     nl: int,
     nf: int,
     nc: int,
@@ -135,12 +135,10 @@ def multalm(
 
 
 def transform_cls(
-    cls: collections.abc.Sequence[
-        npt.NDArray[np.float64] | collections.abc.Sequence[float]
-    ],
+    cls: Cls,
     tfm: str | typing.Callable[[npt.NDArray[np.float64]], npt.NDArray[np.float64]],
     pars: tuple[typing.Any, ...] = (),
-) -> list[list[float] | npt.NDArray[np.float64]]:
+) -> Cls:
     """Transform Cls to Gaussian Cls."""
     gls = []
     for cl in cls:
@@ -159,14 +157,12 @@ def transform_cls(
 
 
 def discretized_cls(
-    cls: collections.abc.Sequence[
-        npt.NDArray[np.float64] | collections.abc.Sequence[float]
-    ],
+    cls: Cls,
     *,
     lmax: int | None = None,
     ncorr: int | None = None,
     nside: int | None = None,
-) -> list[npt.NDArray[np.float64] | collections.abc.Sequence[float]]:
+) -> Cls:
     """
     Apply discretisation effects to angular power spectra.
 
@@ -203,19 +199,15 @@ def discretized_cls(
 
 
 def lognormal_gls(
-    cls: collections.abc.Sequence[
-        npt.NDArray[np.float64] | collections.abc.Sequence[float]
-    ],
+    cls: Cls,
     shift: float = 1.0,
-) -> list[list[float] | npt.NDArray[np.float64]]:
+) -> Cls:
     """Compute Gaussian Cls for a lognormal random field."""
     return transform_cls(cls, "lognormal", (shift,))
 
 
 def generate_gaussian(
-    gls: collections.abc.Sequence[
-        npt.NDArray[np.float64] | collections.abc.Sequence[float]
-    ],
+    gls: Cls,
     nside: int,
     *,
     ncorr: int | None = None,
@@ -300,9 +292,7 @@ def generate_gaussian(
 
 
 def generate_lognormal(
-    gls: collections.abc.Sequence[
-        npt.NDArray[np.float64] | collections.abc.Sequence[float]
-    ],
+    gls: Cls,
     nside: int,
     shift: float = 1.0,
     *,

diff --git a/glass/lensing.py b/glass/lensing.py
@@ -389,7 +389,7 @@ def __init__(self, cosmo: Cosmology) -> None:
         self.z2: float = 0.0
         self.z3: float = 0.0
         self.x3: float = 0.0
-        self.w3: npt.NDArray[np.float64] | float = 0.0
+        self.w3: float = 0.0
         self.r23: float = 1.0
         self.delta3: npt.NDArray[np.float64] = np.array(0.0)
         self.kappa2: npt.NDArray[np.float64] | None = None
@@ -404,13 +404,16 @@ def add_window(self, delta: npt.NDArray[np.float64], w: RadialWindow) -> None:
 
         """
         zsrc = w.zeff
-        lens_weight = np.trapz(  # type: ignore[attr-defined]
-            w.wa,
-            w.za,
-        ) / np.interp(
-            zsrc,
-            w.za,
-            w.wa,
+        lens_weight = float(
+            np.trapz(  # type: ignore[attr-defined]
+                w.wa,
+                w.za,
+            )
+            / np.interp(
+                zsrc,
+                w.za,
+                w.wa,
+            )
         )
 
         self.add_plane(delta, zsrc, lens_weight)
@@ -419,7 +422,7 @@ def add_plane(
         self,
         delta: npt.NDArray[np.float64],
         zsrc: float,
-        wlens: float | npt.NDArray[np.float64] = 1.0,
+        wlens: float = 1.0,
     ) -> None:
         """Add a mass plane at redshift ``zsrc`` to the convergence."""
         if zsrc <= self.z3:
@@ -479,7 +482,7 @@ def delta(self) -> npt.NDArray[np.float64]:
         return self.delta3
 
     @property
-    def wlens(self) -> float | npt.NDArray[np.float64]:
+    def wlens(self) -> float:
         """The weight of the current matter plane."""
         return self.w3
 

diff --git a/glass/shapes.py b/glass/shapes.py
@@ -31,7 +31,7 @@
 def triaxial_axis_ratio(
     zeta: float | npt.NDArray[np.float64],
     xi: float | npt.NDArray[np.float64],
-    size: int | tuple[int, ...] = (),
+    size: int | tuple[int, ...] | None = None,
     *,
     rng: np.random.Generator | None = None,
 ) -> npt.NDArray[np.float64]:
@@ -69,7 +69,7 @@ def triaxial_axis_ratio(
         rng = np.random.default_rng()
 
     # get size from inputs if not explicitly provided
-    if not size:
+    if size is None:
         size = np.broadcast(zeta, xi).shape
 
     # draw random viewing angle (theta, phi)
@@ -101,7 +101,7 @@ def ellipticity_ryden04(  # noqa: PLR0913
     sigma: float | npt.NDArray[np.float64],
     gamma: float | npt.NDArray[np.float64],
     sigma_gamma: float | npt.NDArray[np.float64],
-    size: int | tuple[int, ...] = (),
+    size: int | tuple[int, ...] | None = None,
     *,
     rng: np.random.Generator | None = None,
 ) -> npt.NDArray[np.float64]:
@@ -143,7 +143,7 @@ def ellipticity_ryden04(  # noqa: PLR0913
         rng = np.random.default_rng()
 
     # default size if not given
-    if not size:
+    if size is None:
         size = np.broadcast(mu, sigma, gamma, sigma_gamma).shape
 
     # broadcast all inputs to output shape

diff --git a/glass/shells.py b/glass/shells.py
@@ -44,6 +44,7 @@
 
 from __future__ import annotations
 
+import collections.abc
 import typing
 import warnings
 
@@ -53,10 +54,11 @@
 from glass.core.array import ndinterp
 
 if typing.TYPE_CHECKING:
-    import collections.abc
-
     from cosmology import Cosmology
 
+ArrayLike1D = typing.Union[collections.abc.Sequence[float], npt.NDArray[np.float64]]
+WeightFunc = typing.Callable[[ArrayLike1D], npt.NDArray[np.float64]]
+
 
 def distance_weight(
     z: npt.NDArray[np.float64],
@@ -132,13 +134,9 @@ class RadialWindow(typing.NamedTuple):
 
 
 def tophat_windows(
-    zbins: npt.NDArray[np.float64],
+    zbins: ArrayLike1D,
     dz: float = 1e-3,
-    weight: typing.Callable[
-        [list[float] | npt.NDArray[np.float64]],
-        npt.NDArray[np.float64],
-    ]
-    | None = None,
+    weight: WeightFunc | None = None,
 ) -> list[RadialWindow]:
     """
     Tophat window functions from the given redshift bin edges.
@@ -179,13 +177,7 @@ def tophat_windows(
             stacklevel=2,
         )
 
-    wht: (
-        npt.NDArray[np.float64]
-        | typing.Callable[
-            [list[float] | npt.NDArray[np.float64]],
-            npt.NDArray[np.float64],
-        ]
-    )
+    wht: WeightFunc
     wht = weight if weight is not None else np.ones_like
     ws = []
     for zmin, zmax in zip(zbins, zbins[1:]):
@@ -204,13 +196,9 @@ def tophat_windows(
 
 
 def linear_windows(
-    zgrid: npt.NDArray[np.float64],
+    zgrid: ArrayLike1D,
     dz: float = 1e-3,
-    weight: typing.Callable[
-        [list[float] | npt.NDArray[np.float64]],
-        npt.NDArray[np.float64],
-    ]
-    | None = None,
+    weight: WeightFunc | None = None,
 ) -> list[RadialWindow]:
     """
     Linear interpolation window functions.
@@ -265,13 +253,9 @@ def linear_windows(
 
 
 def cubic_windows(
-    zgrid: npt.NDArray[np.float64],
+    zgrid: ArrayLike1D,
     dz: float = 1e-3,
-    weight: typing.Callable[
-        [list[float] | npt.NDArray[np.float64]],
-        npt.NDArray[np.float64],
-    ]
-    | None = None,
+    weight: WeightFunc | None = None,
 ) -> list[RadialWindow]:
     """
     Cubic interpolation window functions.
@@ -327,8 +311,8 @@ def cubic_windows(
 
 
 def restrict(
-    z: npt.NDArray[np.float64],
-    f: npt.NDArray[np.float64],
+    z: ArrayLike1D,
+    f: ArrayLike1D,
     w: RadialWindow,
 ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
     """

diff --git a/glass/user.py b/glass/user.py
@@ -74,7 +74,7 @@ class _FitsWriter:
     Initialised with the fits object and extension name.
     """
 
-    def __init__(self, fits: fitsio.FITS, ext: str = "") -> None:
+    def __init__(self, fits: fitsio.FITS, ext: str | None = None) -> None:
         """Create a new, uninitialised writer."""
         self.fits = fits
         self.ext = ext
@@ -85,9 +85,9 @@ def _append(
         names: list[str] | None = None,
     ) -> None:
         """Write the FITS file."""
-        if not self.ext or self.ext not in self.fits:
+        if self.ext is None or self.ext not in self.fits:
             self.fits.write_table(data, names=names, extname=self.ext)
-            if not self.ext:
+            if self.ext is None:
                 self.ext = self.fits[-1].get_extnum()
         else:
             hdu = self.fits[self.ext]
@@ -120,7 +120,7 @@ def write(
 def write_catalog(
     filename: pathlib.Path,
     *,
-    ext: str = "",
+    ext: str | None = None,
 ) -> collections.abc.Generator[_FitsWriter]:
     """
     Write a catalogue into a FITS file.