From c255a5e4d111971600c690aaab4685a65a051cfa Mon Sep 17 00:00:00 2001 From: Natasha Badami Date: Mon, 30 Sep 2024 16:20:20 -0700 Subject: [PATCH] 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 | 4 +- kats/consts.py | 6 +-- kats/detectors/bocpd.py | 38 +++++++++---------- kats/detectors/cusum_detection.py | 20 +++++----- kats/detectors/detector_consts.py | 14 +++---- kats/detectors/distribution_distance_model.py | 4 +- kats/detectors/hourly_ratio_detection.py | 2 +- kats/detectors/interval_detector.py | 10 ++--- kats/detectors/meta_learning/hpt_tuning.py | 3 +- kats/detectors/outlier.py | 4 +- kats/detectors/robust_stat_detection.py | 3 +- kats/detectors/rolling_stats_model.py | 12 +++--- kats/detectors/trend_mk.py | 4 +- kats/metrics/metrics.py | 20 +++++++--- kats/models/arima.py | 17 +++++---- kats/models/bayesian_var.py | 10 ++--- kats/models/globalmodel/data_processor.py | 12 +++--- kats/models/globalmodel/ensemble.py | 5 ++- kats/models/globalmodel/model.py | 4 +- kats/models/globalmodel/stdmodel.py | 8 ++-- kats/models/harmonic_regression.py | 8 ++-- kats/models/linear_model.py | 9 +++-- kats/models/lstm.py | 7 ++-- kats/models/metalearner/metalearner_hpt.py | 3 +- .../metalearner/metalearner_modelselect.py | 2 +- .../metalearner/metalearner_predictability.py | 2 +- kats/models/ml_ar.py | 10 ++--- kats/models/neuralprophet.py | 5 ++- kats/models/prophet.py | 6 +-- kats/models/reconciliation/base_models.py | 5 ++- kats/models/reconciliation/thm.py | 10 ++--- kats/models/sarima.py | 15 ++++---- kats/models/simple_heuristic_model.py | 8 ++-- kats/models/stlf.py | 9 +++-- kats/models/theta.py | 11 +++--- kats/models/var.py | 5 ++- kats/tests/metrics/test_metrics.py | 4 +- kats/tests/models/test_globalmodel.py | 2 +- kats/tests/utils/test_emp_confidence_int.py | 3 +- kats/tsfeatures/tsfeatures.py | 6 +-- kats/utils/backtesters.py | 11 +++--- kats/utils/cupik.py | 9 +++-- kats/utils/emp_confidence_int.py | 5 ++- kats/utils/ensemble_predict_interval.py | 4 +- kats/utils/feature_engineering.py | 2 +- 45 files changed, 194 insertions(+), 167 deletions(-) diff --git a/kats/compat/sklearn.py b/kats/compat/sklearn.py index b6ca20db..5f2fc40c 100644 --- a/kats/compat/sklearn.py +++ b/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: diff --git a/kats/consts.py b/kats/consts.py index 7dfee5f8..f53844d6 100644 --- a/kats/consts.py +++ b/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: diff --git a/kats/detectors/bocpd.py b/kats/detectors/bocpd.py index ffad1e27..6907d92f 100644 --- a/kats/detectors/bocpd.py +++ b/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,8 +1444,8 @@ 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 @@ -1453,7 +1453,7 @@ def __init__( 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, diff --git a/kats/detectors/cusum_detection.py b/kats/detectors/cusum_detection.py index 07f7b80a..9381a4e1 100644 --- a/kats/detectors/cusum_detection.py +++ b/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) diff --git a/kats/detectors/detector_consts.py b/kats/detectors/detector_consts.py index f0d53c95..6951c6e4 100644 --- a/kats/detectors/detector_consts.py +++ b/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 diff --git a/kats/detectors/distribution_distance_model.py b/kats/detectors/distribution_distance_model.py index cc67ad14..2e04aa8c 100644 --- a/kats/detectors/distribution_distance_model.py +++ b/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]) diff --git a/kats/detectors/hourly_ratio_detection.py b/kats/detectors/hourly_ratio_detection.py index d8314bba..54fd47af 100644 --- a/kats/detectors/hourly_ratio_detection.py +++ b/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: diff --git a/kats/detectors/interval_detector.py b/kats/detectors/interval_detector.py index 6e379e96..b9436a9d 100644 --- a/kats/detectors/interval_detector.py +++ b/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: diff --git a/kats/detectors/meta_learning/hpt_tuning.py b/kats/detectors/meta_learning/hpt_tuning.py index 4e5da02b..ffc235a1 100644 --- a/kats/detectors/meta_learning/hpt_tuning.py +++ b/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 diff --git a/kats/detectors/outlier.py b/kats/detectors/outlier.py index 9b353338..596f4b06 100644 --- a/kats/detectors/outlier.py +++ b/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__( diff --git a/kats/detectors/robust_stat_detection.py b/kats/detectors/robust_stat_detection.py index 89900a6f..b16ac2d4 100644 --- a/kats/detectors/robust_stat_detection.py +++ b/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( diff --git a/kats/detectors/rolling_stats_model.py b/kats/detectors/rolling_stats_model.py index 7e4040c3..48cdb8f8 100644 --- a/kats/detectors/rolling_stats_model.py +++ b/kats/detectors/rolling_stats_model.py @@ -56,7 +56,7 @@ class RollStatsFunction(str, Enum): } -def calculate_iqr(data_list: npt.NDArray, **kwargs: Any) -> Union[float, np.ndarray]: +def calculate_iqr(data_list: npt.NDArray, **kwargs: Any) -> Union[float, npt.NDArray]: """ Calculate IQR = Q3 - Q1 """ @@ -68,7 +68,7 @@ def calculate_iqr(data_list: npt.NDArray, **kwargs: Any) -> Union[float, np.ndar def calculate_z_scores( data_list: npt.NDArray, **kwargs: Any -) -> Union[float, np.ndarray]: +) -> Union[float, npt.NDArray]: """ Calculate the z-score of the last data point in data_list. Or calculate the z-score of the last data point in each row of the data_list. @@ -85,7 +85,7 @@ def calculate_z_scores( return result[0] if len(result) == 1 else result -def calculate_mad(data_list: npt.NDArray, **kwargs: Any) -> Union[float, np.ndarray]: +def calculate_mad(data_list: npt.NDArray, **kwargs: Any) -> Union[float, npt.NDArray]: """ Calculate MAD: the mean (average) distance between each data value and the mean of the data set. @@ -99,7 +99,7 @@ def calculate_mad(data_list: npt.NDArray, **kwargs: Any) -> Union[float, np.ndar def calculate_modified_z_scores_mad( data_list: npt.NDArray, **kwargs: Any -) -> Union[float, np.ndarray]: +) -> Union[float, npt.NDArray]: """ Calculate Modified z-score: (x-median)/MAD. x: the last point of data_list. @@ -118,7 +118,7 @@ def calculate_modified_z_scores_mad( def calculate_modified_z_scores_iqr( data_list: npt.NDArray, **kwargs: Any -) -> Union[float, np.ndarray]: +) -> Union[float, npt.NDArray]: """ Calculate Modified z-score (iqr version): (x-median)/IQR x: the last point of data_list. @@ -137,7 +137,7 @@ def calculate_modified_z_scores_iqr( def calculate_iqr_median_deviation( data_list: npt.NDArray, **kwargs: Any -) -> Union[float, np.ndarray]: +) -> Union[float, npt.NDArray]: """ Calculate IQR based median-deviation scores. diff --git a/kats/detectors/trend_mk.py b/kats/detectors/trend_mk.py index 8ed0edab..527c440b 100644 --- a/kats/detectors/trend_mk.py +++ b/kats/detectors/trend_mk.py @@ -221,7 +221,7 @@ def _smoothing(self, ts: pd.DataFrame) -> pd.DataFrame: return smoothed_ts - def _preprocessing(self, ts: pd.DataFrame) -> Tuple[np.ndarray, int]: + def _preprocessing(self, ts: pd.DataFrame) -> Tuple[npt.NDArray, int]: """Check and convert the dataframe ts to an numpy array. Args: @@ -249,7 +249,7 @@ def _preprocessing(self, ts: pd.DataFrame) -> Tuple[np.ndarray, int]: return x, c - def _drop_missing_values(self, x: npt.NDArray) -> Tuple[np.ndarray, int]: + def _drop_missing_values(self, x: npt.NDArray) -> Tuple[npt.NDArray, int]: """Drop the missing values in x.""" if x.ndim == 1: # univariate case with 1-dim array/ shape(n,) diff --git a/kats/metrics/metrics.py b/kats/metrics/metrics.py index fc2bf6b0..eb8b550f 100644 --- a/kats/metrics/metrics.py +++ b/kats/metrics/metrics.py @@ -106,7 +106,7 @@ def __call__( ] -def _arrays(*args: Optional[ArrayLike]) -> Generator[np.ndarray, None, None]: +def _arrays(*args: Optional[ArrayLike]) -> Generator[npt.NDArray, None, None]: """Ensure all arguments are arrays of matching size. Args: @@ -135,7 +135,7 @@ def _arrays(*args: Optional[ArrayLike]) -> Generator[np.ndarray, None, None]: def _safe_divide( num: npt.NDArray, - denom: Union[np.ndarray, float], + denom: Union[npt.NDArray, float], negative_infinity: float = -1.0, positive_infinity: float = 1.0, indeterminate: float = 0.0, @@ -182,7 +182,9 @@ def _safe_divide( def _err( - y_true: npt.NDArray, y_pred: npt.NDArray, sample_weight: Optional[np.ndarray] = None + y_true: npt.NDArray, + y_pred: npt.NDArray, + sample_weight: Optional[npt.NDArray] = None, ) -> npt.NDArray: if sample_weight is None: return y_true - y_pred @@ -208,7 +210,9 @@ def error( def _abs_err( - y_true: npt.NDArray, y_pred: npt.NDArray, sample_weight: Optional[np.ndarray] = None + y_true: npt.NDArray, + y_pred: npt.NDArray, + sample_weight: Optional[npt.NDArray] = None, ) -> npt.NDArray: err = np.abs(y_true - y_pred) if sample_weight is None: @@ -235,7 +239,9 @@ def absolute_error( def _pct_err( - y_true: npt.NDArray, y_pred: npt.NDArray, sample_weight: Optional[np.ndarray] = None + y_true: npt.NDArray, + y_pred: npt.NDArray, + sample_weight: Optional[npt.NDArray] = None, ) -> npt.NDArray: err = y_true - y_pred if sample_weight is None: @@ -262,7 +268,9 @@ def percentage_error( def _abs_pct_err( - y_true: npt.NDArray, y_pred: npt.NDArray, sample_weight: Optional[np.ndarray] = None + y_true: npt.NDArray, + y_pred: npt.NDArray, + sample_weight: Optional[npt.NDArray] = None, ) -> npt.NDArray: err = np.abs(y_true - y_pred) if sample_weight is None: diff --git a/kats/models/arima.py b/kats/models/arima.py index 433274b4..65c97821 100644 --- a/kats/models/arima.py +++ b/kats/models/arima.py @@ -25,6 +25,7 @@ from typing import Any, Callable, Dict, List, Optional import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import Params, TimeSeriesData from kats.models.model import Model @@ -51,7 +52,7 @@ class ARIMAParams(Params): p: int d: int q: int - exog: Optional[np.ndarray] = None + exog: Optional[npt.NDArray] = None dates: Optional[pd.DatetimeIndex] = None freq: Optional[str] = None @@ -88,10 +89,10 @@ class ARIMAModel(Model[ARIMAParams]): model: Optional[ARIMAResults] = None include_history: bool = False fcst_df: Optional[pd.DataFrame] = None - y_fcst: Optional[np.ndarray] = None - y_fcst_lower: Optional[np.ndarray] = None - y_fcst_upper: Optional[np.ndarray] = None - start_params: Optional[np.ndarray] = None + y_fcst: Optional[npt.NDArray] = None + y_fcst_lower: Optional[npt.NDArray] = None + y_fcst_upper: Optional[npt.NDArray] = None + start_params: Optional[npt.NDArray] = None transparams: bool = False method: Optional[str] = None trend: Optional[str] = None @@ -99,7 +100,7 @@ class ARIMAModel(Model[ARIMAParams]): maxiter: Optional[int] = None full_output: bool = False disp: Optional[int] = None - callback: Optional[Callable[[np.ndarray], None]] = None + callback: Optional[Callable[[npt.NDArray], None]] = None start_ar_lags: Optional[int] = None dates: Optional[pd.DatetimeIndex] = None @@ -116,7 +117,7 @@ def __init__(self, data: TimeSeriesData, params: ARIMAParams) -> None: def fit( self, - start_params: Optional[np.ndarray] = None, + start_params: Optional[npt.NDArray] = None, transparams: bool = True, method: str = "css-mle", trend: str = "c", @@ -124,7 +125,7 @@ def fit( maxiter: int = 500, full_output: bool = True, disp: int = 5, - callback: Optional[Callable[[np.ndarray], None]] = None, + callback: Optional[Callable[[npt.NDArray], None]] = None, start_ar_lags: Optional[int] = None, **kwargs: Any, ) -> None: diff --git a/kats/models/bayesian_var.py b/kats/models/bayesian_var.py index 584fff7e..c0900861 100644 --- a/kats/models/bayesian_var.py +++ b/kats/models/bayesian_var.py @@ -81,9 +81,9 @@ class BayesianVAR(m.Model[BayesianVARParams]): params: the parameter class defined with `BayesianVARParams` """ - sigma_ols: Optional[np.ndarray] = None - v_posterior: Optional[np.ndarray] = None - mu_posterior: Optional[np.ndarray] = None + sigma_ols: Optional[npt.NDArray] = None + v_posterior: Optional[npt.NDArray] = None + mu_posterior: Optional[npt.NDArray] = None resid: Optional[pd.DataFrame] = None forecast: Optional[Dict[str, TimeSeriesData]] = None forecast_max_time: Optional[datetime] = None @@ -102,7 +102,7 @@ class BayesianVAR(m.Model[BayesianVARParams]): N: int num_mu_coefficients: int fitted: bool = False - forecast_vals: Optional[List[np.ndarray]] = None + forecast_vals: Optional[List[npt.NDArray]] = None def __init__(self, data: TimeSeriesData, params: BayesianVARParams) -> None: if data.is_univariate(): @@ -150,7 +150,7 @@ def _check_get_freq(data: TimeSeriesData) -> str: @staticmethod def _convert_timeseries_np( timeseries: TimeSeriesData, - ) -> Tuple[np.ndarray, np.ndarray]: + ) -> Tuple[npt.NDArray, npt.NDArray]: data_df = timeseries.to_dataframe() Y = data_df.drop(columns=[timeseries.time_col_name]).to_numpy().T X = np.expand_dims(pd.RangeIndex(0, len(timeseries)), axis=0) diff --git a/kats/models/globalmodel/data_processor.py b/kats/models/globalmodel/data_processor.py index 0a488d7d..6e6e28ff 100644 --- a/kats/models/globalmodel/data_processor.py +++ b/kats/models/globalmodel/data_processor.py @@ -80,7 +80,7 @@ def _valid_dataset( self, # pyre-fixme[2]: Parameter annotation cannot contain `Any`. dataset: Union[Dict[Any, TimeSeriesData], List[TimeSeriesData]], - ) -> Tuple[np.ndarray, np.ndarray]: + ) -> Tuple[npt.NDArray, npt.NDArray]: if len(dataset) < 1: msg = "Input dataset should be non-empty." logging.error(msg) @@ -446,11 +446,11 @@ def _get_array( reduced_length: int, reduced_valid_length: int, ) -> Tuple[ - np.ndarray, - np.ndarray, - np.ndarray, - Optional[np.ndarray], - Optional[np.ndarray], + npt.NDArray, + npt.NDArray, + npt.NDArray, + Optional[npt.NDArray], + Optional[npt.NDArray], ]: """ diff --git a/kats/models/globalmodel/ensemble.py b/kats/models/globalmodel/ensemble.py index eade7de8..400b7e6a 100644 --- a/kats/models/globalmodel/ensemble.py +++ b/kats/models/globalmodel/ensemble.py @@ -13,6 +13,7 @@ import joblib import numpy as np +import numpy.typing as npt import pandas as pd import torch from kats.consts import TimeSeriesData @@ -257,7 +258,7 @@ def _combine_fcst( self, # pyre-fixme[2]: Parameter annotation cannot be `Any`. idx: Any, - fcsts: List[np.ndarray], + fcsts: List[npt.NDArray], steps: int, raw: bool, # pyre-fixme[11]: Annotation `Timestamp` is not defined as a type. @@ -295,7 +296,7 @@ def predict( steps: int, test_batch_size: int = 500, raw: bool = False, - ) -> Dict[Any, Union[pd.DataFrame, List[np.ndarray]]]: + ) -> Dict[Any, Union[pd.DataFrame, List[npt.NDArray]]]: """Generate forecasts for the target time series. Args: diff --git a/kats/models/globalmodel/model.py b/kats/models/globalmodel/model.py index e1f49642..d494b6a3 100644 --- a/kats/models/globalmodel/model.py +++ b/kats/models/globalmodel/model.py @@ -511,7 +511,7 @@ def _format_fcst( # pyre-fixme[2]: Parameter annotation cannot contain `Any`. ids: List[Any], # pyre-fixme[2]: Parameter annotation cannot contain `Any`. - fcst_store: Dict[Any, List[np.ndarray]], + fcst_store: Dict[Any, List[npt.NDArray]], steps: int, first_time: npt.NDArray, ) -> Dict[Any, pd.DataFrame]: @@ -568,7 +568,7 @@ def predict( steps: int, test_batch_size: int = 500, raw: bool = False, - ) -> Dict[Any, Union[pd.DataFrame, List[np.ndarray]]]: + ) -> Dict[Any, Union[pd.DataFrame, List[npt.NDArray]]]: """Generate forecasts for target time series. Args: diff --git a/kats/models/globalmodel/stdmodel.py b/kats/models/globalmodel/stdmodel.py index 6d003d93..88594bea 100644 --- a/kats/models/globalmodel/stdmodel.py +++ b/kats/models/globalmodel/stdmodel.py @@ -136,7 +136,7 @@ def _get_period(self, period: Optional[int]) -> Optional[int]: def _decompose( self, ts: TimeSeriesData - ) -> Tuple[Union[ProphetModel, np.ndarray], Dict[str, TimeSeriesData]]: + ) -> Tuple[Union[ProphetModel, npt.NDArray], Dict[str, TimeSeriesData]]: """Decompose time series into trend, seasonality.""" if self.decomposition_model == "prophet": if self.decomposition_params is None: @@ -182,7 +182,7 @@ def _decompose( def _deseasonal( self, ts: TimeSeriesData - ) -> Tuple[Union[ProphetModel, np.ndarray], TimeSeriesData]: + ) -> Tuple[Union[ProphetModel, npt.NDArray], TimeSeriesData]: tsd_model, tsd_res = self._decompose(ts) if self.fit_trend: return tsd_model, tsd_res["trend"] @@ -194,7 +194,7 @@ def _deseasonal( return tsd_model, new_ts def _predict_seasonality( - self, steps: int, tsd_model: Union[ProphetModel, np.ndarray] + self, steps: int, tsd_model: Union[ProphetModel, npt.NDArray] ) -> npt.NDArray: """Predict the future seasonality. @@ -226,7 +226,7 @@ def _predict_seasonality( def _prepare_ts( self, tag: Union[str, int], ts: TimeSeriesData, steps: int - ) -> Tuple[Union[str, int], TimeSeriesData, np.ndarray]: + ) -> Tuple[Union[str, int], TimeSeriesData, npt.NDArray]: """Prepare time series into seasonality and non-seasonal part.""" tsd_model, new_ts = self._deseasonal(ts) new_seasonal = self._predict_seasonality(steps, tsd_model) diff --git a/kats/models/harmonic_regression.py b/kats/models/harmonic_regression.py index 9735be5c..adab4eeb 100644 --- a/kats/models/harmonic_regression.py +++ b/kats/models/harmonic_regression.py @@ -41,8 +41,8 @@ def __post_init__(self) -> None: class HarmonicRegressionModel(Model[Optional[np.ndarray]]): - params: Optional[np.ndarray] = None - harms: Optional[np.ndarray] = None + params: Optional[npt.NDArray] = None + harms: Optional[npt.NDArray] = None def __init__(self, data: TimeSeriesData, params: HarmonicRegressionParams) -> None: super().__init__(data, params) @@ -163,7 +163,7 @@ def fourier_series( @staticmethod def make_harm_eval( harmonics: npt.NDArray, - ) -> Callable[..., np.ndarray]: + ) -> Callable[..., npt.NDArray]: """Defines evaluation function for the optimizer Parameters ---------- @@ -182,7 +182,7 @@ def harm_eval(step: npt.NDArray, *params: Real) -> npt.NDArray: def fit_harmonics( self, period: float, fourier_order: int - ) -> Tuple[np.ndarray, np.ndarray]: + ) -> Tuple[npt.NDArray, npt.NDArray]: """Performs harmonic regression. Harmonic regression fits cosines amplitude*cos(freq*t + phase). Using double angle identity formulas, diff --git a/kats/models/linear_model.py b/kats/models/linear_model.py index 8e5cf31b..5e63c809 100644 --- a/kats/models/linear_model.py +++ b/kats/models/linear_model.py @@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Union import numpy as np +import numpy.typing as npt import pandas as pd import statsmodels.api as sm from kats.consts import Params, TimeSeriesData @@ -74,10 +75,10 @@ def __init__(self, data: TimeSeriesData, params: LinearModelParams) -> None: self._X_future: Optional[List[int]] = None self.past_length: int = len(data.time) self.dates: Optional[pd.DatetimeIndex] = None - self.y_fcst: Optional[Union[pd.Series, np.ndarray]] = None - self.sdev: Optional[Union[np.ndarray, float]] = None - self.y_fcst_lower: Optional[Union[pd.Series, np.ndarray, float]] = None - self.y_fcst_upper: Optional[Union[pd.Series, np.ndarray, float]] = None + self.y_fcst: Optional[Union[pd.Series, npt.NDArray]] = None + self.sdev: Optional[Union[npt.NDArray, float]] = None + self.y_fcst_lower: Optional[Union[pd.Series, npt.NDArray, float]] = None + self.y_fcst_upper: Optional[Union[pd.Series, npt.NDArray, float]] = None def fit(self) -> None: """fit Linear Model.""" diff --git a/kats/models/lstm.py b/kats/models/lstm.py index 115a58c7..5534b872 100644 --- a/kats/models/lstm.py +++ b/kats/models/lstm.py @@ -22,6 +22,7 @@ from typing import Any, Dict, List, Optional, Tuple import numpy as np +import numpy.typing as npt import pandas as pd import torch import torch.nn as nn @@ -120,9 +121,9 @@ class LSTMModel(Model[LSTMParams]): model: Optional[LSTMForecast] = None freq: Optional[str] = None dates: Optional[pd.DatetimeIndex] = None - y_fcst: Optional[np.ndarray] = None - y_fcst_lower: Optional[np.ndarray] = None - y_fcst_upper: Optional[np.ndarray] = None + y_fcst: Optional[npt.NDArray] = None + y_fcst_lower: Optional[npt.NDArray] = None + y_fcst_upper: Optional[npt.NDArray] = None fcst_df: Optional[pd.DataFrame] = None def __init__(self, data: TimeSeriesData, params: LSTMParams) -> None: diff --git a/kats/models/metalearner/metalearner_hpt.py b/kats/models/metalearner/metalearner_hpt.py index 3d62a5a5..a969759b 100644 --- a/kats/models/metalearner/metalearner_hpt.py +++ b/kats/models/metalearner/metalearner_hpt.py @@ -20,6 +20,7 @@ import joblib import matplotlib.pyplot as plt import numpy as np +import numpy.typing as npt import pandas as pd import torch import torch.nn as nn @@ -692,7 +693,7 @@ def pred( return res def pred_by_feature( - self, source_x: Union[np.ndarray, List[np.ndarray], pd.DataFrame] + self, source_x: Union[npt.NDArray, List[npt.NDArray], pd.DataFrame] ) -> List[Dict[str, Any]]: """Predict hyper-parameters for time series features. diff --git a/kats/models/metalearner/metalearner_modelselect.py b/kats/models/metalearner/metalearner_modelselect.py index 6ba5172a..9e83a7a4 100644 --- a/kats/models/metalearner/metalearner_modelselect.py +++ b/kats/models/metalearner/metalearner_modelselect.py @@ -427,7 +427,7 @@ def pred( def pred_by_feature( self, - source_x: Union[np.ndarray, List[np.ndarray], pd.DataFrame], + source_x: Union[npt.NDArray, List[npt.NDArray], pd.DataFrame], n_top: int = 1, ) -> npt.NDArray: """Predict the best forecasting models given a list/dataframe of time series features diff --git a/kats/models/metalearner/metalearner_predictability.py b/kats/models/metalearner/metalearner_predictability.py index 95c86bcf..52e9da2d 100644 --- a/kats/models/metalearner/metalearner_predictability.py +++ b/kats/models/metalearner/metalearner_predictability.py @@ -303,7 +303,7 @@ def pred( return ans def pred_by_feature( - self, source_x: Union[np.ndarray, List[np.ndarray], pd.DataFrame] + self, source_x: Union[npt.NDArray, List[npt.NDArray], pd.DataFrame] ) -> npt.NDArray: """Predict whether a list of time series are predicable or not given their time series features. Args: diff --git a/kats/models/ml_ar.py b/kats/models/ml_ar.py index b9c7a8ae..09449620 100644 --- a/kats/models/ml_ar.py +++ b/kats/models/ml_ar.py @@ -553,7 +553,7 @@ def _normalize_data( def _embed_regressor( self, data: pd.DataFrame - ) -> Tuple[List[np.ndarray], List[str]]: + ) -> Tuple[List[npt.NDArray], List[str]]: """ Embedding external regressors. @@ -598,7 +598,7 @@ def _embed_regressor( def _embed_categorical( self, data: pd.DataFrame, cat_labels: Dict[str, pd.CategoricalDtype] - ) -> Tuple[np.ndarray, List[str]]: + ) -> Tuple[npt.NDArray, List[str]]: cat_encoded_data = pd.DataFrame() if self.params.categoricals: @@ -626,7 +626,7 @@ def _embed_and_gen_past_features_single_series( lags: int, target_vars: Union[List[str], str], cat_labels: Dict[str, pd.CategoricalDtype], - ) -> Tuple[Dict[str, np.ndarray], Dict[str, pd.DataFrame], List[str]]: + ) -> Tuple[Dict[str, npt.NDArray], Dict[str, pd.DataFrame], List[str]]: data = series_data.loc[: series_data[target_vars].last_valid_index()] @@ -874,13 +874,13 @@ def _embed_future_cov( def _merge_past_and_future_reg( self, - norm_in_data: Dict[str, np.ndarray], + norm_in_data: Dict[str, npt.NDArray], norm_out_data: Dict[str, pd.DataFrame], horizons: List[int], cal_feat: pd.DataFrame, emb_fut_cov: pd.DataFrame, gen_meta_data: bool = True, - ) -> Tuple[np.ndarray, pd.DataFrame, List[str]]: + ) -> Tuple[npt.NDArray, pd.DataFrame, List[str]]: offset = pd.tseries.frequencies.to_offset(self.params.freq) num_cols = self.num_hist_reg + cal_feat.shape[1] + emb_fut_cov.shape[1] + 1 diff --git a/kats/models/neuralprophet.py b/kats/models/neuralprophet.py index 91f77f42..904dd6b2 100644 --- a/kats/models/neuralprophet.py +++ b/kats/models/neuralprophet.py @@ -17,6 +17,7 @@ import matplotlib.pyplot as plt import numpy as np +import numpy.typing as npt import pandas as pd import torch @@ -128,7 +129,7 @@ class NeuralProphetParams(Params): Each regressor is a dictionary with required key "names"and optional keys "regularization" and "normalize". """ - changepoints: Optional[Union[List[str], List[np.datetime64], np.ndarray]] + changepoints: Optional[Union[List[str], List[np.datetime64], npt.NDArray]] n_changepoints: int changepoints_range: float trend_reg: float @@ -164,7 +165,7 @@ def __init__( # import np.typing as npt # replace 'np.ndarray' by npt.NDArray['np.datetime64'] changepoints: Optional[ - Union[List[str], List[np.datetime64], np.ndarray] + Union[List[str], List[np.datetime64], npt.NDArray] ] = None, n_changepoints: int = 10, changepoints_range: float = 0.9, diff --git a/kats/models/prophet.py b/kats/models/prophet.py index b39ea917..d88867e7 100644 --- a/kats/models/prophet.py +++ b/kats/models/prophet.py @@ -721,7 +721,7 @@ def sample_model_vectorized( s_a: npt.NDArray, s_m: npt.NDArray, n_samples: int, -) -> Dict[str, np.ndarray]: +) -> Dict[str, npt.NDArray]: """Simulate observations from the extrapolated generative model. Vectorized version of sample_model(). Returns: @@ -746,7 +746,7 @@ def sample_model_vectorized( def sample_posterior_predictive( prophet_model: Prophet, df: pd.DataFrame, vectorized: bool -) -> Dict[str, np.ndarray]: +) -> Dict[str, npt.NDArray]: """Generate posterior samples of a trained Prophet model. Args: @@ -805,7 +805,7 @@ def _make_historical_mat_time( changepoints_t: npt.NDArray, n_row: int, single_diff: float, -) -> Tuple[np.ndarray, np.ndarray]: +) -> Tuple[npt.NDArray, npt.NDArray]: """ Creates a matrix of slope-deltas where these changes occured in training data according to the trained prophet obj """ diff --git a/kats/models/reconciliation/base_models.py b/kats/models/reconciliation/base_models.py index 4912fe2b..b2ff467b 100644 --- a/kats/models/reconciliation/base_models.py +++ b/kats/models/reconciliation/base_models.py @@ -36,6 +36,7 @@ import logging import numpy as np +import numpy.typing as npt import pandas as pd @@ -58,8 +59,8 @@ def __init__( level: int, model_name: Optional[str] = None, model_params: Optional[object] = None, - residuals: Optional[np.ndarray] = None, - fcsts: Optional[np.ndarray] = None, + residuals: Optional[npt.NDArray] = None, + fcsts: Optional[npt.NDArray] = None, ) -> None: if not isinstance(level, int) or level < 1: diff --git a/kats/models/reconciliation/thm.py b/kats/models/reconciliation/thm.py index 6c8b08f3..1a21baa0 100644 --- a/kats/models/reconciliation/thm.py +++ b/kats/models/reconciliation/thm.py @@ -64,8 +64,8 @@ class TemporalHierarchicalModel: # pyre-fixme[24]: Generic type `Model` expects 1 type parameter. models: Optional[Dict[str, Model]] = None - residuals: Optional[Dict[int, np.ndarray]] = None - res_matrix: Optional[np.ndarray] = None + residuals: Optional[Dict[int, npt.NDArray]] = None + res_matrix: Optional[npt.NDArray] = None def __init__(self, data: TimeSeriesData, baseModels: List[BaseTHModel]) -> None: @@ -198,7 +198,7 @@ def _get_residuals(self, model: Model) -> npt.NDArray: return merge[col].values - merge["fcst"].values raise ValueError("Couldn't find residual or forecast values in model") - def _get_all_residuals(self) -> Dict[int, np.ndarray]: + def _get_all_residuals(self) -> Dict[int, npt.NDArray]: """ Calculate residuals for all base models. @@ -308,7 +308,7 @@ def get_W(self, method: str = "struc", eps: float = 1e-5) -> npt.NDArray: raise _log_error(f"{method} is invalid for get_W() method.") # pyre-fixme[2]: Parameter must be annotated. - def _predict_origin(self, steps: int, method="struc") -> Dict[int, np.ndarray]: + def _predict_origin(self, steps: int, method="struc") -> Dict[int, npt.NDArray]: """ Generate original forecasts from each base model (without time index). @@ -370,7 +370,7 @@ def _predict( method="struc", origin_fcst: bool = False, fcst_levels: Optional[List[int]] = None, - ) -> Dict[str, Dict[int, np.ndarray]]: + ) -> Dict[str, Dict[int, npt.NDArray]]: """Generate forecasts for each level (without time index). Args: diff --git a/kats/models/sarima.py b/kats/models/sarima.py index 827512fb..f8f17a6e 100644 --- a/kats/models/sarima.py +++ b/kats/models/sarima.py @@ -9,6 +9,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import Params, TimeSeriesData from kats.models.model import Model @@ -128,7 +129,7 @@ class SARIMAModel(Model[SARIMAParams]): params: :class:`SARIMAParams` for model parameters. """ - start_params: Optional[np.ndarray] = None + start_params: Optional[npt.NDArray] = None transformed: Optional[bool] = None includes_fixed: Optional[bool] = None cov_type: Optional[str] = None @@ -137,7 +138,7 @@ class SARIMAModel(Model[SARIMAParams]): maxiter: Optional[int] = None full_output: Optional[bool] = None disp: Optional[bool] = None - callback: Optional[Callable[[np.ndarray], None]] = None + callback: Optional[Callable[[npt.NDArray], None]] = None return_params: Optional[bool] = None optim_score: Optional[str] = None optim_complex_step: Optional[bool] = None @@ -148,9 +149,9 @@ class SARIMAModel(Model[SARIMAParams]): alpha: float = 0.05 fcst_df: Optional[pd.DataFrame] = None freq: Optional[float] = None - y_fcst: Optional[np.ndarray] = None - y_fcst_lower: Optional[np.ndarray] = None - y_fcst_upper: Optional[np.ndarray] = None + y_fcst: Optional[npt.NDArray] = None + y_fcst_lower: Optional[npt.NDArray] = None + y_fcst_upper: Optional[npt.NDArray] = None dates: Optional[pd.DatetimeIndex] = None def __init__( @@ -167,7 +168,7 @@ def __init__( def fit( self, - start_params: Optional[np.ndarray] = None, + start_params: Optional[npt.NDArray] = None, transformed: bool = True, includes_fixed: bool = False, cov_type: Optional[str] = None, @@ -176,7 +177,7 @@ def fit( maxiter: int = 50, full_output: bool = True, disp: bool = False, - callback: Optional[Callable[[np.ndarray], None]] = None, + callback: Optional[Callable[[npt.NDArray], None]] = None, return_params: bool = False, optim_score: Optional[str] = None, optim_complex_step: bool = True, diff --git a/kats/models/simple_heuristic_model.py b/kats/models/simple_heuristic_model.py index 4d3c40c1..125ccfbe 100644 --- a/kats/models/simple_heuristic_model.py +++ b/kats/models/simple_heuristic_model.py @@ -71,12 +71,12 @@ class SimpleHeuristicModel(Model[SimpleHeuristicModelParams]): params: the parameter class defined with `SimpleHeuristicModelParams` """ - model: Callable[[np.ndarray], np.ndarray] + model: Callable[[npt.NDArray], npt.NDArray] include_history: bool = False dates: Optional[pd.DatetimeIndex] = None - y_fcst: Optional[np.ndarray] = None - y_fcst_lower: Optional[np.ndarray] = None - y_fcst_upper: Optional[np.ndarray] = None + y_fcst: Optional[npt.NDArray] = None + y_fcst_lower: Optional[npt.NDArray] = None + y_fcst_upper: Optional[npt.NDArray] = None fcst_df: pd.DataFrame = pd.DataFrame(data=None) freq: Optional[str] = None diff --git a/kats/models/stlf.py b/kats/models/stlf.py index d44a4beb..d65cd254 100644 --- a/kats/models/stlf.py +++ b/kats/models/stlf.py @@ -29,6 +29,7 @@ from typing import Any, Callable, cast, Dict, List, Optional, Union import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import Params, TimeSeriesData from kats.models import ( @@ -168,9 +169,9 @@ class STLFModel(Model[STLFParams]): ] = None freq: Optional[str] = None alpha: Optional[float] = None - y_fcst: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None - fcst_lower: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None - fcst_upper: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None + y_fcst: Optional[Union[npt.NDArray, pd.Series, pd.DataFrame]] = None + fcst_lower: Optional[Union[npt.NDArray, pd.Series, pd.DataFrame]] = None + fcst_upper: Optional[Union[npt.NDArray, pd.Series, pd.DataFrame]] = None dates: Optional[pd.DatetimeIndex] = None fcst_df: Optional[pd.DataFrame] = None deseasonal_operator: Callable(Union[_operator.truediv, _operator.sub])[ @@ -185,7 +186,7 @@ class STLFModel(Model[STLFParams]): Union[pd.Series, pd.DataFrame], Union[pd.Series, pd.DataFrame], ], - Union[np.ndarray, pd.Series, pd.DataFrame], + Union[npt.NDArray, pd.Series, pd.DataFrame], ] def __init__(self, data: TimeSeriesData, params: STLFParams) -> None: diff --git a/kats/models/theta.py b/kats/models/theta.py index 3ade2844..27e1c7b0 100644 --- a/kats/models/theta.py +++ b/kats/models/theta.py @@ -20,6 +20,7 @@ from typing import Any, Dict, List, Optional import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import Params, TimeSeriesData from kats.models.model import Model @@ -66,13 +67,13 @@ class ThetaModel(Model[ThetaParams]): decomp: Optional[Dict[str, TimeSeriesData]] = None ses_model: Optional[HoltWintersResults] = None drift: Optional[float] = None - fitted_values: Optional[np.ndarray] = None - residuals: Optional[np.ndarray] = None + fitted_values: Optional[npt.NDArray] = None + residuals: Optional[npt.NDArray] = None fcst_df: Optional[pd.DataFrame] = None dates: Optional[pd.DatetimeIndex] = None - y_fcst: Optional[np.ndarray] = None - y_fcst_lower: Optional[np.ndarray] = None - y_fcst_upper: Optional[np.ndarray] = None + y_fcst: Optional[npt.NDArray] = None + y_fcst_lower: Optional[npt.NDArray] = None + y_fcst_upper: Optional[npt.NDArray] = None freq: Optional[str] = None alpha: Optional[float] = None include_history: bool = False diff --git a/kats/models/var.py b/kats/models/var.py index 85291cb0..aac4831c 100644 --- a/kats/models/var.py +++ b/kats/models/var.py @@ -25,6 +25,7 @@ from typing import Any, Dict, List, Optional, Tuple import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import Params, TimeSeriesData from kats.models.model import Model @@ -81,8 +82,8 @@ class VARModel(Model[VARParams]): model: Optional[VARResults] = None k_ar: Optional[int] = None - sigma_u: Optional[np.ndarray] = None - resid: Optional[np.ndarray] = None + sigma_u: Optional[npt.NDArray] = None + resid: Optional[npt.NDArray] = None freq: Optional[str] = None alpha: Optional[float] = None dates: Optional[pd.DatetimeIndex] = None diff --git a/kats/tests/metrics/test_metrics.py b/kats/tests/metrics/test_metrics.py index bc451f39..8b81d247 100644 --- a/kats/tests/metrics/test_metrics.py +++ b/kats/tests/metrics/test_metrics.py @@ -19,7 +19,9 @@ class MetricsTest(TestCase): def validate( - self, expected: Union[float, np.ndarray], result: Union[float, np.ndarray] + self, + expected: Union[float, np_typing.NDArray], + result: Union[float, np_typing.NDArray], ) -> None: if isinstance(expected, float): self.assertTrue(isinstance(result, float)) diff --git a/kats/tests/models/test_globalmodel.py b/kats/tests/models/test_globalmodel.py index 42b08aed..668f5c97 100644 --- a/kats/tests/models/test_globalmodel.py +++ b/kats/tests/models/test_globalmodel.py @@ -103,7 +103,7 @@ def _gm_mock_predict_func( len_quantile: int, raw: bool = True, test_batch_size: int = 500, -) -> Dict[int, List[np.ndarray]]: +) -> Dict[int, List[npt.NDArray]]: """ Helper function for building predict method for mock GMModel. """ diff --git a/kats/tests/utils/test_emp_confidence_int.py b/kats/tests/utils/test_emp_confidence_int.py index 04d6becb..4cdfcb9a 100644 --- a/kats/tests/utils/test_emp_confidence_int.py +++ b/kats/tests/utils/test_emp_confidence_int.py @@ -13,6 +13,7 @@ import kats.utils.emp_confidence_int # noqa import matplotlib.pyplot as plt import numpy as np +import numpy.typing as npt import pandas as pd from kats.compat.pandas import assert_frame_equal from kats.consts import Params, TimeSeriesData @@ -119,7 +120,7 @@ def get_name(funcname: str, num: int, kwargs: Dict[str, Any]) -> str: # backtester.run_backtest() # _RAW_ERRORS = backtester.raw_errors # fmt: off -_RAW_ERRORS: List[np.ndarray] = [ +_RAW_ERRORS: List[npt.NDArray] = [ np.array([ 1.33793031e01, 1.56117160e00, -3.43035881e-02, 1.59373281e01, 1.00730540e01, 3.66296231e01, 4.33868022e01, 3.67005869e01, diff --git a/kats/tsfeatures/tsfeatures.py b/kats/tsfeatures/tsfeatures.py index 11f86e37..4af7191a 100644 --- a/kats/tsfeatures/tsfeatures.py +++ b/kats/tsfeatures/tsfeatures.py @@ -698,7 +698,7 @@ def get_stability(x: npt.NDArray, window_size: int = 20) -> float: # @jit(forceobj=True) def get_statistics( x: npt.NDArray, - dict_features: Optional[Dict[str, Callable[[np.ndarray], float]]] = None, + dict_features: Optional[Dict[str, Callable[[npt.NDArray], float]]] = None, extra_args: Optional[Dict[str, bool]] = None, default_status: bool = True, ) -> Dict[str, float]: @@ -2139,7 +2139,7 @@ def get_features( self, data: Union[TimeSeriesData, pd.Series], raw: bool = False ) -> Union[ pd.DataFrame, - np.ndarray, + npt.NDArray, ]: if isinstance(data, TimeSeriesData): @@ -2211,7 +2211,7 @@ def _get_columns( def get_features( self, data: Union[TimeSeriesData, pd.Series], raw: bool = False - ) -> Union[np.ndarray, pd.DataFrame]: + ) -> Union[npt.NDArray, pd.DataFrame]: if isinstance(data, TimeSeriesData): data = np.array(data.time.astype("int") // 10**9) diff --git a/kats/utils/backtesters.py b/kats/utils/backtesters.py index 0e87be0d..576b8656 100644 --- a/kats/utils/backtesters.py +++ b/kats/utils/backtesters.py @@ -49,6 +49,7 @@ import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import _log_error, Params, TimeSeriesData from kats.metrics.metrics import core_metric, CoreMetric @@ -401,11 +402,11 @@ class BackTesterParent(ABC): params: Params multi: bool offset: int - results: List[Tuple[np.ndarray, np.ndarray, "Model[Params]", np.ndarray]] + results: List[Tuple[npt.NDArray, npt.NDArray, "Model[Params]", npt.NDArray]] errors: Dict[str, float] size: int freq: Optional[str] - raw_errors: List[np.ndarray] + raw_errors: List[npt.NDArray] def __init__( self, @@ -518,7 +519,7 @@ def _create_model( self, training_data_indices: Tuple[int, int], testing_data_indices: Tuple[int, int], - ) -> Optional[Tuple[np.ndarray, np.ndarray, "Model[Params]", np.ndarray]]: + ) -> Optional[Tuple[npt.NDArray, npt.NDArray, "Model[Params]", npt.NDArray]]: """ Trains model, evaluates it, and stores results in results list. """ @@ -1227,10 +1228,10 @@ class CrossValidation: """ size: int - results: List[Tuple[np.ndarray, np.ndarray, "Model[Params]", np.ndarray]] + results: List[Tuple[npt.NDArray, npt.NDArray, "Model[Params]", npt.NDArray]] num_folds: int errors: Dict[str, float] - raw_errors: List[np.ndarray] + raw_errors: List[npt.NDArray] _backtester: BackTesterParent def __init__( diff --git a/kats/utils/cupik.py b/kats/utils/cupik.py index 4453a319..d23dc29e 100644 --- a/kats/utils/cupik.py +++ b/kats/utils/cupik.py @@ -15,6 +15,7 @@ from typing_extensions import Protocol import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import TimeSeriesData @@ -29,8 +30,8 @@ def transform(self, data: TimeSeriesData) -> object: ... def fit( self, - x: Union[pd.DataFrame, np.ndarray], - y: Optional[Union[pd.Series, np.ndarray]], + x: Union[pd.DataFrame, npt.NDArray], + y: Optional[Union[pd.Series, npt.NDArray]], **kwargs: Any, ) -> List[TimeSeriesData]: ... @@ -54,7 +55,7 @@ class Pipeline: remove: bool = False useFeatures: bool = False extra_fitting_params: Optional[Dict[str, Any]] = None - y: Optional[Union[np.ndarray, pd.Series]] = None + y: Optional[Union[npt.NDArray, pd.Series]] = None def __init__(self, steps: List[PipelineStep]) -> None: """ @@ -201,7 +202,7 @@ def _fit_sklearn_( self, step: Step, data: List[TimeSeriesData], - y: Optional[Union[pd.Series, np.ndarray]], + y: Optional[Union[pd.Series, npt.NDArray]], ) -> List[TimeSeriesData]: """ Internal function for fitting sklearn model on a tabular data with features diff --git a/kats/utils/emp_confidence_int.py b/kats/utils/emp_confidence_int.py index 2cebd6fa..e31a217a 100644 --- a/kats/utils/emp_confidence_int.py +++ b/kats/utils/emp_confidence_int.py @@ -18,6 +18,7 @@ import matplotlib.pyplot as plt import numpy as np +import numpy.typing as npt import pandas as pd from kats.consts import Params, TimeSeriesData from kats.utils.backtesters import BackTesterRollingWindow @@ -53,8 +54,8 @@ class EmpConfidenceInt: freq: Optional[str] = None dates: Optional[pd.DatetimeIndex] = None - predicted: Optional[np.ndarray] = None - coefs: Optional[np.ndarray] = None + predicted: Optional[npt.NDArray] = None + coefs: Optional[npt.NDArray] = None df: Optional[pd.DataFrame] = None SE: Optional[pd.DataFrame] = None diff --git a/kats/utils/ensemble_predict_interval.py b/kats/utils/ensemble_predict_interval.py index ed5a1eb8..a2aa3bb5 100644 --- a/kats/utils/ensemble_predict_interval.py +++ b/kats/utils/ensemble_predict_interval.py @@ -168,7 +168,7 @@ def __init__( self.multiprocessing: bool = multiprocessing - def _get_error_matrix_single(self, idx: int) -> Tuple[int, np.ndarray]: + def _get_error_matrix_single(self, idx: int) -> Tuple[int, npt.NDArray]: """ Calculate error vector for each block """ @@ -225,7 +225,7 @@ def _get_error_matrix(self) -> None: def _projection_single( self, idx: int, step: int = 30, rolling_based: bool = False - ) -> Tuple[int, np.ndarray]: + ) -> Tuple[int, npt.NDArray]: """ Get forecasting for future steps for one chain. """ diff --git a/kats/utils/feature_engineering.py b/kats/utils/feature_engineering.py index 340dada1..4bd07aa7 100644 --- a/kats/utils/feature_engineering.py +++ b/kats/utils/feature_engineering.py @@ -61,7 +61,7 @@ } -def _map(index: Union[np.ndarray, pd.Int64Index], map: Dict[int, int]) -> npt.NDArray: +def _map(index: Union[npt.NDArray, pd.Int64Index], map: Dict[int, int]) -> npt.NDArray: """Map values to other values efficiently. Args: