Merge pull request #166 from ModelOriented/avoid-replication-of-z

Avoid replication of x

Author: mayer79

NEWS.md

@@ -1,8 +1,15 @@
 # kernelshap 0.8.1
 
-### Performance improvements
+### API
 
-- `permshap(, exact = TRUE)` is slightly faster by pre-calculating more 
+- The argument `feature_names` can now also be used with matrix input ([#166](https://github.com/ModelOriented/kernelshap/pull/166)).
+
+### Speed and memory improvements
+
+- `permshap()` and `kernelshap()` require about 10% less memory ([#166](https://github.com/ModelOriented/kernelshap/pull/166)).
+- `permshap()` and `kernelshap()` are faster for data.frame input, 
+  and slightly slower for matrix input ([#166](https://github.com/ModelOriented/kernelshap/pull/166)).
+- Additionally, `permshap(, exact = TRUE)` is faster by pre-calculating more 
   elements used across rows [#165](https://github.com/ModelOriented/kernelshap/pull/165)
 
 ### Documentation

R/kernelshap.R

@@ -82,8 +82,7 @@
 #'   Additional (named) arguments are passed via `...`.
 #'   The default, [stats::predict()], will work in most cases.
 #' @param feature_names Optional vector of column names in `X` used to calculate
-#'   SHAP values. By default, this equals `colnames(X)`. Not supported if `X`
-#'   is a matrix.
+#'   SHAP values. By default, this equals `colnames(X)`.
 #' @param bg_w Optional vector of case weights for each row of `bg_X`.
 #'   If `bg_X = NULL`, must be of same length as `X`. Set to `NULL` for no weights.
 #' @param bg_n If `bg_X = NULL`: Size of background data to be sampled from `X`.
@@ -240,12 +239,6 @@ kernelshap.default <- function(
     message(txt)
   }
 
-  # Drop unnecessary columns in bg_X. If X is matrix, also column order is relevant
-  # In what follows, predictions will never be applied directly to bg_X anymore
-  if (!identical(colnames(bg_X), feature_names)) {
-    bg_X <- bg_X[, feature_names, drop = FALSE]
-  }
-
   # Pre-calculations that are identical for each row to be explained
   if (exact || hybrid_degree >= 1L) {
     if (exact) {

R/permshap.R

@@ -137,15 +137,9 @@ permshap.default <- function(
   v0 <- wcolMeans(bg_preds, w = bg_w) # Average pred of bg data: 1 x K
   v1 <- align_pred(pred_fun(object, X, ...)) # Predictions on X:        n x K
 
-  # Drop unnecessary columns in bg_X. If X is matrix, also column order is relevant
-  # Predictions will never be applied directly to bg_X anymore
-  if (!identical(colnames(bg_X), feature_names)) {
-    bg_X <- bg_X[, feature_names, drop = FALSE]
-  }
-
   # Pre-calculations that are identical for each row to be explained
   if (exact) {
-    Z <- exact_Z(p, feature_names = feature_names, keep_extremes = TRUE)
+    Z <- exact_Z(p, feature_names = feature_names)
     m_exact <- nrow(Z) - 2L # We won't evaluate vz for first and last row
     m_eval <- 0L # for consistency with sampling case
     precalc <- list(

R/utils.R

@@ -49,61 +49,56 @@ wcolMeans <- function(x, w = NULL, ...) {
 #'
 #' @param p Number of features.
 #' @param feature_names Feature names.
-#' @param keep_extremes Should extremes be kept? Defaults to `FALSE` (for kernelshap).
 #' @returns An integer matrix of all on-off vectors of length `p`.
-exact_Z <- function(p, feature_names, keep_extremes = FALSE) {
+exact_Z <- function(p, feature_names) {
   if (p < 2L) {
     stop("p must be at least 2 if exact = TRUE.")
   }
   m <- 2^p
-  if (keep_extremes) {
-    M <- seq_len(m) - 1L
-  } else {
-    M <- seq_len(m - 2L)
-    m <- m - 2L
-  }
+  M <- seq_len(m) - 1L
   encoded <- as.integer(intToBits(M))
   dim(encoded) <- c(32L, m)
   Z <- t(encoded[p:1L, , drop = FALSE])
   colnames(Z) <- feature_names
   return(Z)
 }
 
-#' Masker
+#' Masked Predict
 #'
 #' Internal function.
 #' For each on-off vector (rows in `Z`), the (weighted) average prediction is returned.
-#' In Python, this function is called "masker", and Z is the "mask".
 #'
 #' @noRd
 #' @keywords internal
 #'
 #' @inheritParams kernelshap
-#' @param X Row to be explained stacked m*n_bg times.
+#' @param x Row to be explained.
 #' @param bg Background data stacked m times.
-#' @param Z A (m x p) matrix with on-off values.
+#' @param Z A (m x p) matrix with on-off values (logical or integer).
 #' @param w A vector with case weights (of the same length as the unstacked
 #'   background data).
 #' @returns A (m x K) matrix with vz values.
-get_vz <- function(X, bg, Z, object, pred_fun, w, ...) {
+get_vz <- function(x, bg, Z, object, pred_fun, w, ...) {
   m <- nrow(Z)
-  not_Z <- !Z
+  if (!is.logical(Z)) {
+    storage.mode(Z) <- "logical"
+  }
   n_bg <- nrow(bg) / m # because bg was replicated m times
 
-  # Replicate not_Z, so that X, bg, not_Z are all of dimension (m*n_bg x p)
+  # Replicate Z, so that bg and Z are of dimension (m*n_bg x p)
   g <- rep_each(m, each = n_bg)
-  not_Z <- not_Z[g, , drop = FALSE]
+  Z_rep <- Z[g, , drop = FALSE]
 
-  if (is.matrix(X)) {
-    # Remember that columns of X and bg are perfectly aligned in this case
-    X[not_Z] <- bg[not_Z]
-  } else {
-    for (v in colnames(Z)) {
-      s <- not_Z[, v]
-      X[[v]][s] <- bg[[v]][s]
+  for (v in colnames(Z)) {
+    s <- Z_rep[, v]
+    if (is.matrix(x)) {
+      bg[s, v] <- x[, v]
+    } else {
+      bg[[v]][s] <- x[[v]]
     }
   }
-  preds <- align_pred(pred_fun(object, X, ...))
+
+  preds <- align_pred(pred_fun(object, bg, ...))
 
   # Aggregate (distinguishing fast 1-dim case)
   if (ncol(preds) == 1L) {
@@ -337,8 +332,6 @@ basic_checks <- function(X, feature_names, pred_fun) {
     dim(X) >= 1L,
     length(feature_names) >= 1L,
     all(feature_names %in% colnames(X)),
-    "If X is a matrix, feature_names must equal colnames(X)" =
-      !is.matrix(X) || identical(colnames(X), feature_names),
     is.function(pred_fun)
   )
   TRUE

R/utils_kernelshap.R

@@ -1,8 +1,21 @@
 # Kernel SHAP algorithm for a single row x
 # If exact, a single call to predict() is necessary.
 # If sampling is involved, we need at least two additional calls to predict().
-kernelshap_one <- function(x, v1, object, pred_fun, feature_names, bg_w, exact, deg,
-                           m, tol, max_iter, v0, precalc, ...) {
+kernelshap_one <- function(
+    x,
+    v1,
+    object,
+    pred_fun,
+    feature_names,
+    bg_w,
+    exact,
+    deg,
+    m,
+    tol,
+    max_iter,
+    v0,
+    precalc,
+    ...) {
   p <- length(feature_names)
   K <- ncol(v1)
   K_names <- colnames(v1)
@@ -16,14 +29,8 @@ kernelshap_one <- function(x, v1, object, pred_fun, feature_names, bg_w, exact,
     v0_m_exact <- v0[rep.int(1L, m_exact), , drop = FALSE] #  (m_ex x K)
 
     # Most expensive part
-    vz <- get_vz( #  (m_ex x K)
-      X = rep_rows(x, rep.int(1L, nrow(bg_X_exact))), #  (m_ex*n_bg x p)
-      bg = bg_X_exact, #  (m_ex*n_bg x p)
-      Z = Z, #  (m_ex x p)
-      object = object,
-      pred_fun = pred_fun,
-      w = bg_w,
-      ...
+    vz <- get_vz(
+      x = x, bg = bg_X_exact, Z = Z, object = object, pred_fun = pred_fun, w = bg_w, ...
     )
     # Note: w is correctly replicated along columns of (vz - v0_m_exact)
     b_exact <- crossprod(Z, precalc[["w"]] * (vz - v0_m_exact)) #  (p x K)
@@ -37,7 +44,6 @@ kernelshap_one <- function(x, v1, object, pred_fun, feature_names, bg_w, exact,
 
   # Iterative sampling part, always using A_exact and b_exact to fill up the weights
   bg_X_m <- precalc[["bg_X_m"]] #  (m*n_bg x p)
-  X <- rep_rows(x, rep.int(1L, nrow(bg_X_m))) #  (m*n_bg x p)
   v0_m <- v0[rep.int(1L, m), , drop = FALSE] #  (m x K)
   est_m <- array(
     data = 0, dim = c(max_iter, p, K), dimnames = list(NULL, feature_names, K_names)
@@ -60,7 +66,7 @@ kernelshap_one <- function(x, v1, object, pred_fun, feature_names, bg_w, exact,
 
     # Expensive                                                              #  (m x K)
     vz <- get_vz(
-      X = X, bg = bg_X_m, Z = Z, object = object, pred_fun = pred_fun, w = bg_w, ...
+      x = x, bg = bg_X_m, Z = Z, object = object, pred_fun = pred_fun, w = bg_w, ...
     )
 
     # The sum of weights of A_exact and input[["A"]] is 1, same for b
@@ -151,7 +157,8 @@ input_sampling <- function(p, m, deg, feature_names) {
 #      the SHAP kernel distribution
 # - A: Exact matrix A = Z'wZ
 input_exact <- function(p, feature_names) {
-  Z <- exact_Z(p, feature_names = feature_names, keep_extremes = FALSE)
+  Z <- exact_Z(p, feature_names = feature_names)
+  Z <- Z[2L:(nrow(Z) - 1L), , drop = FALSE]
   # Each Kernel weight(j) is divided by the number of vectors z having sum(z) = j
   w <- kernel_weights(p) / choose(p, 1:(p - 1L))
   list(Z = Z, w = w[rowSums(Z)], A = exact_A(p, feature_names = feature_names))

R/utils_permshap.R

@@ -28,7 +28,6 @@ permshap_one <- function(
     max_iter,
     ...) {
   bg <- precalc$bg_X_rep
-  X <- rep_rows(x, rep.int(1L, times = nrow(bg)))
 
   p <- length(feature_names)
   K <- ncol(v1)
@@ -38,7 +37,7 @@ permshap_one <- function(
   if (exact) {
     Z <- precalc$Z # ((m_ex+2) x K)
     vz <- get_vz( # (m_ex x K)
-      X = X,
+      x = x,
       bg = bg,
       Z = Z[2L:(nrow(Z) - 1L), , drop = FALSE], # (m_ex x p)
       object = object,
@@ -52,7 +51,7 @@ permshap_one <- function(
       pos <- precalc$positions[[j]]
       beta_n[j, ] <- wcolMeans(
         vz[pos$on, , drop = FALSE] - vz[pos$off, , drop = FALSE],
-        weights = precalc["shapley_w"][pos$on]
+        w = precalc$shapley_w[pos$on]
       )
     }
     return(list(beta = beta_n))
@@ -68,7 +67,7 @@ permshap_one <- function(
 
   # Pre-calculate part of Z with rowsum 1 or p - 1
   vz_balanced <- get_vz( # (2p x K)
-    X = rep_rows(x, rep.int(1L, times = nrow(precalc$bg_X_balanced))),
+    x = x,
     bg = precalc$bg_X_balanced,
     Z = precalc$Z_balanced,
     object = object,
@@ -90,13 +89,19 @@ permshap_one <- function(
     if (!low_memory) { # predictions for all chains at once
       Z <- do.call(rbind, Z)
       vz <- get_vz(
-        X = X, bg = bg, Z = Z, object = object, pred_fun = pred_fun, w = bg_w, ...
+        x = x, bg = bg, Z = Z, object = object, pred_fun = pred_fun, w = bg_w, ...
       )
     } else { # predictions for each chain separately
       vz <- vector("list", length = p)
       for (j in seq_len(p)) {
         vz[[j]] <- get_vz(
-          X = X, bg = bg, Z = Z[[j]], object = object, pred_fun = pred_fun, w = bg_w, ...
+          x = x,
+          bg = bg,
+          Z = Z[[j]],
+          object = object,
+          pred_fun = pred_fun,
+          w = bg_w,
+          ...
         )
       }
       vz <- do.call(rbind, vz)
@@ -142,20 +147,6 @@ shapley_weights <- function(p, ell) {
   1 / choose(p, ell) / (p - ell)
 }
 
-#' Rowwise Paste
-#'
-#' Fast version of `apply(Z, 1L, FUN = paste0, collapse = "")`.
-#' Required for exact permutation SHAP only.
-#'
-#' @noRd
-#' @keywords internal
-#'
-#' @param Z A (n x p) matrix.
-#' @returns A length n vector.
-rowpaste <- function(Z) {
-  do.call(paste0, asplit(Z, 2L))
-}
-
 #' Balanced Chains
 #'
 #' Creates `p` permutations of the numbers `1:p` such that each number appears once
@@ -249,7 +240,7 @@ init_vzj <- function(p, v0, v1) {
 #' @noRd
 #' @keywords internal
 #'
-#' @param mask The output of `exact_Z(p, feature_names, keep_extremes = TRUE)`.
+#' @param mask The output of `exact_Z(p, feature_names)`.
 #' @returns List with p elements, each containing an `on` and `off` vector.
 positions_for_exact <- function(mask) {
   p <- ncol(mask)

backlog/compare_with_python.R

@@ -31,6 +31,18 @@ system.time(
 )
 ks2
 
+bench::mark(kernelshap(fit, X_small, bg_X = bg_X, verbose=F))
+# 2.17s 1.64GB -> 1.72s 1.44GB
+
+bench::mark(kernelshap(fit, X_small, bg_X = bg_X, verbose=F, exact=F, hybrid_degree = 1))
+# 4.88s  2.79GB ->  4.38s 2.48GB
+
+bench::mark(permshap(fit, X_small, bg_X = bg_X, verbose=F))
+# 1.97s 1.64GB -> 1.8s 1.43GB
+
+bench::mark(permshap(fit, X_small, bg_X = bg_X, verbose=F, exact=F))
+# 3.04s  1.88GB ->  2.9s 1.64GB
+
 # SHAP values of first 2 observations:
 #          carat     clarity     color        cut
 # [1,] -2.050074 -0.28048747 0.1281222 0.01587382
@@ -54,7 +66,7 @@ fit <- lm(
 X_small <- diamonds[seq(1, nrow(diamonds), 53), setdiff(names(diamonds), "price")]
 
 # Exact KernelSHAP on X_small, using X_small as background data
-# (39s for exact, 15s for hybrid deg 2, 8s for hybrid deg 1, 16s for sampling)
+# (38s for exact, 11s for hybrid deg 2, 7s for hybrid deg 1, 11s for sampling)
 system.time(
   ks <- kernelshap(fit, X_small, bg_X = bg_X)
 )

man/kernelshap.Rd

@@ -175,12 +175,6 @@ test_that("Matrix input is fine", {
     expect_no_error( # additional cols in bg are ok
       algo(fit, X[J, x], pred_fun = pred_fun, bg_X = cbind(d = 1, X), verbose = FALSE)
     )
-    expect_error( # feature_names are less flexible
-      algo(fit, X[J, ],
-        pred_fun = pred_fun, bg_X = X,
-        verbose = FALSE, feature_names = "Sepal.Width"
-      )
-    )
   }
 })