Create native Dirichlet multinomial family

DESCRIPTION

@@ -2,8 +2,8 @@ Package: brms
 Encoding: UTF-8
 Type: Package
 Title: Bayesian Regression Models using 'Stan'
-Version: 2.22.8
-Date: 2024-12-05
+Version: 2.22.9
+Date: 2025-02-01
 Authors@R:
     c(person("Paul-Christian", "Bürkner", email = "paul.buerkner@gmail.com",
              role = c("aut", "cre")),

NAMESPACE

@@ -410,6 +410,7 @@ export(dhurdle_negbinomial)
 export(dhurdle_poisson)
 export(dinv_gaussian)
 export(dirichlet)
+export(dirichlet_multinomial)
 export(dlogistic_normal)
 export(dmulti_normal)
 export(dmulti_student_t)

NEWS.md

@@ -5,8 +5,10 @@
 * Fit extended-support Beta models via family `xbeta` 
 thanks to Ioannis Kosmidis. (#1698)
 * Add a `seed` argument to `loo_R2` thanks to Marco Colombo. (#1713)
+* Add family `dirichlet_multinomial` to fit overdispersed 
+multinomial data thanks to Tom Peatman. (#1729)
 * Add the `int_step` R function to match the corresponding Stan 
-function. (#1734)
+function thanks to Daniel Sabanes Bove. (#1734)
 
 ### Bug Fixes
 
@@ -18,7 +20,7 @@ models, `log_lik` will not use `option(mc.cores)` anymore.
 These changes may be reverted once the underlying causes of this
 issue have been fixed. (#1658)
 * Align the definition of the R function `step()` with the definition in Stan,
-such that `step(0) == 1` (#1734)
+such that `step(0) == 1` thanks to Daniel Sabanes Bov. (#1734)
 
 ### Other Changes
 

R/distributions.R

@@ -2185,6 +2185,18 @@ dmultinomial <- function(x, eta, log = FALSE) {
   out
 }
 
+# density of the multinomial distribution with the softmax transform
+# @param x positive integers not greater than ncat
+# @param eta the linear predictor (of length or ncol ncat)
+# @param phi the dispersion parameter
+# @param log return values on the log scale?
+ddirichletmultinomial <- function(x, eta, phi, log = FALSE) {
+  require_package("extraDistr")
+  alpha <- softmax(eta) * phi
+  size <- sum(x)
+  extraDistr::ddirmnom(x, size = size, alpha = alpha, log = log)
+}
+
 # density of the cumulative distribution
 #
 # @param x Integer vector containing response category indices to return the

R/families.R

@@ -18,7 +18,7 @@
 #'   \code{inverse.gaussian}, \code{exponential}, \code{weibull},
 #'   \code{frechet}, \code{Beta}, \code{dirichlet}, \code{von_mises},
 #'   \code{asym_laplace}, \code{gen_extreme_value}, \code{categorical},
-#'   \code{multinomial}, \code{cumulative}, \code{cratio}, \code{sratio},
+#'   \code{multinomial}, \code{dirichlet_multinomial}, \code{cumulative}, \code{cratio}, \code{sratio},
 #'   \code{acat}, \code{hurdle_poisson}, \code{hurdle_negbinomial},
 #'   \code{hurdle_gamma}, \code{hurdle_lognormal}, \code{hurdle_cumulative},
 #'   \code{zero_inflated_binomial}, \code{zero_inflated_beta_binomial},
@@ -51,8 +51,9 @@
 #'   consecutive thresholds to the same value, and
 #'   \code{"sum_to_zero"} ensures the thresholds sum to zero.
 #' @param refcat Optional name of the reference response category used in
-#'   \code{categorical}, \code{multinomial}, \code{dirichlet} and
-#'   \code{logistic_normal} models. If \code{NULL} (the default), the first
+#'   \code{categorical}, \code{multinomial}, \code{dirichlet},
+#'   \code{dirichlet_multinomial} and \code{logistic_normal} models.
+#'   If \code{NULL} (the default), the first
 #'   category is used as the reference. If \code{NA}, all categories will be
 #'   predicted, which requires strong priors or carefully specified predictor
 #'   terms in order to lead to an identified model.
@@ -76,8 +77,9 @@
 #'   can be used for binary regression (i.e., most commonly logistic
 #'   regression).}
 #'
-#'   \item{Families \code{categorical} and \code{multinomial} can be used for
-#'   multi-logistic regression when there are more than two possible outcomes.}
+#'   \item{Families \code{categorical}, \code{multinomial} and
+#'   \code{dirichlet_multinomial} can be used for multi-logistic regression
+#'   when there are more than two possible outcomes.}
 #'
 #'   \item{Families \code{cumulative}, \code{cratio} ('continuation ratio'),
 #'   \code{sratio} ('stopping ratio'), and \code{acat} ('adjacent category')
@@ -150,8 +152,8 @@
 #'   \code{acat}, and \code{hurdle_cumulative} support \code{logit},
 #'   \code{probit}, \code{probit_approx}, \code{cloglog}, and \code{cauchit}.}
 #'
-#'   \item{Families \code{categorical}, \code{multinomial}, and \code{dirichlet}
-#'   support \code{logit}.}
+#'   \item{Families \code{categorical}, \code{multinomial},
+#'   \code{dirichlet_multinomial} and \code{dirichlet} support \code{logit}.}
 #'
 #'   \item{Families \code{Gamma}, \code{weibull}, \code{exponential},
 #'   \code{frechet}, and \code{hurdle_gamma} support
@@ -812,6 +814,15 @@ multinomial <- function(link = "logit", refcat = NULL) {
   .brmsfamily("multinomial", link = link, slink = slink, refcat = refcat)
 }
 
+#' @rdname brmsfamily
+#' @export
+dirichlet_multinomial <- function(link = "logit", link_phi = "log",
+                                  refcat = NULL) {
+  slink <- substitute(link)
+  .brmsfamily("dirichlet_multinomial", link = link, slink = slink,
+              link_phi = link_phi, refcat = refcat)
+}
+
 #' @rdname brmsfamily
 #' @export
 cumulative <- function(link = "logit", link_disc = "log",

R/family-lists.R

@@ -138,6 +138,20 @@
   )
 }
 
+.family_dirichlet_multinomial <- function() {
+  list(
+    links = "logit",
+    dpars = "phi",
+    multi_dpars = "mu",  # size determined by the data
+    type = "int", ybounds = c(-Inf, Inf),
+    closed = c(NA, NA),
+    ad = c("weights", "subset", "trials", "index"),
+    specials = c("multinomial", "joint_link"),
+    include = "fun_dirichlet_multinomial_logit.stan",
+    normalized = ""
+  )
+}
+
 .family_beta <- function() {
   list(
     links = c(

R/log_lik.R

@@ -837,6 +837,17 @@ log_lik_multinomial <- function(i, prep) {
   log_lik_weight(out, i = i, prep = prep)
 }
 
+log_lik_dirichlet_multinomial <- function(i, prep) {
+  stopifnot(prep$family$link == "logit")
+  eta <- get_Mu(prep, i = i)
+  eta <- insert_refcat(eta, refcat = prep$refcat)
+  phi <- get_dpar(prep, "phi", i = i)
+  cats <- seq_len(prep$data$ncat)
+  alpha <- dcategorical(cats, eta = eta) * phi
+  out <- ddirichletmultinomial(prep$data$Y[i, ], eta = eta, phi = phi, log = TRUE)
+  log_lik_weight(out, i = i, prep = prep)
+}
+
 log_lik_dirichlet <- function(i, prep) {
   stopifnot(prep$family$link == "logit")
   eta <- get_Mu(prep, i = i)

R/posterior_epred.R

@@ -587,6 +587,21 @@ posterior_epred_multinomial <- function(prep) {
   out
 }
 
+posterior_epred_dirichlet_multinomial <- function(prep) {
+  get_counts <- function(i) {
+    eta <- insert_refcat(slice_col(eta, i), refcat = prep$refcat)
+    dcategorical(cats, eta = eta) * trials[i]
+  }
+  # dirichlet part included in mu
+  eta <- get_Mu(prep)
+  cats <- seq_len(prep$data$ncat)
+  trials <- prep$data$trials
+  out <- abind(lapply(seq_cols(eta), get_counts), along = 3)
+  out <- aperm(out, perm = c(1, 3, 2))
+  dimnames(out)[[3]] <- prep$cats
+  out
+}
+
 posterior_epred_dirichlet <- function(prep) {
   get_probs <- function(i) {
     eta <- insert_refcat(slice_col(eta, i), refcat = prep$refcat)

R/posterior_predict.R

@@ -882,6 +882,16 @@ posterior_predict_multinomial <- function(i, prep, ...) {
   rblapply(seq_rows(p), function(s) t(rmultinom(1, size, p[s, ])))
 }
 
+posterior_predict_dirichlet_multinomial <- function(i, prep, ...) {
+  eta <- get_Mu(prep, i = i)
+  eta <- insert_refcat(eta, refcat = prep$refcat)
+  phi <- get_dpar(prep, "phi", i = i)
+  alpha <- dcategorical(seq_len(prep$data$ncat), eta = eta) * phi
+  p <- rdirichlet(prep$ndraws, alpha = alpha)
+  size <- prep$data$trials[i]
+  rblapply(seq_rows(p), function(s) t(rmultinom(1, size, p[s, ])))
+}
+
 posterior_predict_dirichlet <- function(i, prep, ...) {
   eta <- get_Mu(prep, i = i)
   eta <- insert_refcat(eta, refcat = prep$refcat)

R/stan-likelihood.R

@@ -777,6 +777,14 @@ stan_log_lik_multinomial <- function(bterms, ...) {
   sdist("multinomial_logit2", p$mu, vec = FALSE)
 }
 
+stan_log_lik_dirichlet_multinomial <- function(bterms, ...) {
+  stopifnot(bterms$family$link == "logit")
+  mu <- stan_log_lik_dpars(bterms, reqn = TRUE, dpars = "mu", type = "multi")$mu
+  reqn_phi <- is_pred_dpar(bterms, "phi")
+  phi <- stan_log_lik_dpars(bterms, reqn = reqn_phi, dpars = "phi")$phi
+  sdist("dirichlet_multinomial_logit2", mu, phi, vec = FALSE)
+}
+
 stan_log_lik_dirichlet <- function(bterms, ...) {
   stopifnot(bterms$family$link == "logit")
   mu <- stan_log_lik_dpars(bterms, reqn = TRUE, dpars = "mu", type = "multi")$mu

inst/chunks/fun_dirichlet_multinomial_logit.stan

@@ -0,0 +1,19 @@
+  /* dirichlet-multinomial-logit log-PMF
+   * Args:
+   *   y: array of integer response values
+   *   mu: vector of category logit probabilities
+   *   phi: precision parameter (sum of Dirichlet alphas)
+   * Returns:
+   *   a scalar to be added to the log posterior
+  */
+  real dirichlet_multinomial_logit2_lpmf(array[] int y, vector mu, real phi) {
+    // get Dirichlet alphas
+    int N = num_elements(mu);
+    vector[N] alpha = phi * softmax(mu);
+
+    // get trials from y
+    real T = sum(y);
+
+    return lgamma(phi) + lgamma(T + 1.0) - lgamma(T + phi) +
+      sum(lgamma(to_vector(y) + alpha)) - sum(lgamma(alpha)) - sum(lgamma(to_vector(y) + 1));
+  }

tests/local/tests.models-4.R

@@ -181,6 +181,37 @@ test_that("multinomial models work correctly", suppressWarnings({
   expect_ggplot(plot(ce, ask = FALSE)[[1]])
 }))
 
+test_that("dirichlet_multinomial models work correctly", suppressWarnings({
+  require("extraDistr")
+  set.seed(1245)
+  N <- 100
+  dat <- as.data.frame(extraDistr::rdirmnom(N, 10, c(10, 5, 1)))
+  names(dat) <- paste0("y", 1:3)
+  dat$size <- with(dat, y1 + y2 + y3)
+  dat$x <- rnorm(N)
+  dat$y <- with(dat, cbind(y1, y2, y3))
+
+  fit <- brm(
+    y | trials(size) ~ x, data = dat,
+    family = dirichlet_multinomial(),
+    prior = prior("exponential(0.01)", "phi")
+  )
+  print(summary(fit))
+
+  pred <- predict(fit)
+  expect_equal(dim(pred), c(nobs(fit), 4, 3))
+  expect_equal(dimnames(pred)[[3]], c("y1", "y2", "y3"))
+
+  pred_mean <- fitted(fit)
+  expect_equal(dim(pred_mean), c(nobs(fit), 4, 3))
+
+  waic <- waic(fit)
+  expect_range(waic$estimates[3, 1], 550, 650)
+
+  ce <- conditional_effects(fit, categorical = TRUE)
+  expect_ggplot(plot(ce, ask = FALSE)[[1]])
+}))
+
 test_that("dirichlet models work correctly", suppressWarnings({
   set.seed(1246)
   N <- 100

tests/testthat/tests.log_lik.R

@@ -445,6 +445,12 @@ test_that("log_lik for categorical and related models runs without erros", {
   ll <- sapply(1:nobs, brms:::log_lik_multinomial, prep = prep)
   expect_equal(dim(ll), c(ns, nobs))
 
+  prep$data$trials <- sample(1:20, nobs)
+  prep$dpars$phi <- rexp(ns, 10)
+  prep$family <- dirichlet_multinomial()
+  ll <- sapply(1:nobs, brms:::log_lik_dirichlet_multinomial, prep = prep)
+  expect_equal(dim(ll), c(ns, nobs))
+
   prep$data$Y <- prep$data$Y / rowSums(prep$data$Y)
   prep$dpars$phi <- rexp(ns, 10)
   prep$family <- dirichlet()

tests/testthat/tests.posterior_epred.R

@@ -182,7 +182,7 @@ test_that("posterior_epred for advanced count data distributions runs without er
   expect_equal(dim(pred), c(ns, nobs))
 })
 
-test_that("posterior_epred for multinomial and dirichlet models runs without errors", {
+test_that("posterior_epred for multinomial, dirichlet_multinomial and dirichlet models runs without errors", {
   ns <- 15
   nobs <- 8
   ncat <- 3
@@ -198,6 +198,10 @@ test_that("posterior_epred for multinomial and dirichlet models runs without err
   pred <- brms:::posterior_epred_multinomial(prep = prep)
   expect_equal(dim(pred), c(ns, nobs, ncat))
 
+  prep$family <- dirichlet_multinomial()
+  pred <- brms:::posterior_epred_dirichlet_multinomial(prep = prep)
+  expect_equal(dim(pred), c(ns, nobs, ncat))
+
   prep$family <- dirichlet()
   pred <- brms:::posterior_epred_dirichlet(prep = prep)
   expect_equal(dim(pred), c(ns, nobs, ncat))

tests/testthat/tests.posterior_predict.R

@@ -346,6 +346,12 @@ test_that("posterior_predict for categorical and related models runs without err
   pred <- brms:::posterior_predict_multinomial(i = sample(1:nobs, 1), prep = prep)
   expect_equal(dim(pred), c(ns, ncat))
 
+  prep$data$trials <- sample(1:20, nobs)
+  prep$dpars$phi <- rexp(ns, 1)
+  prep$family <- dirichlet_multinomial()
+  pred <- brms:::posterior_predict_dirichlet_multinomial(i = sample(1:nobs, 1), prep = prep)
+  expect_equal(dim(pred), c(ns, ncat))
+
   prep$dpars$phi <- rexp(ns, 1)
   prep$family <- dirichlet()
   pred <- brms:::posterior_predict_dirichlet(i = sample(1:nobs, 1), prep = prep)