From f00fa87e75509ef0bd1add3c081cabdddf228546 Mon Sep 17 00:00:00 2001
From: Ben Tupper <btupper@bigelow.org>
Date: Thu, 27 Feb 2025 13:15:54 -0500
Subject: [PATCH] add xgb.Booster methods to feature_effects() and
 partial_dependence()`

---
 NAMESPACE                 |  2 ++
 R/feature_effects.R       | 47 +++++++++++++++++++++++++++++++++++++++
 R/partial_dependence.R    | 40 +++++++++++++++++++++++++++++++++
 man/feature_effects.Rd    | 24 ++++++++++++++++++++
 man/partial_dependence.Rd | 20 +++++++++++++++++
 5 files changed, 133 insertions(+)

diff --git a/NAMESPACE b/NAMESPACE
index 62b59b1..8d5ffe9 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -11,10 +11,12 @@ S3method(feature_effects,H2OModel)
 S3method(feature_effects,default)
 S3method(feature_effects,explainer)
 S3method(feature_effects,ranger)
+S3method(feature_effects,xgb.Booster)
 S3method(partial_dependence,H2OModel)
 S3method(partial_dependence,default)
 S3method(partial_dependence,explainer)
 S3method(partial_dependence,ranger)
+S3method(partial_dependence,xgb.Booster)
 S3method(plot,EffectData)
 S3method(print,EffectData)
 S3method(update,EffectData)
diff --git a/R/feature_effects.R b/R/feature_effects.R
index b65275a..06bb2b1 100644
--- a/R/feature_effects.R
+++ b/R/feature_effects.R
@@ -473,6 +473,53 @@ feature_effects.H2OModel <- function(
   )
 }
 
+#' @describeIn feature_effects Method for xgb.Booster models.
+#' @export
+feature_effects.xgb.Booster <- function(
+    object,
+    v,
+    data,
+    y = NULL,
+    pred = NULL,
+    pred_fun = stats::predict,
+    trafo = NULL,
+    which_pred = NULL,
+    w = NULL,
+    breaks = "Sturges",
+    right = TRUE,
+    discrete_m = 13L,
+    outlier_iqr = 2,
+    calc_pred = TRUE,
+    pd_n = 500L,
+    ale_n = 50000L,
+    ale_bin_size = 200L,
+    ...
+) {
+  if (!inherits(data, "matrix")) {
+    data = as.matrix(data)
+  }
+  feature_effects.default(
+    object,
+    v = v,
+    data = data,
+    y = y,
+    pred = pred,
+    pred_fun = pred_fun,
+    trafo = trafo,
+    which_pred = which_pred,
+    w = w,
+    breaks = breaks,
+    right = right,
+    discrete_m = discrete_m,
+    outlier_iqr = outlier_iqr,
+    calc_pred = calc_pred,
+    pd_n = pd_n,
+    ale_n = ale_n,
+    ale_bin_size = ale_bin_size,
+    ...
+  )
+}
+
 #' Workhorse of feature_effects()
 #'
 #' Internal function used to calculate the output of `feature_effects()` for one
diff --git a/R/partial_dependence.R b/R/partial_dependence.R
index beae4cb..48f6a38 100644
--- a/R/partial_dependence.R
+++ b/R/partial_dependence.R
@@ -204,6 +204,46 @@ partial_dependence.H2OModel <- function(
   )
 }
 
+#' @describeIn partial_dependence Method for xgb.Booster models.
+#' @export
+partial_dependence.xgb.Booster <- function(
+    object,
+    v,
+    data,
+    pred_fun = stats::predict,
+    trafo = NULL,
+    which_pred = NULL,
+    w = NULL,
+    breaks = "Sturges",
+    right = TRUE,
+    discrete_m = 13L,
+    outlier_iqr = 2,
+    pd_n = 500L,
+    seed = NULL,
+    ...
+) {
+
+  if (!inherits(data, "matrix")) {
+    data = as.matrix(data)
+  }
+  partial_dependence.default(
+    object = object,
+    v = v,
+    data = data,
+    pred_fun = pred_fun,
+    trafo = trafo,
+    which_pred = which_pred,
+    w = w,
+    breaks = breaks,
+    right = right,
+    discrete_m = discrete_m,
+    outlier_iqr = outlier_iqr,
+    pd_n = pd_n,
+    seed = seed,
+    ...
+  )
+}
+
 #' Barebone Partial Dependence
 #'
 #' This is a barebone implementation of Friedman's partial dependence
diff --git a/man/feature_effects.Rd b/man/feature_effects.Rd
index b6f0a8f..a68a3cf 100644
--- a/man/feature_effects.Rd
+++ b/man/feature_effects.Rd
@@ -6,6 +6,7 @@
 \alias{feature_effects.ranger}
 \alias{feature_effects.explainer}
 \alias{feature_effects.H2OModel}
+\alias{feature_effects.xgb.Booster}
 \title{Feature Effects}
 \usage{
 feature_effects(object, ...)
@@ -94,6 +95,27 @@ feature_effects(object, ...)
   ale_bin_size = 200L,
   ...
 )
+
+\method{feature_effects}{xgb.Booster}(
+  object,
+  v,
+  data,
+  y = NULL,
+  pred = NULL,
+  pred_fun = stats::predict,
+  trafo = NULL,
+  which_pred = NULL,
+  w = NULL,
+  breaks = "Sturges",
+  right = TRUE,
+  discrete_m = 13L,
+  outlier_iqr = 2,
+  calc_pred = TRUE,
+  pd_n = 500L,
+  ale_n = 50000L,
+  ale_bin_size = 200L,
+  ...
+)
 }
 \arguments{
 \item{object}{Fitted model.}
@@ -232,6 +254,8 @@ If you need only one specific statistic, you can use the simplified APIs of
 
 \item \code{feature_effects(H2OModel)}: Method for H2O models.
 
+\item \code{feature_effects(xgb.Booster)}: Method for xgb.Booster models.
+
 }}
 \examples{
 fit <- lm(Sepal.Length ~ ., data = iris)
diff --git a/man/partial_dependence.Rd b/man/partial_dependence.Rd
index 6495d3e..4b9f122 100644
--- a/man/partial_dependence.Rd
+++ b/man/partial_dependence.Rd
@@ -6,6 +6,7 @@
 \alias{partial_dependence.ranger}
 \alias{partial_dependence.explainer}
 \alias{partial_dependence.H2OModel}
+\alias{partial_dependence.xgb.Booster}
 \title{Partial Dependence}
 \usage{
 partial_dependence(object, ...)
@@ -77,6 +78,23 @@ partial_dependence(object, ...)
   seed = NULL,
   ...
 )
+
+\method{partial_dependence}{xgb.Booster}(
+  object,
+  v,
+  data,
+  pred_fun = stats::predict,
+  trafo = NULL,
+  which_pred = NULL,
+  w = NULL,
+  breaks = "Sturges",
+  right = TRUE,
+  discrete_m = 13L,
+  outlier_iqr = 2,
+  pd_n = 500L,
+  seed = NULL,
+  ...
+)
 }
 \arguments{
 \item{object}{Fitted model.}
@@ -175,6 +193,8 @@ and eventually weighted by \code{w}.
 
 \item \code{partial_dependence(H2OModel)}: Method for H2O models.
 
+\item \code{partial_dependence(xgb.Booster)}: Method for xgb.Booster models.
+
 }}
 \examples{
 fit <- lm(Sepal.Length ~ ., data = iris)