Create weighted data grids

DESCRIPTION

@@ -1,7 +1,7 @@
 Type: Package
 Package: modelbased
 Title: Estimation of Model-Based Predictions, Contrasts and Means
-Version: 0.16.0.13
+Version: 0.16.0.14
 Authors@R:
     c(person(given = "Dominique",
              family = "Makowski",

R/get_marginalmeans.R

@@ -155,6 +155,22 @@ get_marginalmeans <- function(
     fun_args$type <- predict_args$predict
   }
 
+  ## TODO: document "fast" argument?
+
+  # fast mode?
+  # ---------------------------
+
+  # create weighted (reduced) data grid, if "fast = TRUE"
+  tmp <- .create_weighted_datagrid(
+    model = model,
+    estimate = estimate,
+    fun_args = fun_args,
+    dots = dots
+  )
+  # if we made any changes, update arguments
+  fun_args <- tmp$fun_args
+  dots <- tmp$dots
+
   # weights?
   # ---------------------------
 
@@ -796,3 +812,47 @@ get_marginalmeans <- function(
     }
   }
 }
+
+
+# create weighted data grid, for faster computation of predictions. Works best
+# if reduced data grid is considerably smaller than the original model data,
+# and is more efficient if model contains fewer categorical predictors.
+
+.create_weighted_datagrid <- function(model, estimate, fun_args, dots) {
+  if (isTRUE(dots$fast) || is.numeric(dots$fast)) {
+    # this overrides the existing data grid, if any. this means, "fast = TRUE"
+    # only works with marginal predictions (when `estimate` is "average" or
+    # "population"), not conditional, datagrid-based predictions (when
+    # `estimate` is "specific" or "typical")
+    if (estimate %in% c("specific", "typical")) {
+      insight::format_error(
+        "`fast` only works for marginal predictions, i.e. when `estimate` is set to \"average\" or \"population\"."
+      )
+    }
+    # do we have bins?
+    n_bins <- dots$n_bins
+    # if not specified, we default to 5 bins, or to the value provided in "fast"
+    if (is.null(n_bins)) {
+      if (is.numeric(dots$fast)) {
+        n_bins <- dots$fast
+      } else {
+        n_bins <- 5
+      }
+    }
+    # does model have weights?
+    model_weights <- insight::find_weights(model)
+    # if not, we default to TRUE to trigger weighted data grid. Else, we
+    # pass the name of the weights variable
+    if (is.null(model_weights)) {
+      model_weights <- TRUE
+    }
+    # create weighted (reduced) data grid
+    fast_grid <- insight::get_datagrid(model, n_bins = n_bins, weighted = model_weights)
+    # update newdata with reduced data grid for faster computation
+    fun_args$newdata <- fast_grid
+    dots$weights <- fast_grid$Weight
+    # clean-up dots
+    dots$fast <- dots$n_bins <- NULL
+  }
+  list(fun_args = fun_args, dots = dots)
+}

tests/testthat/test-weighted_datagrid.R

@@ -2,7 +2,7 @@ skip_on_cran()
 skip_if_not_installed("marginaleffects")
 skip_if_not_installed("insight", minimum_version = "1.5.3")
 
-test_that("weighted data grids work", {
+test_that("weighted data grids work for average", {
   data(penguins)
 
   # one factor
@@ -17,5 +17,117 @@ test_that("weighted data grids work", {
     data = dg,
     weights = dg$Weight
   )
+  emm3 <- estimate_means(model, "species", estimate = "average", fast = TRUE)
+
+  expect_equal(emm1$Mean, emm2$Mean, tolerance = 1e-4)
+  expect_equal(emm1$Mean, emm3$Mean, tolerance = 1e-4)
+})
+
+
+test_that("weighted data grids work for average, binning of numerics", {
+  data(penguins)
+
+  # one factor
+  model <- lm(bill_len ~ species + sex + island + body_mass, data = penguins)
+  dg <- insight::get_datagrid(model, weighted = TRUE)
+
+  emm1 <- estimate_means(model, "species", estimate = "average")
+  emm2 <- estimate_means(
+    model,
+    "species",
+    estimate = "average",
+    data = dg,
+    weights = dg$Weight
+  )
+  emm3 <- estimate_means(model, "species", estimate = "average", fast = TRUE)
+
+  # need lower tolerance, due to binning not being exact to the empirical
+  # average of numeric values in the data
+  expect_equal(emm1$Mean, emm2$Mean, tolerance = 1e-3)
+  expect_equal(emm2$Mean, emm3$Mean, tolerance = 1e-5)
+
+  dg <- insight::get_datagrid(model, n_bins = 15, weighted = TRUE)
+  emm2 <- estimate_means(
+    model,
+    "species",
+    estimate = "average",
+    data = dg,
+    weights = dg$Weight
+  )
+  emm3 <- estimate_means(model, "species", estimate = "average", fast = 15)
+
+  # need lower tolerance, due to binning not being exact to the empirical
+  # average of numeric values in the data, but tolerance is stricter
+  # due to more precise binning
+  expect_equal(emm1$Mean, emm2$Mean, tolerance = 1e-4)
+  expect_equal(emm2$Mean, emm3$Mean, tolerance = 1e-5)
+})
+
+
+test_that("weighted data grids work for average and model weights, binning of numerics", {
+  set.seed(123)
+  d <- penguins
+  d$weights <- abs(rnorm(nrow(d), 1, 0.2))
+  model <- lm(body_mass ~ species + sex + bill_len, data = d, weights = weights)
+
+  dg <- insight::get_datagrid(model, weighted = "weights")
+
+  emm1 <- estimate_means(
+    model,
+    "species",
+    estimate = "average",
+    data = dg,
+    weights = dg$Weight
+  )
+  emm2 <- estimate_means(model, "species", estimate = "average", fast = TRUE)
+  expect_equal(emm1$Mean, emm2$Mean, tolerance = 1e-5)
+
+  dg <- insight::get_datagrid(model, weighted = "weights", n_bins = 10)
+
+  emm1 <- estimate_means(
+    model,
+    "species",
+    estimate = "average",
+    data = dg,
+    weights = dg$Weight
+  )
+  emm2 <- estimate_means(model, "species", estimate = "average", fast = 10)
+  expect_equal(emm1$Mean, emm2$Mean, tolerance = 1e-5)
+})
+
+
+test_that("weighted data grids work for population", {
+  data(penguins)
+
+  # one factor
+  model <- lm(bill_len ~ species + sex + island, data = penguins)
+  dg <- insight::get_datagrid(model, weighted = TRUE)
+
+  emm1 <- estimate_means(model, "species", estimate = "population")
+  emm2 <- estimate_means(
+    model,
+    "species",
+    estimate = "population",
+    data = dg,
+    weights = dg$Weight
+  )
+  emm3 <- estimate_means(model, "species", estimate = "population", fast = TRUE)
+
   expect_equal(emm1$Mean, emm2$Mean, tolerance = 1e-4)
+  expect_equal(emm1$Mean, emm3$Mean, tolerance = 1e-4)
+})
+
+
+test_that("weighted data grids errors for other estimate options", {
+  data(penguins)
+
+  # one factor
+  model <- lm(bill_len ~ species + sex + island, data = penguins)
+  dg <- insight::get_datagrid(model, weighted = TRUE)
+
+  expect_error(
+    estimate_means(model, "species", fast = TRUE),
+    regex = "`fast` only works for marginal",
+    fixed = TRUE
+  )
 })