Use shap_names instead

NorskRegnesentral · Dec 13, 2024 · ebcc1b8 · ebcc1b8
1 parent ebc00ed
commit ebcc1b8
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/R/print_iter.R b/R/print_iter.R
@@ -79,13 +79,9 @@ print_iter <- function(internal) {
   }
 
   if ("shapley" %in% verbose) {
-    if (internal$parameters$type == "forecast") {
-      dt_shapley_est <- internal$iter_list[[iter]]$dt_shapley_est[, -c("explain_idx", "horizon")]
-      dt_shapley_sd <- internal$iter_list[[iter]]$dt_shapley_sd[, -c("explain_idx", "horizon")]
-    } else {
-      dt_shapley_est <- internal$iter_list[[iter]]$dt_shapley_est[, -"explain_id"]
-      dt_shapley_sd <- internal$iter_list[[iter]]$dt_shapley_sd[, -"explain_id"]
-    }
+    shap_names_with_none <- c("none", internal$parameters$shap_names)
+    dt_shapley_est <- internal$iter_list[[iter]]$dt_shapley_est[, ..shap_names_with_none]
+    dt_shapley_sd <- internal$iter_list[[iter]]$dt_shapley_sd[, ..shap_names_with_none]
 
     # Printing the current Shapley values
     matrix1 <- format(round(dt_shapley_est, 3), nsmall = 2, justify = "right")