fix ndarray type annotations for fbcode/kats

Summary: Change numpy.ndarray type annotations to numpy.typing.NDArray to unblock PSS2 upgrades.

Reviewed By: florazzz

Differential Revision: D63653820

fbshipit-source-id: d6cf42ce71f4fd5ca6c614f09fc29a11a393da59

kats/compat/sklearn.py

@@ -19,7 +19,7 @@
 def mean_squared_error(
     y_true: npt.NDArray,
     y_pred: npt.NDArray,
-    sample_weight: Optional[np.ndarray] = None,
+    sample_weight: Optional[npt.NDArray] = None,
     multioutput: str = "uniform_average",
     squared: bool = True,
 ) -> float:
@@ -36,7 +36,7 @@ def mean_squared_error(
 def mean_squared_log_error(
     y_true: npt.NDArray,
     y_pred: npt.NDArray,
-    sample_weight: Optional[np.ndarray] = None,
+    sample_weight: Optional[npt.NDArray] = None,
     multioutput: str = "uniform_average",
     squared: bool = True,
 ) -> float:

kats/consts.py

@@ -239,7 +239,7 @@ def __init__(  # noqa C901
         use_unix_time: bool = False,
         unix_time_units: str = "ns",
         tz: Optional[str] = None,
-        tz_ambiguous: Union[str, np.ndarray] = "raise",
+        tz_ambiguous: Union[str, npt.NDArray] = "raise",
         tz_nonexistent: str = "raise",
         categorical_var: Optional[List[str]] = None,
         drop_duplicate_time: bool = False,
@@ -574,7 +574,7 @@ def _set_time_format(
         use_unix_time: Optional[bool],
         unix_time_units: Optional[str],
         tz: Optional[str] = None,
-        tz_ambiguous: Union[str, np.ndarray] = "raise",
+        tz_ambiguous: Union[str, npt.NDArray] = "raise",
         tz_nonexistent: str = "raise",
         cache_datetimes: bool = True,
     ) -> pd.core.series.Series:
@@ -1092,7 +1092,7 @@ def convert_timezone(self, tz: str) -> None:
     def set_timezone(
         self,
         tz: str,
-        tz_ambiguous: Union[str, np.ndarray] = "raise",
+        tz_ambiguous: Union[str, npt.NDArray] = "raise",
         tz_nonexistent: str = "raise",
         sort_by_time: bool = True,
     ) -> None:

kats/detectors/bocpd.py

@@ -230,7 +230,7 @@ class TrendChangeParameters(BOCPDModelParameters):
             debug.
     """
 
-    mu_prior: Optional[np.ndarray] = None
+    mu_prior: Optional[npt.NDArray] = None
     num_likelihood_samples: int = 100
     num_points_prior: int = _MIN_POINTS
     readjust_sigma_prior: bool = False
@@ -641,7 +641,7 @@ def group_changepoints_by_timeseries(
 
         return dict(change_points_per_ts)
 
-    def get_change_prob(self) -> Dict[str, np.ndarray]:
+    def get_change_prob(self) -> Dict[str, npt.NDArray]:
         """Returns the probability of being a changepoint.
 
         Args:
@@ -658,7 +658,7 @@ def get_change_prob(self) -> Dict[str, np.ndarray]:
             raise ValueError("detector needs to be run before getting prob")
         return self.change_prob
 
-    def get_run_length_matrix(self) -> Dict[str, np.ndarray]:
+    def get_run_length_matrix(self) -> Dict[str, npt.NDArray]:
         """Returns the entire run-time posterior.
         Args:
             None.
@@ -724,11 +724,11 @@ class _BayesOnlineChangePoint(Detector):
         maximum values diagonally.
     """
 
-    rt_posterior: Optional[np.ndarray] = None
-    pred_mean_arr: Optional[np.ndarray] = None
-    pred_std_arr: Optional[np.ndarray] = None
-    next_pred_prob: Optional[np.ndarray] = None
-    threshold: Optional[np.ndarray] = None
+    rt_posterior: Optional[npt.NDArray] = None
+    pred_mean_arr: Optional[npt.NDArray] = None
+    pred_std_arr: Optional[npt.NDArray] = None
+    next_pred_prob: Optional[npt.NDArray] = None
+    threshold: Optional[npt.NDArray] = None
     posterior_predictive: npt.NDArray
     T: int
     P: int
@@ -778,8 +778,8 @@ def __init__(
     def detector(
         self,
         model: Union[SupportedModelType, "_PredictiveModel"],
-        threshold: Union[float, np.ndarray] = 0.5,
-        changepoint_prior: Union[float, np.ndarray] = 0.01,
+        threshold: Union[float, npt.NDArray] = 0.5,
+        changepoint_prior: Union[float, npt.NDArray] = 0.01,
     ) -> Dict[str, Any]:
         """Runs the actual BOCPD detection algorithm.
 
@@ -953,7 +953,7 @@ def _find_posterior(
 
     def plot(
         self,
-        threshold: Optional[Union[float, np.ndarray]] = None,
+        threshold: Optional[Union[float, npt.NDArray]] = None,
         lag: Optional[int] = None,
         ts_names: Optional[List[str]] = None,
         **kwargs: Any,
@@ -1421,7 +1421,7 @@ class _BayesianLinReg(_PredictiveModel):
         parameters: Specifying all the priors.
     """
 
-    mu_prior: Optional[np.ndarray] = None
+    mu_prior: Optional[npt.NDArray] = None
     prior_regression_numpoints: Optional[int] = None
 
     def __init__(
@@ -1444,16 +1444,16 @@ def __init__(
             f"sigma prior adjustment {readjust_sigma_prior}, "
             f"and plot prior regression {plot_regression_prior}"
         )
-        self._x: Optional[np.ndarray] = None
-        self._y: Optional[np.ndarray] = None
+        self._x: Optional[npt.NDArray] = None
+        self._y: Optional[npt.NDArray] = None
         self.t = 0
 
         # Random numbers I tried out to make the sigma_squared values really large
         self.a_0 = 0.1  # TODO find better priors?
         self.b_0 = 200.0  # TODO
 
         self.all_time: npt.NDArray = np.array(range(data.time.shape[0]))
-        self.all_vals: Union[pd.DataFrame, pd.Series, np.ndarray] = data.value
+        self.all_vals: Union[pd.DataFrame, pd.Series, npt.NDArray] = data.value
 
         self.lambda_prior: npt.NDArray = np.multiply(2e-7, np.identity(2))
 
@@ -1530,7 +1530,7 @@ def setup(self) -> None:
     @staticmethod
     def _plot_regression(
         x: npt.NDArray,
-        y: Union[np.ndarray, pd.DataFrame, pd.Series],
+        y: Union[npt.NDArray, pd.DataFrame, pd.Series],
         intercept: float,
         slope: float,
     ) -> None:
@@ -1557,7 +1557,7 @@ def _sample_bayesian_linreg(
         a_n: float,
         b_n: float,
         num_samples: int,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    ) -> Tuple[npt.NDArray, npt.NDArray]:
 
         # this is to make sure the results are consistent
         # and tests don't break randomly
@@ -1580,7 +1580,7 @@ def _sample_bayesian_linreg(
     @staticmethod
     def _compute_bayesian_likelihood(
         beta: npt.NDArray, sigma_squared: npt.NDArray, x: npt.NDArray, val: float
-    ) -> Tuple[float, np.ndarray]:
+    ) -> Tuple[float, npt.NDArray]:
         prediction = np.matmul(beta, x)
         bayesian_likelihoods = norm.pdf(
             val, loc=prediction, scale=np.sqrt(sigma_squared)
@@ -1597,7 +1597,7 @@ def _sample_likelihood(
         x: npt.NDArray,
         val: float,
         num_samples: int,
-    ) -> Tuple[float, np.ndarray, np.ndarray]:
+    ) -> Tuple[float, npt.NDArray, npt.NDArray]:
         (
             all_sample_betas,
             sample_sigma_squared,

kats/detectors/cusum_detection.py

@@ -165,9 +165,9 @@ def __init__(
         confidence: float,
         direction: str,
         cp_index: int,
-        mu0: Union[float, np.ndarray],
-        mu1: Union[float, np.ndarray],
-        delta: Union[float, np.ndarray],
+        mu0: Union[float, npt.NDArray],
+        mu1: Union[float, npt.NDArray],
+        delta: Union[float, npt.NDArray],
         llr_int: float,
         llr: float,
         regression_detected: bool,
@@ -197,15 +197,15 @@ def cp_index(self) -> int:
         return self._cp_index
 
     @property
-    def mu0(self) -> Union[float, np.ndarray]:
+    def mu0(self) -> Union[float, npt.NDArray]:
         return self._mu0
 
     @property
-    def mu1(self) -> Union[float, np.ndarray]:
+    def mu1(self) -> Union[float, npt.NDArray]:
         return self._mu1
 
     @property
-    def delta(self) -> Union[float, np.ndarray]:
+    def delta(self) -> Union[float, npt.NDArray]:
         return self._delta
 
     @property
@@ -840,10 +840,10 @@ def _get_llr(
     def _log_llr_multi(
         self,
         x: npt.NDArray,
-        mu0: Union[float, np.ndarray],
-        sigma0: Union[float, np.ndarray],
-        mu1: Union[float, np.ndarray],
-        sigma1: Union[float, np.ndarray],
+        mu0: Union[float, npt.NDArray],
+        sigma0: Union[float, npt.NDArray],
+        mu1: Union[float, npt.NDArray],
+        sigma1: Union[float, npt.NDArray],
     ) -> float:
         try:
             sigma0_inverse = np.linalg.inv(sigma0)

kats/detectors/detector_consts.py

@@ -191,10 +191,10 @@ class PercentageChange:
         min_perc_change: float, minimum percentage change, for a non zero score. Score will be clipped to zero if the absolute value of the percentage chenge is less than this value
     """
 
-    upper: Optional[Union[float, np.ndarray]]
-    lower: Optional[Union[float, np.ndarray]]
-    _t_score: Optional[Union[float, np.ndarray]]
-    _p_value: Optional[Union[float, np.ndarray]]
+    upper: Optional[Union[float, npt.NDArray]]
+    lower: Optional[Union[float, npt.NDArray]]
+    _t_score: Optional[Union[float, npt.NDArray]]
+    _p_value: Optional[Union[float, npt.NDArray]]
     num_series: int
 
     def __init__(
@@ -225,7 +225,7 @@ def __init__(
         self.min_perc_change = min_perc_change
 
     @property
-    def ratio_estimate(self) -> Union[float, np.ndarray]:
+    def ratio_estimate(self) -> Union[float, npt.NDArray]:
         # pyre-ignore[6]: Expected float for 1st positional only parameter to call float.__truediv__ but got Union[float, np.ndarray].
         return self.current.mean_val / self.previous.mean_val
 
@@ -300,11 +300,11 @@ def p_value(self) -> float:
         return cast(float, self._p_value)
 
     @property
-    def mean_previous(self) -> Union[float, np.ndarray]:
+    def mean_previous(self) -> Union[float, npt.NDArray]:
         return self.previous.mean_val
 
     @property
-    def mean_difference(self) -> Union[float, np.ndarray]:
+    def mean_difference(self) -> Union[float, npt.NDArray]:
         # pyre-ignore[6]: Expected `float` for 1st param but got `Union[float,
         #  np.ndarray]`.
         _mean_diff = self.current.mean_val - self.previous.mean_val

kats/detectors/distribution_distance_model.py

@@ -34,7 +34,7 @@
 _log: logging.Logger = logging.getLogger("distribution_distance_model")
 
 
-def _merge_percentile(l1: npt.NDArray) -> Tuple[np.ndarray, np.ndarray]:
+def _merge_percentile(l1: npt.NDArray) -> Tuple[npt.NDArray, npt.NDArray]:
     """
     handle equal percentile:
     [-2.5, -1.3, -1.3, 1.2, 1.2] -> [-2.5, -1.3, 1.2] with prob [0.2, 0.4, 0.4]
@@ -53,7 +53,7 @@ def _merge_percentile(l1: npt.NDArray) -> Tuple[np.ndarray, np.ndarray]:
 
 def _percentile_to_prob(
     l1: npt.NDArray, l2: npt.NDArray
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> Tuple[npt.NDArray, npt.NDArray]:
     """
     convert decile to probability distribution: ([3,4,5,6,7], [1,2,3,4,5])
     to ([0.2, 0.2, 0.2, 0.2, 0.2], [0.6, 0.2, 0.2, 0, 0])

kats/detectors/hourly_ratio_detection.py

@@ -160,7 +160,7 @@ def _mahalanobis_test(
         median: npt.NDArray,
         cov: npt.NDArray,
         alpha: float = 0.01,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    ) -> Tuple[npt.NDArray, npt.NDArray]:
         """mahalanobis test function.
 
         Args:

kats/detectors/interval_detector.py

@@ -904,8 +904,8 @@ def _get_lowest_p(
     def _mvn_mvnun(
         lower: npt.NDArray,
         upper: npt.NDArray,
-        mean: Optional[np.ndarray] = None,
-        cov: Union[int, np.ndarray] = 1,
+        mean: Optional[npt.NDArray] = None,
+        cov: Union[int, npt.NDArray] = 1,
         allow_singular: bool = False,
         maxpts: Optional[int] = None,
         abseps: float = 1e-6,
@@ -1031,7 +1031,7 @@ def _w_two_tailed(m: int, p: float, cov: npt.NDArray) -> npt.NDArray:
 
     @staticmethod
     def _w(
-        m: int, p: float, test_type: TestType, cov: Optional[np.ndarray] = None
+        m: int, p: float, test_type: TestType, cov: Optional[npt.NDArray] = None
     ) -> npt.NDArray:
         if cov is None:
             return IntervalDetectorModel._w_independent(m=m, p=p)
@@ -1048,7 +1048,7 @@ def _probability_of_at_least_one_m_run_in_n_trials(
         n: int,
         m: int,
         test_type: TestType = TestType.ONE_SIDED_UPPER,
-        cov: Optional[np.ndarray] = None,
+        cov: Optional[npt.NDArray] = None,
     ) -> float:
         """P(at least 1 run of m consecutive rejections) in a vectorized formulation.
 
@@ -1203,7 +1203,7 @@ def _check_shapes(x: pd.Series) -> None:
             )
 
     @staticmethod
-    def _get_true_run_indices(x: npt.NDArray) -> Tuple[np.ndarray, np.ndarray]:
+    def _get_true_run_indices(x: npt.NDArray) -> Tuple[npt.NDArray, npt.NDArray]:
         """Helper function that finds consecutive runs of `True` values.
 
         Example:

kats/detectors/meta_learning/hpt_tuning.py

@@ -21,6 +21,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 from kats.consts import TimeSeriesData
 from kats.detectors.detector import DetectorModel
@@ -303,7 +304,7 @@ def get_hpt_from_ts(self, ts: TimeSeriesData) -> pd.DataFrame:
         return self.get_hpt_from_features(features_array)
 
     def get_hpt_from_features(
-        self, source_x: Union[np.ndarray, List[np.ndarray], pd.DataFrame]
+        self, source_x: Union[npt.NDArray, List[npt.NDArray], pd.DataFrame]
     ) -> pd.DataFrame:
         assert self._meta_hpt_model is not None
 

kats/detectors/outlier.py

@@ -171,8 +171,8 @@ class MultivariateAnomalyDetector(Detector):
         model_type: The type of multivariate anomaly detector (currently 'BAYESIAN_VAR' and 'VAR' options available)
     """
 
-    resid: Optional[np.ndarray] = None
-    sigma_u: Optional[np.ndarray] = None
+    resid: Optional[npt.NDArray] = None
+    sigma_u: Optional[npt.NDArray] = None
     anomaly_score_df: Optional[pd.DataFrame] = None
 
     def __init__(

kats/detectors/robust_stat_detection.py

@@ -10,6 +10,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 from kats.consts import TimeSeriesChangePoint, TimeSeriesData
 from kats.detectors.detector import Detector
@@ -48,7 +49,7 @@ def __init__(self, data: TimeSeriesData) -> None:
             )
             logging.error(msg)
             raise ValueError(msg)
-        self.zscore: Optional[np.ndarray] = None
+        self.zscore: Optional[npt.NDArray] = None
 
     # pyre-fixme[14]: `detector` overrides method defined in `Detector` inconsistently.
     def detector(