Added func to test/score Sklearn model

Author: nmaarnio

eis_toolkit/prediction/model_utils.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 from beartype import beartype
-from beartype.typing import List, Literal, Optional, Sequence, Tuple, Union
+from beartype.typing import Dict, List, Literal, Optional, Sequence, Tuple, Union
 from scipy import sparse
 from sklearn.base import BaseEstimator, is_classifier, is_regressor
 from sklearn.metrics import (
@@ -91,14 +91,55 @@ def split_data(
 
 
 @beartype
-def predict(model: Union[BaseEstimator, keras.Model], data: np.ndarray) -> np.ndarray:
+def test_model(
+    X_test: Union[np.ndarray, pd.DataFrame],
+    y_test: Union[np.ndarray, pd.Series],
+    model: Union[BaseEstimator, keras.Model],
+    metrics: Optional[Sequence[Literal["mse", "rmse", "mae", "r2", "accuracy", "precision", "recall", "f1"]]] = None,
+) -> Dict[str, Number]:
+    """
+    Test and score a trained model.
+
+    TODO: Implement for Keras models.
+
+    Args:
+        X_test: Test data.
+        y_test: Target labels for test data.
+        model: Trained Sklearn classifier or regressor.
+        metrics: Metrics to use for scoring the model. Defaults to "accuracy" for a classifier
+            and to "mse" for a regressor.
+
+    Returns:
+        Test metric scores as a dictionary.
+    """
+    x_size = X_test.index.size if isinstance(X_test, pd.DataFrame) else X_test.shape[0]
+    if x_size != y_test.size:
+        raise exceptions.NonMatchingParameterLengthsException(
+            f"X and y must have the length {x_size} != {y_test.size}."
+        )
+
+    if metrics is None:
+        metrics = ["accuracy"] if is_classifier(model) else ["mse"]
+
+    y_pred = model.predict(X_test)
+
+    out_metrics = {}
+    for metric in metrics:
+        score = _score_model(model, y_test, y_pred, metric)
+        out_metrics[metric] = score
+
+    return out_metrics
+
+
+@beartype
+def predict(data: Union[np.ndarray, pd.DataFrame], model: Union[BaseEstimator, keras.Model]) -> np.ndarray:
     """
     Predict with a trained model.
 
     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).
-        data: Data used to make predictions.
 
     Returns:
         Predictions.

tests/prediction/model_utils_test.py

@@ -14,6 +14,7 @@
     save_model,
     split_data,
 )
+from eis_toolkit.prediction.model_utils import test_model as model_test
 
 TEST_DIR = Path(__file__).parent.parent
 
@@ -50,7 +51,7 @@ def test_train_and_evaluate_with_split():
     )
 
     assert isinstance(model, RandomForestClassifier)
-    assert len(out_metrics) == 4
+    np.testing.assert_equal(len(out_metrics), 4)
 
 
 def test_train_and_evaluate_with_kfold_cv():
@@ -60,7 +61,7 @@ def test_train_and_evaluate_with_kfold_cv():
     )
 
     assert isinstance(model, RandomForestClassifier)
-    assert len(out_metrics) == 4
+    np.testing.assert_equal(len(out_metrics), 4)
 
 
 def test_train_and_evaluate_with_skfold_cv():
@@ -70,7 +71,7 @@ def test_train_and_evaluate_with_skfold_cv():
     )
 
     assert isinstance(model, RandomForestClassifier)
-    assert len(out_metrics) == 4
+    np.testing.assert_equal(len(out_metrics), 4)
 
 
 def test_binary_classification():
@@ -97,7 +98,7 @@ def test_binary_classification():
     )
 
     assert isinstance(model, RandomForestClassifier)
-    assert len(out_metrics) == 4
+    np.testing.assert_equal(len(out_metrics), 4)
 
 
 def test_splitting():
@@ -109,6 +110,18 @@ def test_splitting():
     np.testing.assert_equal(len(y_test), len(Y_IRIS) * 0.2)
 
 
+def test_test_model_sklearn():
+    """Test that test model works as expected with a Sklearn model."""
+    X_train, X_test, y_train, y_test = split_data(X_IRIS, Y_IRIS, split_size=0.2)
+
+    model, _ = _train_and_validate_sklearn_model(
+        X_train, y_train, model=RF_MODEL, validation_method="none", metrics=CLF_METRICS, random_state=42
+    )
+
+    out_metrics = model_test(X_test, y_test, model)
+    np.testing.assert_equal(out_metrics["accuracy"], 1.0)
+
+
 def test_predict_sklearn():
     """Test that predict works as expected with a Sklearn model."""
     X_train, X_test, y_train, y_test = split_data(X_IRIS, Y_IRIS, split_size=0.2)
@@ -117,8 +130,8 @@ def test_predict_sklearn():
         X_train, y_train, model=RF_MODEL, validation_method="none", metrics=CLF_METRICS, random_state=42
     )
 
-    predicted_labels = predict(model, X_test)
-    assert len(predicted_labels) == len(y_test)
+    predicted_labels = predict(X_test, model)
+    np.testing.assert_equal(len(predicted_labels), len(y_test))
 
 
 def test_save_and_load_model():