Regression Models Explainability (XAI) using SHAP (#2766)

* SHAP for regression models

* fixed Xai

* fix colours of axis

* raw explain DF

* naming convention

* removed printing. exports to csv now

* i18n update

* re-add back in. need to set new fix

Co-authored-by: James Maslek <jmaslek11@gmail.com>
Co-authored-by: Colin Delahunty <72827203+colin99d@users.noreply.github.com>

i18n/en.yml

@@ -1057,8 +1057,8 @@ en:
   forecast/arima: Arima (Non-darts)
   forecast/expo: Probabilistic Exponential Smoothing
   forecast/theta: Theta Method
-  forecast/linregr: Probabilistic Linear Regression (feat. Past covariates)
-  forecast/regr: Regression (feat. Past covariates)
+  forecast/linregr: Probabilistic Linear Regression (feat. Past covariates and Explainability)
+  forecast/regr: Regression (feat. Past covariates and Explainability)
   forecast/rnn: Probabilistic Recurrent Neural Network (RNN, LSTM, GRU)
   forecast/brnn: Block Recurrent Neural Network (RNN, LSTM, GRU) (feat. Past covariates)
   forecast/nbeats: Neural Bayesian Estimation (feat. Past covariates)

openbb_terminal/forecast/forecast_controller.py

@@ -381,6 +381,7 @@ class ForecastController(BaseController):
         residuals: bool = False,
         forecast_only: bool = False,
         naive: bool = False,
+        explainability_raw: bool = False,
     ):
         if hidden_size:
             parser.add_argument(
@@ -638,6 +639,15 @@ class ForecastController(BaseController):
                 default=False,
                 dest="forecast_only",
             )
+        if explainability_raw:
+            parser.add_argument(
+                "--explainability-raw",
+                action="store_true",
+                dest="explainability_raw",
+                default=False,
+                help="Prints out a raw dataframe showing explainability results.",
+            )
+
             # if user does not put in --target-dataset
         return super().parse_known_args_and_warn(
             parser, other_args, export_allowed, raw, limit
@@ -2047,6 +2057,7 @@ class ForecastController(BaseController):
             start=True,
             end=True,
             naive=True,
+            explainability_raw=True,
         )
 
         if ns_parser:
@@ -2074,6 +2085,7 @@ class ForecastController(BaseController):
                 start_date=ns_parser.s_start_date,
                 end_date=ns_parser.s_end_date,
                 naive=ns_parser.naive,
+                explainability_raw=ns_parser.explainability_raw,
             )
 
     @log_start_end(log=logger)
@@ -2108,6 +2120,7 @@ class ForecastController(BaseController):
             start=True,
             end=True,
             naive=True,
+            explainability_raw=True,
         )
 
         if ns_parser:
@@ -2134,6 +2147,7 @@ class ForecastController(BaseController):
                 start_date=ns_parser.s_start_date,
                 end_date=ns_parser.s_end_date,
                 naive=ns_parser.naive,
+                explainability_raw=ns_parser.explainability_raw,
             )
 
     @log_start_end(log=logger)

openbb_terminal/forecast/helpers.py

@@ -18,6 +18,7 @@ from darts.dataprocessing.transformers import MissingValuesFiller, Scaler
 from darts.utils.statistics import plot_residuals_analysis
 from darts import TimeSeries
 from darts.metrics import mape
+from darts.explainability.shap_explainer import ShapExplainer
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
 from openbb_terminal.rich_config import console
 from openbb_terminal.config_terminal import theme
@@ -529,9 +530,52 @@ def plot_forecast(
 
     print_pretty_prediction(numeric_forecast, data[target_col].iloc[-1])
 
+    # TODO: This needs to get fixed
     export_data(export, os.path.dirname(os.path.abspath(__file__)), "expo")
 
 
+def plot_explainability(
+    model, explainability_raw=False, external_axes: Optional[List[plt.axes]] = None
+):
+    """Use SHAP to explain the model's predictions.
+    Args:
+        model (Linregr or Regr): Trained model
+    """
+    if not external_axes:
+        _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI)
+    else:
+        ax = external_axes[0]
+
+    shap_explain = ShapExplainer(model)
+    shap_explain.summary_plot(horizons=1)
+    if explainability_raw:
+        console.print("")
+        console.print("[green]Exporting Raw Explainability DataFrame[/green]")
+        raw_df = shap_explain.explain().get_explanation(horizon=1).pd_dataframe()
+        export_data(
+            "csv",
+            os.path.dirname(os.path.abspath(__file__)),
+            "explainability_raw",
+            raw_df,
+        )
+
+    ax.yaxis.set_label_position("left")
+    ax.yaxis.tick_left()
+
+    # change the colour of the y axis tick labels
+    for t in ax.get_yticklabels():
+        t.set_color("white")
+
+    # change the colour of the x axis tick labels
+    for t in ax.get_xticklabels():
+        t.set_color("white")
+
+    theme.style_primary_axis(ax)
+
+    if not external_axes:
+        theme.visualize_output()
+
+
 def dt_format(x) -> str:
     """Convert any Timestamp to YYYY-MM-DD
     Args:

openbb_terminal/forecast/linregr_view.py

@@ -33,6 +33,7 @@ def display_linear_regression(
     start_date: Optional[datetime] = None,
     end_date: Optional[datetime] = None,
     naive: bool = False,
+    explainability_raw: bool = False,
     external_axes: Optional[List[plt.axes]] = None,
 ):
     """Display Linear Regression Forecasting
@@ -120,3 +121,6 @@ def display_linear_regression(
         helpers.plot_residuals(
             _model, past_covariates, ticker_series, forecast_horizon=forecast_horizon
         )
+
+    # SHAP
+    helpers.plot_explainability(_model, explainability_raw)

openbb_terminal/forecast/regr_view.py

@@ -33,6 +33,7 @@ def display_regression(
     start_date: Optional[datetime] = None,
     end_date: Optional[datetime] = None,
     naive: bool = False,
+    explainability_raw: bool = False,
     external_axes: Optional[List[plt.axes]] = None,
 ):
     """Display Regression Forecasting
@@ -119,3 +120,6 @@ def display_regression(
         helpers.plot_residuals(
             _model, past_covariates, ticker_series, forecast_horizon=forecast_horizon
         )
+
+    # SHAP
+    helpers.plot_explainability(_model, explainability_raw)