Added Type Hints for dilated shapelets #1910

aeon/transformations/collection/shapelet_based/_dilated_shapelet_transform.py

@@ -8,9 +8,11 @@
 __all__ = ["RandomDilatedShapeletTransform"]
 
 import warnings
+from typing import Dict, List, Optional, Union
 
 import numpy as np
 from numba import njit, prange, set_num_threads
+from numba.core.registry import CPUDispatcher
 from numba.typed import List
 from sklearn.preprocessing import LabelEncoder
 
@@ -148,16 +150,15 @@ class RandomDilatedShapeletTransform(BaseCollectionTransformer):
     }
 
     def __init__(
-        self,
-        max_shapelets=10_000,
-        shapelet_lengths=None,
-        proba_normalization=0.8,
-        threshold_percentiles=None,
-        alpha_similarity=0.5,
-        use_prime_dilations=False,
-        random_state=None,
-        distance="manhattan",
-        n_jobs=1,
+        max_shapelets: int = 10_000,
+        shapelet_lengths: Optional[Union[List[int], np.ndarray]] = None,
+        proba_normalization: float = 0.8,
+        threshold_percentiles: Optional[Union[List[float], np.ndarray]] = None,
+        alpha_similarity: float = 0.5,
+        use_prime_dilations: bool = False,
+        random_state: Optional[int] = None,
+        distance: CPUDispatcher = "manhattan",
+        n_jobs: int = 1,
     ):
         self.max_shapelets = max_shapelets
         self.shapelet_lengths = shapelet_lengths
@@ -171,7 +172,7 @@ def __init__(
 
         super().__init__()
 
-    def _fit(self, X, y=None):
+    def _fit(self, X: np.ndarray, y: Optional[Union[np.ndarray, List]] = None):
         """Fit the random dilated shapelet transform to a specified X and y.
 
         Parameters
@@ -245,7 +246,7 @@ def _fit(self, X, y=None):
 
         return self
 
-    def _transform(self, X, y=None):
+    def _transform(self, X: np.ndarray, y: Optional[Union[np.ndarray, List]] = None):
         """Transform X according to the extracted shapelets.
 
         Parameters
@@ -345,7 +346,9 @@ def _check_input_params(self):
             self.threshold_percentiles_ = np.asarray(self.threshold_percentiles_)
 
     @classmethod
-    def get_test_params(cls, parameter_set="default"):
+    def get_test_params(
+        cls, parameter_set: str = "default"
+    ) -> "Union[Dict, List[Dict]]":
         """Return testing parameter settings for the estimator.
 
         Parameters
@@ -365,6 +368,7 @@ def get_test_params(cls, parameter_set="default"):
         """
         if parameter_set == "default":
             params = {"max_shapelets": 10}
+            params = {"max_shapelets": 10}
         else:
             raise NotImplementedError(
                 f"The parameter set {parameter_set} is not yet implemented"
@@ -427,8 +431,9 @@ def _init_random_shapelet_params(
     # Lengths of the shapelets
     # test dtypes correctness
     lengths = np.random.choice(shapelet_lengths, size=max_shapelets).astype(np.int32)
+    lengths = np.random.choice(shapelet_lengths, size=max_shapelets).astype(np.int32)
     # Upper bound values for dilations
-    dilations = np.zeros(max_shapelets, dtype=np.int32)
+    dilations: np.ndarray = np.zeros(max_shapelets, dtype=np.int32)
     upper_bounds = np.log2(np.floor_divide(n_timepoints - 1, lengths - 1))
 
     if use_prime_dilations:
@@ -484,16 +489,16 @@ def _get_admissible_sampling_point(current_mask):
 
 @njit(fastmath=True, cache=True, parallel=True)
 def random_dilated_shapelet_extraction(
-    X,
-    y,
-    max_shapelets,
-    shapelet_lengths,
-    proba_normalization,
-    threshold_percentiles,
-    alpha_similarity,
-    use_prime_dilations,
-    seed,
-    distance,
+    X: np.ndarray,
+    y: np.ndarray,
+    max_shapelets: int,
+    shapelet_lengths: Union[np.ndarray, List[int]],
+    proba_normalization: float,
+    threshold_percentiles: Union[np.ndarray, List[float]],
+    alpha_similarity: float,
+    use_prime_dilations: bool,
+    seed: int,
+    distance: CPUDispatcher,
 ):
     """Randomly generate a set of shapelets given the input parameters.
 
@@ -554,9 +559,9 @@ def random_dilated_shapelet_extraction(
         - stds : array, shape (max_shapelets, n_channels)
             Standard deviation of the shapelets
     """
-    n_cases = len(X)
-    n_channels = X[0].shape[0]
-    n_timepointss = np.zeros(n_cases, dtype=np.int64)
+    n_cases: int = len(X)
+    n_channels: int = X[0].shape[0]
+    n_timepointss: np.ndarray = np.zeros(n_cases, dtype=np.int64)
     for i in range(n_cases):
         n_timepointss[i] = X[i].shape[1]
     min_n_timepoints = n_timepointss.min()
@@ -685,7 +690,19 @@ def random_dilated_shapelet_extraction(
 
 
 @njit(fastmath=True, cache=True, parallel=True)
-def dilated_shapelet_transform(X, shapelets, distance):
+def dilated_shapelet_transform(
+    X: np.ndarray,
+    shapelets: tuple[
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+    ],
+    distance: CPUDispatcher,
+):
     """Perform the shapelet transform with a set of shapelets and a set of time series.
 
     Parameters
@@ -797,7 +814,7 @@ def normalize_subsequences(X_subs, X_means, X_stds):
 
 
 @njit(fastmath=True, cache=True)
-def get_all_subsequences(X, length, dilation):
+def get_all_subsequences(X: np.ndarray, length: int, dilation: int) -> np.ndarray:
     """
     Generate subsequences from a time series given the length and dilation parameters.
 
@@ -828,7 +845,13 @@ def get_all_subsequences(X, length, dilation):
 
 
 @njit(fastmath=True, cache=True)
-def compute_shapelet_features(X_subs, values, length, threshold, distance):
+def compute_shapelet_features(
+    X_subs: np.ndarray,
+    values: np.ndarray,
+    length: int,
+    threshold: float,
+    distance: CPUDispatcher,
+):
     """Extract the features from a shapelet distance vector.
 
     Given a shapelet and a time series, extract three features from the resulting
@@ -875,7 +898,9 @@ def compute_shapelet_features(X_subs, values, length, threshold, distance):
 
 
 @njit(fastmath=True, cache=True)
-def compute_shapelet_dist_vector(X_subs, values, length, distance):
+def compute_shapelet_dist_vector(
+    X_subs: np.ndarray, values: np.ndarray, length: int, distance: CPUDispatcher
+):
     """Extract the features from a shapelet distance vector.
 
     Given a shapelet and a time series, extract three features from the resulting