feat: add predict_uncertainty() for ensemble-based regression intervals

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)

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
 

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"]))