Create weighted data grids

[strengejacke]

1. @strengejacke Do you think we could implement a way in modelbased to automatize the weighted approach? I.e., compute a weighted datagrid upfront

Originally posted by @DominiqueMakowski in #639

[DominiqueMakowski]

Idea:

1. Create a "fast" mode
2. First, stresstest by making it available only if all predictors are categorical (in the future, we could imagine binning continuous variables to approximate results, it could still be useful for massive datasets)
3. This fast mode creates a weighted datagrid (perhaps "weighted" is a feature to be added to get_datagrid, which would add a weight column)
4. These weights are then used by marginaleffects
5. Question: I reckon this only applies to some type of marginalization approaches

[strengejacke]

• For estimate = "specific" (similar to estimate_relation()) and estimate = "typical" (default), by creates a data grid that is passed to newdata (i.e., filtering works "by design").
• For estimate = "average", by is used for the by argument in marginaleffects, where the full data is used. Possible filtering in by requires that the "filter values/levels" are present in the data. newdata is not used.
• For estimate = "population", by is used to create a data grid, and the unique values for each variable in that data grid are passed to the variables argument in marginaleffects (i.e., filtering works "by design"). newdata is not used.

However: you can always specify newdata (or our alias data). If explicitly set, it will then be used as newdata way for all estimate options, meaning that the "fast" option should always work (unless marginaleffects handles newdata differently than expected).

[strengejacke]

Do we want to explicitly have an option/argument to internally trigger a weighted data grid, or is it sufficient to have it in insight::get_datagrid() only? Although we can simply pass the weighted argument to insight::get_datagrid() and document it in modelbased.

[DominiqueMakowski]

Do we want to explicitly have an option/argument to internally trigger a weighted data grid

You mean from like estimate_means/contrasts/slopes? Perhaps a fast argument could be quite explicit - especially if we verify that it indeed leads to speedups. Alternatively, exposing a weighted argument should work too, and we can document it by saying that it can be used for potential memory/efficiency benefits

[strengejacke]

This is not large-scale, but still a speed improvement:

library(glmmTMB)
library(modelbased)

m <- glmmTMB(
  count ~ mined + (1 | site),
  zi = ~mined,
  family = poisson,
  data = Salamanders
)

weighted_dg <- insight::get_datagrid(m, include_random = TRUE, weighted = TRUE)

estimate_means(m, "mined", estimate = "average")
#> Average Predictions
#> 
#> mined | Mean |   SE |       95% CI |     z
#> ------------------------------------------
#> yes   | 0.27 | 0.05 | [0.17, 0.37] |  5.47
#> no    | 2.27 | 0.22 | [1.83, 2.71] | 10.19
#> 
#> Variable predicted: count
#> Predictors modulated: mined
#> Predictions are on the response-scale.
estimate_means(m, "mined", estimate = "average", data = weighted_dg, weights = "weight")
#> Average Predictions
#> 
#> mined | Mean |   SE |       95% CI |     z
#> ------------------------------------------
#> yes   | 0.27 | 0.05 | [0.17, 0.37] |  5.47
#> no    | 2.27 | 0.22 | [1.83, 2.71] | 10.19
#> 
#> Variable predicted: count
#> Predictors modulated: mined
#> Predictions are on the response-scale.

microbenchmark::microbenchmark(
  estimate_means(m, "mined", estimate = "average"),
  estimate_means(
    m,
    "mined",
    estimate = "average",
    data = weighted_dg,
    weights = "weight"
  ),
  times = 20
)
#> Unit: milliseconds
#>                                                                                           expr
#>                                               estimate_means(m, "mined", estimate = "average")
#>  estimate_means(m, "mined", estimate = "average", data = weighted_dg,      weights = "weight")
#>       min       lq     mean   median       uq      max neval cld
#>  466.0961 478.1445 494.1027 482.5915 503.7244 560.9094    20   a
#>  436.3481 460.9678 483.5811 466.8222 482.0480 645.3554    20   a

^{Created on 2026-07-24 with reprex v2.1.1}

[strengejacke]

library(glmmTMB)
library(modelbased)
data(penguins)

m <- glmmTMB(
  body_mass ~ species + sex + (1 | island),
  data = penguins
)

weighted_dg <- insight::get_datagrid(m, include_random = TRUE, weighted = TRUE)

estimate_means(m, "species", estimate = "average")
#> Average Predictions
#> 
#> species   |    Mean |    SE |             95% CI |      z
#> ---------------------------------------------------------
#> Adelie    | 3706.16 | 26.04 | [3655.12, 3757.21] | 142.30
#> Chinstrap | 3733.09 | 38.16 | [3658.29, 3807.89] |  97.82
#> Gentoo    | 5092.44 | 28.85 | [5035.89, 5148.98] | 176.52
#> 
#> Variable predicted: body_mass
#> Predictors modulated: species
estimate_means(m, "species", estimate = "average", data = weighted_dg, weights = "weight")
#> Average Predictions
#> 
#> species   |    Mean |    SE |             95% CI |      z
#> ---------------------------------------------------------
#> Adelie    | 3706.16 | 26.04 | [3655.12, 3757.21] | 142.30
#> Chinstrap | 3733.09 | 38.16 | [3658.29, 3807.89] |  97.82
#> Gentoo    | 5092.44 | 28.85 | [5035.89, 5148.98] | 176.52
#> 
#> Variable predicted: body_mass
#> Predictors modulated: species

microbenchmark::microbenchmark(
  estimate_means(m, "species", estimate = "average"),
  estimate_means(
    m,
    "species",
    estimate = "average",
    data = weighted_dg,
    weights = "weight"
  ),
  times = 20
)
#> Unit: milliseconds
#>                                                                                             expr
#>                                               estimate_means(m, "species", estimate = "average")
#>  estimate_means(m, "species", estimate = "average", data = weighted_dg,      weights = "weight")
#>       min       lq     mean   median       uq      max neval cld
#>  446.5185 455.0209 473.8285 466.7677 476.0482 628.4250    20  a 
#>  439.9400 444.5771 451.7498 449.3822 455.1578 482.2854    20   b

^{Created on 2026-07-24 with reprex v2.1.1}

[strengejacke]

Based on the discussion in easystats/insight#1208, do we just want to have this feature in insight:get_datagrid(), and then maybe add some docs/examples how to do this manually? If it's not worth the effort to do a separate implementation in modelbased?

[DominiqueMakowski]

On this example we gain ~40% of time so it's not negligible...

library(brms)
library(modelbased)
library(marginaleffects)
library(microbenchmark)

data <- data.frame(Outcome = rnorm(10000), Condition = rep_len(c("A", "B", "C"), 10000), Participant = paste0("S", rep_len(1:800, 10000)))

m <- brms::brm(
  formula = Outcome ~ Condition + (Condition | Participant),
  data = data,
  prior = c(
    set_prior("normal(0, 1)", class = "b"),
    set_prior("cauchy(0, 1)", class = "sd"),
    set_prior("cauchy(0, 1)", class = "sigma")
  ),
  sample_prior = "only",
  refresh = 0,
  chains = 8,
  cores = 8,
  iter = 4000,
  algorithm = "sampling",
  backend = "cmdstanr"
)

weighted_dg <- insight::get_datagrid(m, include_random = TRUE, weighted = TRUE)

# estimate_means(m, "Condition", estimate = "average")
# estimate_means(m, "Condition", estimate = "average", data = weighted_dg, weights = "weight")


microbenchmark::microbenchmark(
  estimate_means(m, "Condition", estimate = "average"),
  estimate_means(
    m,
    "Condition",
    estimate = "average",
    data = weighted_dg,
    weights = "weight"
  ),
  times = 5
)

A modelbased's implementation seems relatively trial to add, but since it only works with one specific type of marginalization ("average"), I'm not sure what's the best API.

1. I missed that earlier but the weight column is lower case, and we usually capitalize our col names in easystats so perhaps worth considering renaming to Weight?
2. API idea:

in estimate_means(...):

• if weighted=True is passed & estimate="average", then use it. Don't expose weighted but document that behaviour inside the docs as a "potential" way to make the estimation faster (assuming backend = marginaleffects & number of combinations << number of observations )

[DominiqueMakowski]

microbenchmark::microbenchmark(
  estimate_contrasts(m, "Condition", estimate = "average"),
  estimate_contrasts(
    m,
    "Condition",
    estimate = "average",
    data = weighted_dg,
    weights = "weight"
  ),
  times = 5
)

nit: seconds
                                                                                                  expr      min       lq     mean
                                              estimate_contrasts(m, "Condition", estimate = "average") 19.13018 20.14403 20.57303
 estimate_contrasts(m, "Condition", estimate = "average", data = weighted_dg,      weights = "weight") 14.45850 15.44106 15.95625
   median       uq      max neval
 20.22104 20.82911 22.54076     5
 16.02166 16.71115 17.14888     5

On contrasts too

[strengejacke]

but since it only works with one specific type of marginalization ("average")

Not exactly. Whenever we explicitly pass data (or newdata), in combination with by, it's passed to by and newdata in marginaleffects, thus generating average marginal means, and the estimate option is overridden - unless estimate = "population", in this case, arguments are passed to variables and newdata in marginaleffects, creating counterfactual predictions

[strengejacke]

Ok, based on easystats/insight#1211 (now merged into insight main), I have opened a PR that implements a "fast" mode: #650

Benchmark for your example (no binning, default = 5 bins for numerics, fast = 10 = 10 bins for numerics):

Unit: seconds
                                                              expr       min       lq     mean   median       uq      max neval cld
              estimate_means(m, "Condition", estimate = "average") 16.017515 16.85124 16.94802 17.00653 17.10221 17.76261     5  a 
 estimate_means(m, "Condition", estimate = "average", fast = TRUE)  9.538026 10.20541 11.19938 10.89997 12.55424 12.79929     5   b
   estimate_means(m, "Condition", estimate = "average", fast = 10) 10.747865 11.95850 12.61095 12.59092 13.62962 14.12786     5   b

Accuracy is also fine:

# accuracy
estimate_means(m, "Condition", estimate = "average") |> print(digits = 4)
estimate_means(m, "Condition", estimate = "average", fast = TRUE) |> print(digits = 4)
estimate_means(m, "Condition", estimate = "average", fast = 10) |> print(digits = 4)