[ENH] Add type hints for deep learning regression classes

aeon/regression/deep_learning/_cnn.py

@@ -1,18 +1,26 @@
 """Time Convolutional Neural Network (TimeCNN) regressor."""
 
+from __future__ import annotations
+
 __maintainer__ = ["hadifawaz1999"]
 __all__ = ["TimeCNNRegressor"]
 
 import gc
 import os
 import time
 from copy import deepcopy
+from typing import TYPE_CHECKING, Any
 
+import numpy as np
 from sklearn.utils import check_random_state
 
 from aeon.networks import TimeCNNNetwork
 from aeon.regression.deep_learning.base import BaseDeepRegressor
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    from tensorflow.keras.callbacks import Callback
+
 
 class TimeCNNRegressor(BaseDeepRegressor):
     """Time Series Convolutional Neural Network (CNN).
@@ -120,39 +128,39 @@ class TimeCNNRegressor(BaseDeepRegressor):
     >>> X, y = make_example_3d_numpy(n_cases=10, n_channels=1, n_timepoints=12,
     ...                              return_y=True, regression_target=True,
     ...                              random_state=0)
-    >>> rgs = TimeCNNRegressor(n_epochs=20, bacth_size=4) # doctest: +SKIP
+    >>> rgs = TimeCNNRegressor(n_epochs=20, batch_size=4) # doctest: +SKIP
     >>> rgs.fit(X, y) # doctest: +SKIP
     TimeCNNRegressor(...)
     """
 
     def __init__(
         self,
-        n_layers=2,
-        kernel_size=7,
-        n_filters=None,
-        avg_pool_size=3,
-        activation="sigmoid",
-        padding="valid",
-        strides=1,
-        dilation_rate=1,
-        n_epochs=2000,
-        batch_size=16,
-        callbacks=None,
-        file_path="./",
-        save_best_model=False,
-        save_last_model=False,
-        save_init_model=False,
-        best_file_name="best_model",
-        last_file_name="last_model",
-        init_file_name="init_model",
-        verbose=False,
-        loss="mean_squared_error",
-        output_activation="linear",
-        metrics="mean_squared_error",
-        random_state=None,
-        use_bias=True,
-        optimizer=None,
-    ):
+        n_layers: int = 2,
+        kernel_size: int | list[int] = 7,
+        n_filters: int | list[int] | None = None,
+        avg_pool_size: int | list[int] = 3,
+        activation: str | list[str] = "sigmoid",
+        padding: str | list[str] = "valid",
+        strides: int | list[int] = 1,
+        dilation_rate: int | list[int] = 1,
+        n_epochs: int = 2000,
+        batch_size: int = 16,
+        callbacks: Callback | list[Callback] | None = None,
+        file_path: str = "./",
+        save_best_model: bool = False,
+        save_last_model: bool = False,
+        save_init_model: bool = False,
+        best_file_name: str = "best_model",
+        last_file_name: str = "last_model",
+        init_file_name: str = "init_model",
+        verbose: bool = False,
+        loss: str = "mean_squared_error",
+        output_activation: str = "linear",
+        metrics: str | list[str] = "mean_squared_error",
+        random_state: int | np.random.RandomState | None = None,
+        use_bias: bool | list[bool] = True,
+        optimizer: tf.keras.optimizers.Optimizer | None = None,
+    ) -> None:
         self.n_layers = n_layers
         self.avg_pool_size = avg_pool_size
         self.padding = padding
@@ -196,7 +204,9 @@ def __init__(
             use_bias=self.use_bias,
         )
 
-    def build_model(self, input_shape, **kwargs):
+    def build_model(
+        self, input_shape: tuple[int, ...], **kwargs: Any
+    ) -> tf.keras.Model:
         """Construct a compiled, un-trained, keras model that is ready for training.
 
         In aeon, time series are stored in numpy arrays of shape (d,m), where d
@@ -213,7 +223,6 @@ def build_model(self, input_shape, **kwargs):
         -------
         output : a compiled Keras Model
         """
-        import numpy as np
         import tensorflow as tf
         from tensorflow import keras
 
@@ -239,7 +248,7 @@ def build_model(self, input_shape, **kwargs):
         )
         return model
 
-    def _fit(self, X, y):
+    def _fit(self, X: np.ndarray, y: np.ndarray) -> TimeCNNRegressor:
         """Fit the regressor on the training set (X, y).
 
         Parameters
@@ -316,7 +325,9 @@ def _fit(self, X, y):
         return self
 
     @classmethod
-    def _get_test_params(cls, parameter_set="default"):
+    def _get_test_params(
+        cls, parameter_set: str = "default"
+    ) -> dict[str, Any] | list[dict[str, Any]]:
         """Return testing parameter settings for the estimator.
 
         Parameters

aeon/regression/deep_learning/_disjoint_cnn.py

@@ -1,18 +1,26 @@
 """DisjointCNN regressor."""
 
+from __future__ import annotations
+
 __maintainer__ = ["hadifawaz1999"]
 __all__ = ["DisjointCNNRegressor"]
 
 import gc
 import os
 import time
 from copy import deepcopy
+from typing import TYPE_CHECKING, Any
 
+import numpy as np
 from sklearn.utils import check_random_state
 
 from aeon.networks import DisjointCNNNetwork
 from aeon.regression.deep_learning.base import BaseDeepRegressor
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    from tensorflow.keras.callbacks import Callback
+
 
 class DisjointCNNRegressor(BaseDeepRegressor):
     """Disjoint Convolutional Neural Netowkr regressor.
@@ -159,37 +167,37 @@ class DisjointCNNRegressor(BaseDeepRegressor):
 
     def __init__(
         self,
-        n_layers=4,
-        n_filters=64,
-        kernel_size=None,
-        dilation_rate=1,
-        strides=1,
-        padding="same",
-        activation="elu",
-        use_bias=True,
-        kernel_initializer="he_uniform",
-        pool_size=5,
-        pool_strides=None,
-        pool_padding="valid",
-        hidden_fc_units=128,
-        activation_fc="relu",
-        n_epochs=2000,
-        batch_size=16,
-        use_mini_batch_size=False,
-        random_state=None,
-        verbose=False,
-        output_activation="linear",
-        loss="mean_squared_error",
-        metrics="mean_squared_error",
-        optimizer=None,
-        file_path="./",
-        save_best_model=False,
-        save_last_model=False,
-        save_init_model=False,
-        best_file_name="best_model",
-        last_file_name="last_model",
-        init_file_name="init_model",
-        callbacks=None,
+        n_layers: int = 4,
+        n_filters: int | list[int] = 64,
+        kernel_size: int | list[int] | None = None,
+        dilation_rate: int | list[int] = 1,
+        strides: int | list[int] = 1,
+        padding: str | list[str] = "same",
+        activation: str | list[str] = "elu",
+        use_bias: bool | list[bool] = True,
+        kernel_initializer: str | list[str] = "he_uniform",
+        pool_size: int = 5,
+        pool_strides: int | None = None,
+        pool_padding: str = "valid",
+        hidden_fc_units: int = 128,
+        activation_fc: str = "relu",
+        n_epochs: int = 2000,
+        batch_size: int = 16,
+        use_mini_batch_size: bool = False,
+        random_state: int | np.random.RandomState | None = None,
+        verbose: bool = False,
+        output_activation: str = "linear",
+        loss: str = "mean_squared_error",
+        metrics: str | list[str] = "mean_squared_error",
+        optimizer: tf.keras.optimizers.Optimizer | None = None,
+        file_path: str = "./",
+        save_best_model: bool = False,
+        save_last_model: bool = False,
+        save_init_model: bool = False,
+        best_file_name: str = "best_model",
+        last_file_name: str = "last_model",
+        init_file_name: str = "init_model",
+        callbacks: Callback | list[Callback] | None = None,
     ):
         self.n_layers = n_layers
         self.n_filters = n_filters
@@ -247,7 +255,9 @@ def __init__(
             activation_fc=self.activation_fc,
         )
 
-    def build_model(self, input_shape, **kwargs):
+    def build_model(
+        self, input_shape: tuple[int, ...], **kwargs: Any
+    ) -> tf.keras.Model:
         """Construct a compiled, un-trained, keras model that is ready for training.
 
         In aeon, time series are stored in numpy arrays of shape (d,m), where d
@@ -266,7 +276,6 @@ def build_model(self, input_shape, **kwargs):
         -------
         output : a compiled Keras Model
         """
-        import numpy as np
         import tensorflow as tf
 
         rng = check_random_state(self.random_state)
@@ -291,7 +300,7 @@ def build_model(self, input_shape, **kwargs):
 
         return model
 
-    def _fit(self, X, y):
+    def _fit(self, X: np.ndarray, y: np.ndarray) -> DisjointCNNRegressor:
         """Fit the regressor on the training set (X, y).
 
         Parameters
@@ -376,7 +385,9 @@ def _fit(self, X, y):
         return self
 
     @classmethod
-    def _get_test_params(cls, parameter_set="default"):
+    def _get_test_params(
+        cls, parameter_set: str = "default"
+    ) -> dict[str, Any] | list[dict[str, Any]]:
         """Return testing parameter settings for the estimator.
 
         Parameters

aeon/regression/deep_learning/_encoder.py

@@ -1,18 +1,27 @@
 """Encoder Regressor."""
 
+from __future__ import annotations
+
 __author__ = ["AnonymousCodes911", "hadifawaz1999"]
 __all__ = ["EncoderRegressor"]
 
+
 import gc
 import os
 import time
 from copy import deepcopy
+from typing import TYPE_CHECKING, Any
 
+import numpy as np
 from sklearn.utils import check_random_state
 
 from aeon.networks import EncoderNetwork
 from aeon.regression.deep_learning.base import BaseDeepRegressor
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    from tensorflow.keras.callbacks import Callback
+
 
 class EncoderRegressor(BaseDeepRegressor):
     """
@@ -111,31 +120,31 @@ class EncoderRegressor(BaseDeepRegressor):
 
     def __init__(
         self,
-        n_epochs=100,
-        batch_size=12,
-        kernel_size=None,
-        n_filters=None,
-        dropout_proba=0.2,
-        activation="sigmoid",
-        output_activation="linear",
-        max_pool_size=2,
-        padding="same",
-        strides=1,
-        fc_units=256,
-        callbacks=None,
-        file_path="./",
-        save_best_model=False,
-        save_last_model=False,
-        save_init_model=False,
-        best_file_name="best_model",
-        last_file_name="last_model",
-        init_file_name="init_model",
-        verbose=False,
-        loss="mean_squared_error",
-        metrics="mean_squared_error",
-        use_bias=True,
-        optimizer=None,
-        random_state=None,
+        n_epochs: int = 100,
+        batch_size: int = 12,
+        kernel_size: list[int] | None = None,
+        n_filters: list[int] | None = None,
+        dropout_proba: float = 0.2,
+        activation: str = "sigmoid",
+        output_activation: str = "linear",
+        max_pool_size: int = 2,
+        padding: str = "same",
+        strides: int = 1,
+        fc_units: int = 256,
+        callbacks: Callback | list[Callback] | None = None,
+        file_path: str = "./",
+        save_best_model: bool = False,
+        save_last_model: bool = False,
+        save_init_model: bool = False,
+        best_file_name: str = "best_model",
+        last_file_name: str = "last_model",
+        init_file_name: str = "init_model",
+        verbose: bool = False,
+        loss: str = "mean_squared_error",
+        metrics: str | list[str] = "mean_squared_error",
+        use_bias: bool = True,
+        optimizer: tf.keras.optimizers.Optimizer | None = None,
+        random_state: int | np.random.RandomState | None = None,
     ):
         self.n_filters = n_filters
         self.max_pool_size = max_pool_size
@@ -179,7 +188,9 @@ def __init__(
             activation=self.activation,
         )
 
-    def build_model(self, input_shape, **kwargs):
+    def build_model(
+        self, input_shape: tuple[int, ...], **kwargs: Any
+    ) -> tf.keras.Model:
         """Construct a compiled, un-trained, keras model that is ready for training.
 
         In aeon, time series are stored in numpy arrays of shape (d, m), where d
@@ -195,7 +206,6 @@ def build_model(self, input_shape, **kwargs):
         -------
         output : a compiled Keras Model
         """
-        import numpy as np
         import tensorflow as tf
 
         rng = check_random_state(self.random_state)
@@ -222,7 +232,7 @@ def build_model(self, input_shape, **kwargs):
 
         return model
 
-    def _fit(self, X, y):
+    def _fit(self, X: np.ndarray, y: np.ndarray) -> EncoderRegressor:
         """Fit the classifier on the training set (X, y).
 
         Parameters
@@ -299,7 +309,9 @@ def _fit(self, X, y):
         return self
 
     @classmethod
-    def _get_test_params(cls, parameter_set="default"):
+    def _get_test_params(
+        cls, parameter_set: str = "default"
+    ) -> dict[str, Any] | list[dict[str, Any]]:
         """Return testing parameter settings for the estimator.
 
         Parameters