Formatting, formatted docstrings for Sphinx, minor fixes, added typing,

Author: SamoraHunter

ml_grid/model_classes/H2OAutoMLClassifier.py

@@ -1,22 +1,52 @@
+from typing import Any, Dict, Optional
+
 import h2o
 from h2o.automl import H2OAutoML
 import numpy as np
 import pandas as pd
 from sklearn.base import BaseEstimator, ClassifierMixin
-from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
+from sklearn.utils.validation import check_is_fitted
 
 
 class H2OAutoMLClassifier(BaseEstimator, ClassifierMixin):
-    def __init__(self, max_runtime_secs=360, nfolds=2, seed=1):
+    """A scikit-learn compatible wrapper for H2O's AutoML.
+
+    This class allows H2O's AutoML to be used as a standard scikit-learn
+    classifier, making it compatible with tools like GridSearchCV and
+    BayesSearchCV.
+    """
+
+    def __init__(
+        self, max_runtime_secs: int = 360, nfolds: int = 2, seed: int = 1
+    ):
+        """Initializes the H2OAutoMLClassifier.
+
+        Args:
+            max_runtime_secs (int): Maximum time in seconds to run the AutoML process.
+            nfolds (int): Number of folds for cross-validation.
+            seed (int): Random seed for reproducibility.
+        """
         self.max_runtime_secs = max_runtime_secs
         self.nfolds = nfolds
         self.seed = seed
-        self.automl = None
+        self.automl: Optional[H2OAutoML] = None
+        self.classes_: Optional[np.ndarray] = None
+
+    def fit(self, X: pd.DataFrame, y: pd.Series) -> "H2OAutoMLClassifier":
+        """Fits the H2O AutoML model.
 
-    def fit(self, X, y):
-        # X, y = check_X_y(X, y)
+        This method initializes an H2O cluster, converts the pandas DataFrame
+        and Series to H2O Frames, and then trains the AutoML model.
+
+        Args:
+            X (pd.DataFrame): The training input samples.
+            y (pd.Series): The target values.
+
+        Returns:
+            H2OAutoMLClassifier: The fitted estimator.
+        """
         self.classes_ = np.unique(y)
-        # outcome_var = y.columns[0]
+
         try:
             outcome_var = y.columns[0]
         except:
@@ -31,10 +61,7 @@ def fit(self, X, y):
             pass
 
         h2o.init()
-        # train_df = pd.concat([pd.DataFrame(X), pd.Series(y)], axis=1)
-
         train_df = pd.concat([X, y], axis=1)
-
         train_h2o = h2o.H2OFrame(train_df)
 
         train_h2o[y_n] = train_h2o[y_n].asfactor()
@@ -49,33 +76,69 @@ def fit(self, X, y):
         self.automl.train(y=y_n, x=x, training_frame=train_h2o)
         return self
 
-    def predict(self, X):
+    def predict(self, X: pd.DataFrame) -> np.ndarray:
+        """Predicts class labels for samples in X.
+
+        Args:
+            X (pd.DataFrame): The input samples to predict.
+
+        Returns:
+            np.ndarray: The predicted class labels.
+        """
         check_is_fitted(self)
-        # X = check_array(X)
-        # test_h2o = h2o.H2OFrame(pd.DataFrame(X))
         test_h2o = h2o.H2OFrame(X)
         predictions = self.automl.leader.predict(test_h2o)
 
-        # return predictions[:,0]
         return predictions["predict"].as_data_frame().values
 
-    def predict_proba(self, X):
+    def predict_proba(self, X: pd.DataFrame) -> np.ndarray:
+        """Predicts class probabilities for samples in X.
+
+        Note:
+            This method is not implemented for H2O AutoML.
+
+        Args:
+            X (pd.DataFrame): The input samples.
+
+        Raises:
+            NotImplementedError: H2O AutoML does not support predict_proba.
+        """
         raise NotImplementedError("H2O AutoML does not support predict_proba.")
 
-    def get_params(self, deep=True):
+    def get_params(self, deep: bool = True) -> Dict[str, Any]:
+        """Gets parameters for this estimator.
+
+        Args:
+            deep (bool): If True, will return the parameters for this estimator and
+                contained subobjects that are estimators.
 
+        Returns:
+            Dict[str, Any]: Parameter names mapped to their values.
+        """
         return {
             "max_runtime_secs": self.max_runtime_secs,
             "nfolds": self.nfolds,
             "seed": self.seed,
         }
 
-    def set_params(self, **params):
+    def set_params(self, **params: Any) -> "H2OAutoMLClassifier":
+        """Sets the parameters of this estimator.
+
+        Args:
+            **params (Any): Estimator parameters.
+
+        Returns:
+            H2OAutoMLClassifier: The instance with updated parameters.
+        """
         for param, value in params.items():
             setattr(self, param, value)
         return self
 
-    def get_leader_params(
-        self,
-    ):
+    def get_leader_params(self) -> Dict[str, Any]:
+        """Gets the parameters of the best model found by AutoML.
+
+        Returns:
+            Dict[str, Any]: A dictionary of the leader model's parameters.
+        """
+        check_is_fitted(self)
         return self.automl.leader.params

ml_grid/model_classes/NeuralNetworkClassifier_class.py

@@ -1,5 +1,8 @@
-"""Define NeuralNetworkClassifier class"""
+"""Defines the NeuralNetworkClassifier model class."""
 
+from typing import Optional
+
+import pandas as pd
 from ml_grid.util import param_space
 
 # from ml_grid.model_classes.nni_sklearn_wrapper import *
@@ -9,14 +12,24 @@
 
 
 class NeuralNetworkClassifier_class:
-    """NeuralNetworkClassifier."""
+    """NeuralNetworkClassifier with a predefined parameter space."""
 
-    def __init__(self, X=None, y=None, parameter_space_size=None, global_parameters=None):
-        """_summary_
+    def __init__(
+        self,
+        X: Optional[pd.DataFrame] = None,
+        y: Optional[pd.Series] = None,
+        parameter_space_size: Optional[str] = None,
+    ):
+        """Initializes the NeuralNetworkClassifier_class.
 
         Args:
-            X_train (_type_): _description_
-            y_train (_type_): _description_
+            X (Optional[pd.DataFrame]): Feature matrix for training.
+                Defaults to None.
+            y (Optional[pd.Series]): Target vector for training.
+                Defaults to None.
+            parameter_space_size (Optional[str]): Size of the parameter space for
+                optimization. This is not used in the current implementation
+                as the parameter space is hardcoded. Defaults to None.
         """
         self.X = X
         self.y = y

ml_grid/model_classes/adaboost_classifier_class.py

@@ -1,3 +1,6 @@
+from typing import Optional
+
+import pandas as pd
 from sklearn.ensemble import AdaBoostClassifier
 from ml_grid.util import param_space
 from ml_grid.util.global_params import global_parameters
@@ -8,14 +11,21 @@
 class adaboost_class:
     """AdaBoostClassifier with support for both Bayesian and non-Bayesian parameter spaces."""
 
-    def __init__(self, X=None, y=None, parameter_space_size=None):
-        """
-        Initialize the adaboost_class.
+    def __init__(
+        self,
+        X: Optional[pd.DataFrame] = None,
+        y: Optional[pd.Series] = None,
+        parameter_space_size: Optional[str] = None,
+    ):
+        """Initializes the adaboost_class.
 
         Args:
-            X (_type_): Feature matrix for training (optional).
-            y (_type_): Target vector for training (optional).
-            parameter_space_size (_type_): Size of the parameter space for optimization.
+            X (Optional[pd.DataFrame]): Feature matrix for training.
+                Defaults to None.
+            y (Optional[pd.Series]): Target vector for training.
+                Defaults to None.
+            parameter_space_size (Optional[str]): Size of the parameter space for
+                optimization. Defaults to None.
         """
         global_params = global_parameters
         self.X = X

ml_grid/model_classes/catboostClassifier.py

@@ -1,16 +1,56 @@
+from typing import Any, Union
+
+import numpy as np
+import pandas as pd
 from catboost import CatBoostClassifier
 from sklearn.base import BaseEstimator, ClassifierMixin
 
 
 class CatBoostSKLearnWrapper(BaseEstimator, ClassifierMixin):
-    def __init__(self, **kwargs):
+    """A scikit-learn compatible wrapper for the CatBoostClassifier."""
+
+    def __init__(self, **kwargs: Any):
+        """Initializes the CatBoostSKLearnWrapper.
+
+        Args:
+            **kwargs (Any): Keyword arguments passed directly to the
+                `catboost.CatBoostClassifier`.
+        """
         self.model = CatBoostClassifier(**kwargs)
 
-    def fit(self, X, y):
+    def fit(
+        self, X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray]
+    ) -> "CatBoostSKLearnWrapper":
+        """Fits the CatBoost model.
+
+        Args:
+            X (Union[pd.DataFrame, np.ndarray]): The training input samples.
+            y (Union[pd.Series, np.ndarray]): The target values.
+
+        Returns:
+            CatBoostSKLearnWrapper: The fitted estimator.
+        """
         self.model.fit(X, y)
+        return self
 
-    def predict(self, X):
+    def predict(self, X: Union[pd.DataFrame, np.ndarray]) -> np.ndarray:
+        """Predicts class labels for samples in X.
+
+        Args:
+            X (Union[pd.DataFrame, np.ndarray]): The input samples to predict.
+
+        Returns:
+            np.ndarray: The predicted class labels.
+        """
         return self.model.predict(X)
 
-    def predict_proba(self, X):
+    def predict_proba(self, X: Union[pd.DataFrame, np.ndarray]) -> np.ndarray:
+        """Predicts class probabilities for samples in X.
+
+        Args:
+            X (Union[pd.DataFrame, np.ndarray]): The input samples.
+
+        Returns:
+            np.ndarray: The class probabilities of the input samples.
+        """
         return self.model.predict_proba(X)

ml_grid/model_classes/catboost_classifier_class.py

@@ -1,4 +1,7 @@
+from typing import Optional
+
 import numpy as np
+import pandas as pd
 from catboost import CatBoostClassifier
 from skopt.space import Categorical, Real, Integer
 from ml_grid.util import param_space
@@ -7,14 +10,21 @@
 class CatBoost_class:
     """CatBoost Classifier with hyperparameter tuning."""
 
-    def __init__(self, X=None, y=None, parameter_space_size=None):
-        """
-        Initialize the CatBoost_class.
+    def __init__(
+        self,
+        X: Optional[pd.DataFrame] = None,
+        y: Optional[pd.Series] = None,
+        parameter_space_size: Optional[str] = None,
+    ):
+        """Initializes the CatBoost_class.
 
         Args:
-            X (_type_): Feature matrix for training (optional).
-            y (_type_): Target vector for training (optional).
-            parameter_space_size (_type_): Size of the parameter space for optimization.
+            X (Optional[pd.DataFrame]): Feature matrix for training.
+                Defaults to None.
+            y (Optional[pd.Series]): Target vector for training.
+                Defaults to None.
+            parameter_space_size (Optional[str]): Size of the parameter space for
+                optimization. Defaults to None.
         """
         global_params = global_parameters  # Fetch global parameters
         self.X = X