From cd76fe182baa35e7db9b45e45a79628b7d2d64de Mon Sep 17 00:00:00 2001 From: Debi-Prasad-Panda Date: Thu, 11 Jun 2026 23:50:48 +0530 Subject: [PATCH] feat: add predict_uncertainty() for ensemble-based regression intervals --- supervised/automl.py | 29 ++++ supervised/base_automl.py | 126 ++++++++++++++++++ .../tests_automl/test_predict_uncertainty.py | 124 +++++++++++++++++ 3 files changed, 279 insertions(+) create mode 100644 tests/tests_automl/test_predict_uncertainty.py diff --git a/supervised/automl.py b/supervised/automl.py index b3702f62..e3b32456 100644 --- a/supervised/automl.py +++ b/supervised/automl.py @@ -601,3 +601,32 @@ def need_retrain( bool: Decides if there is a need to retrain the AutoML. """ return self._need_retrain(X, y, sample_weight, decrease) + + def predict_uncertainty( + self, + X: Union[List, numpy.ndarray, pandas.DataFrame], + alpha: float = 0.05, + ) -> pandas.DataFrame: + """Computes the ensemble-based uncertainty intervals for a regression task. + + This method estimates prediction uncertainty by analyzing the disagreement/variance + among the individual sub-models selected within the trained Ensemble. + + Note: + This represents an ensemble uncertainty interval based on model disagreement, + not a statistically calibrated confidence interval. + + Args: + X (Union[List, numpy.ndarray, pandas.DataFrame]): Input features for predictions. + alpha (float): Significance level for the interval (e.g., alpha=0.05 for a 95% interval). + Must be between 0 and 1 exclusive. Defaults to 0.05. + + Returns: + pandas.DataFrame: A DataFrame containing: + - `prediction`: Weighted mean of the sub-models' predictions. + - `prediction_std`: Standard deviation of the sub-models' predictions. + - `prediction_variance`: Variance of the sub-models' predictions. + - `lower`: Lower bound of the uncertainty interval. + - `upper`: Upper bound of the uncertainty interval. + """ + return self._predict_uncertainty(X, alpha) \ No newline at end of file diff --git a/supervised/base_automl.py b/supervised/base_automl.py index 5310304b..aaf24755 100644 --- a/supervised/base_automl.py +++ b/supervised/base_automl.py @@ -11,6 +11,7 @@ import joblib import numpy as np import pandas as pd +import scipy.stats from sklearn.base import BaseEstimator from sklearn.metrics import accuracy_score, r2_score from sklearn.utils.validation import check_array @@ -1567,6 +1568,131 @@ def _predict_all(self, X): # Make and return predictions return self._base_predict(X) + def _predict_uncertainty(self, X, alpha=0.05): + """Internal method to compute ensemble-based uncertainty intervals. + + Args: + X (Union[List, numpy.ndarray, pandas.DataFrame]): Input features. + alpha (float): Significance level for the interval. + + Returns: + pandas.DataFrame: Prediction interval results. + """ + if self._best_model is None: + try: + self.load(self.results_path) + except Exception: + pass + + if self._best_model is None: + raise AutoMLException("This model has not been fitted yet. Please call `fit()` first.") + + if self._ml_task != REGRESSION: + raise AutoMLException( + f"Predict uncertainty is only supported for regression tasks. Current task is '{self._ml_task}'." + ) + + # Use the Ensemble model for uncertainty; fall back from best model if needed + uncertainty_model = self._best_model + if uncertainty_model.get_type() != "Ensemble": + ensemble_model = None + for m in self._models: + if m.get_type() == "Ensemble": + ensemble_model = m + break + + # Attempt to load Ensemble from disk if not found in memory + if ensemble_model is None and hasattr(self, "_model_subpaths") and "Ensemble" in self._model_subpaths: + try: + models_map = {m.get_name(): m for m in self._models} + needed_submodels = self.get_ensemble_models("Ensemble") + for submodel_name in needed_submodels: + if submodel_name not in models_map: + m = ModelFramework.load(self._results_path, submodel_name, False) + self._models.append(m) + models_map[m.get_name()] = m + ens = Ensemble.load(self._results_path, "Ensemble", models_map) + self._models.append(ens) + ensemble_model = ens + except Exception as e: + logger.warning(f"Failed to load Ensemble model: {str(e)}") + + if ensemble_model is not None: + uncertainty_model = ensemble_model + logger.warning( + f"The best selected model is '{self._best_model.get_type()}', not an Ensemble. " + f"Using the trained Ensemble model '{ensemble_model.get_name()}' to compute uncertainty intervals." + ) + else: + raise AutoMLException( + f"Predict uncertainty is only supported when an Ensemble model is available. " + f"Current best model type is '{self._best_model.get_type()}' and no Ensemble model was found." + ) + + if not (0 < alpha < 1): + raise AutoMLException("Parameter alpha must be between 0 and 1 exclusive.") + + if not hasattr(uncertainty_model, "selected_models") or not uncertainty_model.selected_models: + raise AutoMLException("The ensemble model does not contain any submodels.") + + X = self._build_dataframe(X) + if not isinstance(X.columns[0], str): + X.columns = [str(c) for c in X.columns] + + input_columns = X.columns.tolist() + for column in self._data_info["columns"]: + if column not in input_columns: + raise AutoMLException( + f"Missing column: {column} in input data. Cannot predict" + ) + + X = X[self._data_info["columns"]] + self._validate_X_predict(X) + + X_stacked = None + if uncertainty_model._is_stacked: + self._perform_model_stacking() + X_stacked = self.get_stacked_data(X, mode="predict") + + # Collect predictions from each sub-model in the ensemble + predictions_list = [] + weights = [] + for selected in uncertainty_model.selected_models: + submodel = selected["model"] + weight = selected["repeat"] + weights.append(weight) + + if submodel._is_stacked: + pred = submodel.predict(X_stacked) + else: + pred = submodel.predict(X) + + predictions_list.append(pred["prediction"].to_numpy()) + + preds_array = np.array(predictions_list) # shape: (n_models, n_samples) + weights_array = np.array(weights) + + # Weighted mean, variance, and standard deviation across sub-models + mu = np.average(preds_array, weights=weights_array, axis=0) + var = np.average((preds_array - mu) ** 2, weights=weights_array, axis=0) + std = np.sqrt(var) + + # Compute interval bounds using the normal distribution z-score + z = scipy.stats.norm.ppf(1.0 - alpha / 2.0) + lower = mu - z * std + upper = mu + z * std + + return pd.DataFrame( + { + "prediction": mu, + "prediction_std": std, + "prediction_variance": var, + "lower": lower, + "upper": upper, + }, + index=X.index, + ) + def _score(self, X, y=None, sample_weight=None): # y default must be None for scikit-learn compatibility diff --git a/tests/tests_automl/test_predict_uncertainty.py b/tests/tests_automl/test_predict_uncertainty.py new file mode 100644 index 00000000..5966aaea --- /dev/null +++ b/tests/tests_automl/test_predict_uncertainty.py @@ -0,0 +1,124 @@ +import os +import shutil +import unittest +import numpy as np +import pandas as pd +import pytest +from sklearn import datasets +from supervised import AutoML +from supervised.exceptions import AutoMLException + +class AutoMLPredictUncertaintyTest(unittest.TestCase): + automl_dir = "AutoML_Predict_Uncertainty_Test" + + def tearDown(self): + shutil.rmtree(self.automl_dir, ignore_errors=True) + + def setUp(self): + shutil.rmtree(self.automl_dir, ignore_errors=True) + + def test_predict_uncertainty_before_fit(self): + """Should raise AutoMLException when called before fit""" + automl = AutoML(results_path=self.automl_dir) + X, y = datasets.make_regression(n_samples=50, n_features=5, random_state=42) + with self.assertRaises(AutoMLException) as context: + automl.predict_uncertainty(X) + self.assertIn("has not been fitted yet", str(context.exception)) + + def test_predict_uncertainty_classification(self): + """Should raise AutoMLException for classification tasks""" + X, y = datasets.make_classification(n_samples=50, n_features=5, n_classes=2, random_state=42) + automl = AutoML( + results_path=self.automl_dir, + algorithms=["Decision Tree"], + explain_level=0, + verbose=0 + ) + automl.fit(X, y) + with self.assertRaises(AutoMLException) as context: + automl.predict_uncertainty(X) + self.assertIn("only supported for regression tasks", str(context.exception)) + + def test_predict_uncertainty_no_ensemble(self): + """Should raise AutoMLException when the best model is not an Ensemble and no Ensemble was trained""" + X, y = datasets.make_regression(n_samples=50, n_features=5, random_state=42) + automl = AutoML( + results_path=self.automl_dir, + algorithms=["Decision Tree"], + train_ensemble=False, + explain_level=0, + verbose=0, + random_state=42 + ) + automl.fit(X, y) + + self.assertNotEqual(automl._best_model.get_type(), "Ensemble") + + with self.assertRaises(AutoMLException) as context: + automl.predict_uncertainty(X) + self.assertIn("Ensemble model is available", str(context.exception)) + + def test_predict_uncertainty_ensemble_fallback(self): + """Should fallback to Ensemble model if the best model is not an Ensemble but Ensemble was trained""" + X, y = datasets.make_regression(n_samples=50, n_features=5, random_state=42) + automl = AutoML( + results_path=self.automl_dir, + algorithms=["Linear", "Decision Tree"], + train_ensemble=True, + explain_level=0, + verbose=0, + random_state=42 + ) + automl.fit(X, y) + + for m in automl._models: + if m.get_type() == "Decision Tree": + automl._best_model = m + break + + self.assertEqual(automl._best_model.get_type(), "Decision Tree") + + res = automl.predict_uncertainty(X) + self.assertIsInstance(res, pd.DataFrame) + self.assertIn("prediction", res.columns) + + def test_predict_uncertainty_regression_success(self): + """Should successfully compute uncertainty for regression when ensemble is available""" + np.random.seed(42) + X = np.random.rand(120, 5) + y = X[:, 0]**2 + np.sin(X[:, 1] * np.pi) + np.exp(X[:, 2]) + np.random.normal(0, 0.05, 120) + X = pd.DataFrame(X, columns=[f"f_{i}" for i in range(5)]) + + automl = AutoML( + results_path=self.automl_dir, + algorithms=["Linear", "Decision Tree", "Random Forest"], + explain_level=0, + verbose=0, + random_state=42 + ) + automl.fit(X, y) + + self.assertEqual(automl._best_model.get_type(), "Ensemble") + + alpha = 0.05 + res = automl.predict_uncertainty(X, alpha=alpha) + + # Assertions + self.assertIsInstance(res, pd.DataFrame) + + expected_cols = ["prediction", "prediction_std", "prediction_variance", "lower", "upper"] + for col in expected_cols: + self.assertIn(col, res.columns) + + self.assertEqual(len(res), len(X)) + + # Check standard deviation and variance are non-negative + self.assertTrue((res["prediction_std"] >= 0).all()) + self.assertTrue((res["prediction_variance"] >= 0).all()) + + # Check lower <= prediction <= upper + self.assertTrue((res["lower"] <= res["prediction"]).all()) + self.assertTrue((res["prediction"] <= res["upper"]).all()) + + # Check standard deviation calculation is close to square root of variance + np.testing.assert_array_almost_equal(res["prediction_std"], np.sqrt(res["prediction_variance"]))