[DOC] Added Type Hints for _dilated_shapelet_transform

aeon/transformations/collection/shapelet_based/_dilated_shapelet_transform.py

@@ -8,9 +8,13 @@
 __all__ = ["RandomDilatedShapeletTransform"]
 
 import warnings
+from typing import Dict
+from typing import List as TypingList
+from typing import 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
 
@@ -151,15 +155,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[TypingList[int], np.ndarray]] = None,
+        proba_normalization: float = 0.8,
+        threshold_percentiles: Optional[Union[TypingList[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
@@ -173,7 +177,7 @@ def __init__(
 
         super().__init__()
 
-    def _fit(self, X, y=None):
+    def _fit(self, X: np.ndarray, y: Optional[Union[np.ndarray, TypingList]] = None):
         """Fit the random dilated shapelet transform to a specified X and y.
 
         Parameters
@@ -247,7 +251,9 @@ 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, TypingList]] = None
+    ):
         """Transform X according to the extracted shapelets.
 
         Parameters
@@ -347,7 +353,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, TypingList[Dict]]":
         """Return testing parameter settings for the estimator.
 
         Parameters
@@ -376,12 +384,12 @@ def get_test_params(cls, parameter_set="default"):
 
 @njit(fastmath=True, cache=True)
 def _init_random_shapelet_params(
-    max_shapelets,
-    shapelet_lengths,
-    proba_normalization,
-    use_prime_dilations,
-    n_channels,
-    n_timepoints,
+    max_shapelets: int,
+    shapelet_lengths: np.ndarray,
+    proba_normalization: float,
+    use_prime_dilations: bool,
+    n_channels: int,
+    n_timepoints: int,
 ):
     """Randomly initialize the parameters of the shapelets.
 
@@ -505,16 +513,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: np.ndarray,
+    proba_normalization: float,
+    threshold_percentiles: np.ndarray,
+    alpha_similarity: float,
+    use_prime_dilations: bool,
+    seed: int,
+    distance: CPUDispatcher,
 ):
     """Randomly generate a set of shapelets given the input parameters.
 
@@ -721,7 +729,21 @@ 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,
+        np.ndarray,
+        np.ndarray,
+    ],
+    distance: CPUDispatcher,
+):
     """Perform the shapelet transform with a set of shapelets and a set of time series.
 
     Parameters
@@ -808,7 +830,11 @@ def dilated_shapelet_transform(X, shapelets, distance):
 
 
 @njit(fastmath=True, cache=True)
-def normalize_subsequences(X_subs, X_means, X_stds):
+def normalize_subsequences(
+    X_subs: np.ndarray[np.float64],
+    X_means: np.ndarray[np.float64],
+    X_stds: np.ndarray[np.float64],
+):
     """
     Generate subsequences from a time series given the length and dilation parameters.
 
@@ -839,7 +865,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.
 
@@ -870,7 +896,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
@@ -917,7 +949,12 @@ 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