update shap unit tests

Author: dennisbader

darts/tests/explainability/test_shap_explainer.py

@@ -16,13 +16,19 @@
 from darts.explainability.explainability_result import ShapExplainabilityResult
 from darts.explainability.shap_explainer import ShapExplainer
 from darts.models import (
+    CatBoostModel,
     ExponentialSmoothing,
+    LightGBMModel,
     LinearRegressionModel,
+    NotImportedModule,
     RegressionModel,
     XGBModel,
 )
 from darts.tests.base_test_class import DartsBaseTestClass
 
+lgbm_available = not isinstance(LightGBMModel, NotImportedModule)
+cb_available = not isinstance(CatBoostModel, NotImportedModule)
+
 
 class ShapExplainerTestCase(DartsBaseTestClass):
     np.random.seed(42)
@@ -128,9 +134,9 @@ class ShapExplainerTestCase(DartsBaseTestClass):
     )
 
     def test_creation(self):
-
+        model_cls = LightGBMModel if lgbm_available else XGBModel
         # Model should be fitted first
-        m = XGBModel(
+        m = model_cls(
             lags=4,
             lags_past_covariates=[-1, -2, -3],
             lags_future_covariates=[0],
@@ -157,7 +163,7 @@ def test_creation(self):
                 self.target_ts,
             )
 
-        m = XGBModel(
+        m = model_cls(
             lags=4,
             lags_past_covariates=[-1, -2, -3],
             lags_future_covariates=[0],
@@ -248,7 +254,8 @@ def test_creation(self):
         )
 
         # CatBoost
-        m = XGBModel(
+        model_cls = CatBoostModel if cb_available else XGBModel
+        m = model_cls(
             lags=4,
             lags_past_covariates=[-1, -2, -6],
             lags_future_covariates=[0],
@@ -269,7 +276,8 @@ def test_creation(self):
             ShapExplainer(m, shap_method="bad_choice")
 
     def test_explain(self):
-        m = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        m = model_cls(
             lags=4,
             lags_past_covariates=[-1, -2, -3],
             lags_future_covariates=[0],
@@ -428,7 +436,8 @@ def test_explain(self):
         self.assertTrue(isinstance(shap_explain.explain(), ShapExplainabilityResult))
 
     def test_explain_with_lags_future_covariates_series_of_same_length_as_target(self):
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             lags_past_covariates=[-1, -2, -3],
             lags_future_covariates=[2],
@@ -464,7 +473,8 @@ def test_explain_with_lags_future_covariates_series_extending_into_future(self):
         fut_cov = np.random.normal(0, 1, len(days)).astype("float32")
         fut_cov_ts = TimeSeries.from_times_and_values(days, fut_cov.reshape(-1, 1))
 
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             lags_past_covariates=[-1, -2, -3],
             lags_future_covariates=[2],
@@ -499,7 +509,8 @@ def test_explain_with_lags_covariates_series_older_timestamps_than_target(self):
         past_cov = np.random.normal(0, 1, len(days)).astype("float32")
         past_cov_ts = TimeSeries.from_times_and_values(days, past_cov.reshape(-1, 1))
 
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=None,
             lags_past_covariates=[-1, -2],
             lags_future_covariates=[-1, -2],
@@ -525,7 +536,8 @@ def test_explain_with_lags_covariates_series_older_timestamps_than_target(self):
             self.assertEqual(explanation.start_time(), self.target_ts.start_time())
 
     def test_plot(self):
-        m_0 = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        m_0 = model_cls(
             lags=4,
             lags_past_covariates=[-1, -2, -3],
             lags_future_covariates=[0],
@@ -618,7 +630,8 @@ def test_plot(self):
         plt.close()
 
     def test_feature_values_align_with_input(self):
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             output_chunk_length=1,
         )
@@ -644,7 +657,8 @@ def test_feature_values_align_with_input(self):
         )
 
     def test_feature_values_align_with_raw_output_shap(self):
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             output_chunk_length=1,
         )
@@ -670,7 +684,8 @@ def test_feature_values_align_with_raw_output_shap(self):
         ), "The shape of the feature values should be the same as the shap values"
 
     def test_shap_explanation_object_validity(self):
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             lags_past_covariates=2,
             lags_future_covariates=[1],
@@ -692,7 +707,8 @@ def test_shap_explanation_object_validity(self):
         )
 
     def test_shap_selected_components(self):
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             lags_past_covariates=2,
             lags_future_covariates=[1],
@@ -728,7 +744,8 @@ def test_shap_selected_components(self):
 
     def test_shapley_with_static_cov(self):
         ts = self.target_ts_with_static_covs
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             output_chunk_length=1,
         )
@@ -772,7 +789,8 @@ def test_shapley_with_static_cov(self):
             ]
 
     def test_shapley_multiple_series_with_different_static_covs(self):
-        model = XGBModel(
+        model_cls = LightGBMModel if lgbm_available else XGBModel
+        model = model_cls(
             lags=4,
             output_chunk_length=1,
         )