astropy · pllim · Jul 3, 2024 · Jun 14, 2024 · Jun 14, 2024 · Jun 14, 2024
@@ -45,19 +45,36 @@
    https://www.tandfonline.com/doi/abs/10.1080/01621459.1969.10501038
 """
 
+from __future__ import annotations
+
 import warnings
 from inspect import signature
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from astropy.utils.exceptions import AstropyUserWarning
 
+if TYPE_CHECKING:
+    from collections.abc import KeysView
+    from typing import Literal
+
+    from numpy.typing import ArrayLike, NDArray
+
+# TODO: typing: use a custom-defined 'ArrayLike-but-not-a-scalar' type for `float | ArrayLike` or `ArrayLike | float` hints
+
 # TODO: implement other fitness functions from appendix C of Scargle 2013
 
 __all__ = ["FitnessFunc", "Events", "RegularEvents", "PointMeasures", "bayesian_blocks"]
 
 
-def bayesian_blocks(t, x=None, sigma=None, fitness="events", **kwargs):
+def bayesian_blocks(
+    t: ArrayLike,
+    x: ArrayLike | None = None,
+    sigma: ArrayLike | float | None = None,
+    fitness: Literal["events", "regular_events", "measures"] | FitnessFunc = "events",
+    **kwargs,
+) -> NDArray[float]:
     r"""Compute optimal segmentation of data with Scargle's Bayesian Blocks.
 
     This is a flexible implementation of the Bayesian Blocks algorithm
@@ -210,12 +227,22 @@ class FitnessFunc:
        https://ui.adsabs.harvard.edu/abs/2013ApJ...764..167S
     """
 
-    def __init__(self, p0=0.05, gamma=None, ncp_prior=None):
+    def __init__(
+        self,
+        p0: float = 0.05,
+        gamma: float | None = None,
+        ncp_prior: float | None = None,
+    ) -> None:
         self.p0 = p0
         self.gamma = gamma
         self.ncp_prior = ncp_prior
 
-    def validate_input(self, t, x=None, sigma=None):
+    def validate_input(
+        self,
+        t: ArrayLike,
+        x: ArrayLike | None = None,
+        sigma: float | ArrayLike | None = None,
+    ) -> tuple[NDArray[float], NDArray[float], NDArray[float]]:
         """Validate inputs to the model.
 
         Parameters
@@ -229,7 +256,7 @@ def validate_input(self, t, x=None, sigma=None):
 
         Returns
         -------
-        t, x, sigma : array-like, float or None
+        t, x, sigma : array-like, float
             validated and perhaps modified versions of inputs
         """
         # validate array input
@@ -246,7 +273,7 @@ def validate_input(self, t, x=None, sigma=None):
             if sigma is not None:
                 raise ValueError("If sigma is specified, x must be specified")
             else:
-                sigma = 1
+                sigma = 1.0
 
             if len(unq_t) == len(t):
                 x = np.ones_like(t)
@@ -273,7 +300,7 @@ def validate_input(self, t, x=None, sigma=None):
 
         # verify the given sigma value
         if sigma is None:
-            sigma = 1
+            sigma = 1.0
         else:
             sigma = np.asarray(sigma, dtype=float)
             if sigma.shape not in [(), (1,), (t.size,)]:
@@ -284,7 +311,7 @@ def validate_input(self, t, x=None, sigma=None):
     def fitness(self, **kwargs):
         raise NotImplementedError()
 
-    def p0_prior(self, N):
+    def p0_prior(self, N: int) -> float:
         """Empirical prior, parametrized by the false alarm probability ``p0``.
 
         See eq. 21 in Scargle (2013).
@@ -298,10 +325,10 @@ def p0_prior(self, N):
     # the fitness_args property will return the list of arguments accepted by
     # the method fitness().  This allows more efficient computation below.
     @property
-    def _fitness_args(self):
+    def _fitness_args(self) -> KeysView[str]:
         return signature(self.fitness).parameters.keys()
 
-    def compute_ncp_prior(self, N):
+    def compute_ncp_prior(self, N: int) -> float:
         """
         If ``ncp_prior`` is not explicitly defined, compute it from ``gamma``
         or ``p0``.
@@ -316,7 +343,12 @@ def compute_ncp_prior(self, N):
                 "``gamma`` nor ``p0`` is defined."
             )
 
-    def fit(self, t, x=None, sigma=None):
+    def fit(
+        self,
+        t: ArrayLike,
+        x: ArrayLike | None = None,
+        sigma: ArrayLike | float | None = None,
+    ) -> NDArray[float]:
         """Fit the Bayesian Blocks model given the specified fitness function.
 
         Parameters
@@ -439,11 +471,16 @@ class Events(FitnessFunc):
         If ``ncp_prior`` is specified, ``gamma`` and ``p0`` is ignored.
     """
 
-    def fitness(self, N_k, T_k):
+    def fitness(self, N_k: NDArray[float], T_k: NDArray[float]) -> NDArray[float]:
         # eq. 19 from Scargle 2013
         return N_k * (np.log(N_k / T_k))
 
-    def validate_input(self, t, x, sigma):
+    def validate_input(
+        self,
+        t: ArrayLike,
+        x: ArrayLike | None,
+        sigma: float | ArrayLike | None,
+    ) -> tuple[NDArray[float], NDArray[float], NDArray[float]]:
         t, x, sigma = super().validate_input(t, x, sigma)
         if x is not None and np.any(x % 1 > 0):
             raise ValueError("x must be integer counts for fitness='events'")
@@ -465,24 +502,39 @@ class RegularEvents(FitnessFunc):
         False alarm probability, used to compute the prior on :math:`N_{\rm
         blocks}` (see eq. 21 of Scargle 2013). If gamma is specified, p0 is
         ignored.
+    gamma : float, optional
+        If specified, then use this gamma to compute the general prior form,
+        :math:`p \sim {\tt gamma}^{N_{\rm blocks}}`.  If gamma is specified, p0
+        is ignored.
     ncp_prior : float, optional
         If specified, use the value of ``ncp_prior`` to compute the prior as
         above, using the definition :math:`{\tt ncp\_prior} = -\ln({\tt
         gamma})`.  If ``ncp_prior`` is specified, ``gamma`` and ``p0`` are
         ignored.
     """
 
-    def __init__(self, dt, p0=0.05, gamma=None, ncp_prior=None):
+    def __init__(
+        self,
+        dt: float,
+        p0: float | None = 0.05,
+        gamma: float | None = None,
+        ncp_prior: float | None = None,
+    ) -> None:
         self.dt = dt
         super().__init__(p0, gamma, ncp_prior)
 
-    def validate_input(self, t, x, sigma):
+    def validate_input(
+        self,
+        t: ArrayLike,
+        x: ArrayLike | None = None,
+        sigma: float | ArrayLike | None = None,
+    ) -> tuple[NDArray[float], NDArray[float], NDArray[float]]:
         t, x, sigma = super().validate_input(t, x, sigma)
         if not np.all((x == 0) | (x == 1)):
             raise ValueError("Regular events must have only 0 and 1 in x")
         return t, x, sigma
 
-    def fitness(self, T_k, N_k):
+    def fitness(self, T_k: NDArray[float], N_k: NDArray[float]) -> NDArray[float]:
         # Eq. C23 of Scargle 2013
         M_k = T_k / self.dt
         N_over_M = N_k / M_k
@@ -510,21 +562,35 @@ class PointMeasures(FitnessFunc):
         False alarm probability, used to compute the prior on :math:`N_{\rm
         blocks}` (see eq. 21 of Scargle 2013). If gamma is specified, p0 is
         ignored.
+    gamma : float, optional
+        If specified, then use this gamma to compute the general prior form,
+        :math:`p \sim {\tt gamma}^{N_{\rm blocks}}`.  If gamma is specified, p0
+        is ignored.
     ncp_prior : float, optional
         If specified, use the value of ``ncp_prior`` to compute the prior as
         above, using the definition :math:`{\tt ncp\_prior} = -\ln({\tt
         gamma})`.  If ``ncp_prior`` is specified, ``gamma`` and ``p0`` are
         ignored.
     """
 
-    def __init__(self, p0=0.05, gamma=None, ncp_prior=None):
+    def __init__(
+        self,
+        p0: float | None = 0.05,
+        gamma: float | None = None,
+        ncp_prior: float | None = None,
+    ) -> None:
         super().__init__(p0, gamma, ncp_prior)
 
-    def fitness(self, a_k, b_k):
+    def fitness(self, a_k: NDArray[float], b_k: ArrayLike) -> NDArray[float]:
         # eq. 41 from Scargle 2013
         return (b_k * b_k) / (4 * a_k)
 
-    def validate_input(self, t, x, sigma):
+    def validate_input(
+        self,
+        t: ArrayLike,
+        x: ArrayLike | None,
+        sigma: float | ArrayLike | None,
+    ) -> tuple[NDArray[float], NDArray[float], NDArray[float]]:
         if x is None:
             raise ValueError("x must be specified for point measures")
         return super().validate_input(t, x, sigma)
@@ -4,10 +4,21 @@
 Tukey's biweight function.
 """
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import numpy as np
 
 from .funcs import median_absolute_deviation
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from numpy.typing import ArrayLike, NDArray
+
+# TODO: typing: use a custom-defined 'ArrayLike-but-not-a-scalar' type for `float | ArrayLike` or `ArrayLike | float` hints
+
 __all__ = [
     "biweight_location",
     "biweight_scale",
@@ -17,7 +28,11 @@
 ]
 
 
-def _stat_functions(data, ignore_nan=False):
+def _stat_functions(
+    data: ArrayLike,
+    ignore_nan: bool | None = False,
+) -> tuple[Callable[..., NDArray[float]], Callable[..., NDArray[float]]]:
+    # TODO: typing: update return Callables with custom callback protocol (https://mypy.readthedocs.io/en/stable/protocols.html#callback-protocols)
     if isinstance(data, np.ma.MaskedArray):
         median_func = np.ma.median
         sum_func = np.ma.sum
@@ -31,7 +46,14 @@ def _stat_functions(data, ignore_nan=False):
     return median_func, sum_func
 
 
-def biweight_location(data, c=6.0, M=None, axis=None, *, ignore_nan=False):
+def biweight_location(
+    data: ArrayLike,
+    c: float | None = 6.0,
+    M: float | ArrayLike | None = None,
+    axis: int | tuple[int, ...] | None = None,
+    *,
+    ignore_nan: bool | None = False,
+) -> float | NDArray[float]:
     r"""
     Compute the biweight location.
 
@@ -73,7 +95,7 @@ def biweight_location(data, c=6.0, M=None, axis=None, *, ignore_nan=False):
         ``axis`` of the input array.  If `None` (default), then the
         median of the input array will be used (or along each ``axis``,
         if specified).
-    axis : None, int, or tuple of int, optional
+    axis : int or tuple of int, optional
         The axis or axes along which the biweight locations are
         computed.  If `None` (default), then the biweight location of
         the flattened input array will be computed.
@@ -160,8 +182,14 @@ def biweight_location(data, c=6.0, M=None, axis=None, *, ignore_nan=False):
 
 
 def biweight_scale(
-    data, c=9.0, M=None, axis=None, modify_sample_size=False, *, ignore_nan=False
-):
+    data: ArrayLike,
+    c: float | None = 9.0,
+    M: float | ArrayLike | None = None,
+    axis: int | tuple[int, ...] | None = None,
+    modify_sample_size: bool | None = False,
+    *,
+    ignore_nan: bool | None = False,
+) -> float | NDArray[float]:
     r"""
     Compute the biweight scale.
 
@@ -222,7 +250,7 @@ def biweight_scale(
         containing the location estimate along each ``axis`` of the
         input array.  If `None` (default), then the median of the input
         array will be used (or along each ``axis``, if specified).
-    axis : None, int, or tuple of int, optional
+    axis : int or tuple of int, optional
         The axis or axes along which the biweight scales are computed.
         If `None` (default), then the biweight scale of the flattened
         input array will be computed.
@@ -280,8 +308,14 @@ def biweight_scale(
 
 
 def biweight_midvariance(
-    data, c=9.0, M=None, axis=None, modify_sample_size=False, *, ignore_nan=False
-):
+    data: ArrayLike,
+    c: float | None = 9.0,
+    M: float | ArrayLike | None = None,
+    axis: int | tuple[int, ...] | None = None,
+    modify_sample_size: bool | None = False,
+    *,
+    ignore_nan: bool | None = False,
+) -> float | NDArray[float]:
     r"""
     Compute the biweight midvariance.
 
@@ -341,7 +375,7 @@ def biweight_midvariance(
         containing the location estimate along each ``axis`` of the
         input array.  If `None` (default), then the median of the input
         array will be used (or along each ``axis``, if specified).
-    axis : None, int, or tuple of int, optional
+    axis : int or tuple of int, optional
         The axis or axes along which the biweight midvariances are
         computed.  If `None` (default), then the biweight midvariance of
         the flattened input array will be computed.
@@ -456,7 +490,12 @@ def biweight_midvariance(
         return where_func(mad.squeeze() == 0, 0.0, value)
 
 
-def biweight_midcovariance(data, c=9.0, M=None, modify_sample_size=False):
+def biweight_midcovariance(
+    data: ArrayLike,
+    c: float | None = 9.0,
+    M: float | ArrayLike | None = None,
+    modify_sample_size: bool | None = False,
+) -> NDArray[float]:
     r"""
     Compute the biweight midcovariance between pairs of multiple
     variables.
@@ -667,7 +706,13 @@ def biweight_midcovariance(data, c=9.0, M=None, modify_sample_size=False):
         return value
 
 
-def biweight_midcorrelation(x, y, c=9.0, M=None, modify_sample_size=False):
+def biweight_midcorrelation(
+    x: ArrayLike,
+    y: ArrayLike,
+    c: float | None = 9.0,
+    M: float | ArrayLike | None = None,
+    modify_sample_size: bool | None = False,
+) -> float:
     r"""
     Compute the biweight midcorrelation between two variables.