Add model type exception, add classification threshold, return one-dimensional arrays from prediction funcs

Author: nmaarnio

eis_toolkit/exceptions.py

@@ -34,6 +34,10 @@ class InvalidDataShapeException(Exception):
     """Exception error for datasets with invalid shapes."""
 
 
+class InvalidModelTypeException(Exception):
+    """Exception error for invalid model type."""
+
+
 class InvalidParameterValueException(Exception):
     """Exception error class for invalid parameter values."""
 

eis_toolkit/prediction/machine_learning_predict.py

@@ -2,40 +2,55 @@
 import pandas as pd
 from beartype import beartype
 from beartype.typing import Tuple, Union
-from sklearn.base import BaseEstimator
+from sklearn.base import BaseEstimator, is_classifier
 from tensorflow import keras
 
+from eis_toolkit.exceptions import InvalidModelTypeException
+
 
 @beartype
 def predict_classifier(
-    data: Union[np.ndarray, pd.DataFrame], model: Union[BaseEstimator, keras.Model], include_probabilities: bool = True
+    data: Union[np.ndarray, pd.DataFrame],
+    model: Union[BaseEstimator, keras.Model],
+    classification_threshold: float = 0.5,
+    include_probabilities: bool = True,
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """
-    Predict with a trained model.
+    Predict with a trained classifier model.
+
+    Only works for binary classification currently.
 
     Args:
         data: Data used to make predictions.
         model: Trained classifier or regressor. Can be any machine learning model trained with
             EIS Toolkit (Sklearn and Keras models).
+        classification_threshold: Threshold for classifying based on probabilities. Defaults to 0.5.
         include_probabilities: If the probability array should be returned too. Defaults to True.
 
     Returns:
-        Predicted labels and optionally predicted probabilities by a classifier model.
+        Predicted labels and optionally predicted probabilities as one-dimensional arrays by a classifier model.
+
+    Raises:
+        InvalidModelTypeException: Input model is not a classifier model.
     """
     if isinstance(model, keras.Model):
-        probabilities = model.predict(data)
-        labels = probabilities.argmax(axis=-1)
+        probabilities = model.predict(data).squeeze()
+        labels = probabilities >= classification_threshold
         if include_probabilities:
-            return labels, probabilities
+            return labels, probabilities.astype(np.float32)
         else:
             return labels
     elif isinstance(model, BaseEstimator):
-        labels = model.predict(data)
+        if not is_classifier(model):
+            raise InvalidModelTypeException(f"Expected a classifier model: {type(model)}.")
+        probabilities = model.predict_proba(data)[:, 1]
+        labels = (probabilities >= classification_threshold).astype(np.float32)
         if include_probabilities:
-            probabilities = model.predict_proba(data)
-            return labels, probabilities
+            return labels, probabilities.astype(np.float32)
         else:
             return labels
+    else:
+        raise InvalidModelTypeException(f"Model type not recognized: {type(model)}.")
 
 
 @beartype
@@ -44,7 +59,7 @@ def predict_regressor(
     model: Union[BaseEstimator, keras.Model],
 ) -> np.ndarray:
     """
-    Predict with a trained model.
+    Predict with a trained regressor model.
 
     Args:
         data: Data used to make predictions.
@@ -53,6 +68,11 @@ def predict_regressor(
 
     Returns:
         Regression model prediction array.
+
+    Raises:
+        InvalidModelTypeException: Input model is not a regressor model.
     """
+    if is_classifier(model):
+        raise InvalidModelTypeException(f"Expected a regressor model: {type(model)}.")
     result = model.predict(data)
     return result