Dirichlet-multinomial custom family in brms

Hi all

Below I have produced a small simulated data set and used brms to fit a multinomial without any problems. But I am doing something wrong trying to get a custom_family working for the Dirichlet-multinomial (DM). I found the a closed-form PMF for the DM here (Transforming a multinomial model into a dirichlet-multinomial - #2 by nhuurre) and I have been fiddling around with the variables in custom_family and read up as much as I can but just can’t find what I need. Any pointers would be much appreciated.

library(tidyverse)
library(brms)

# Simulate some data

n <- 100
rdata <- t(rmultinom(n = n, size = 5, prob = 1:10))
data <- data.frame(x = factor(sample.int(n = 5, size = n, replace = TRUE)), 
                   size = rowSums(rdata))
data$Y <- rdata

# Multinomial model - this works OK

fit1 <- brm(formula = Y | trials(size) ~ x, data = data, 
            iter = 10, refresh = 1, family = multinomial())

# Dirichlet-multinomial model - cannot get to run

dirichlet_multinomial <- custom_family(
  name = "dirichlet_multinomial", 
  dpars = c("mu", "phi"),
  # multi_dpars = "mu",
  links = "logit", 
  type = "int", 
  lb = c(NA, 0),
  # vars = "vint1",
  # vars = "trials[n]",
  loop = TRUE
)

stan_density <- "
  real dirichlet_multinomial_lpmf(int[] y, vector mu, real phi) {
    real sum_mu = sum(mu);
    real output = lgamma(sum_mu) - lgamma(sum(y) + sum_mu) +
           // + lgamma(sum(y) + 1) - sum(lgamma(to_vector(y) + 1) // constant, may omit
           sum(lgamma(to_vector(y) + mu)) - sum(lgamma(mu));
    return phi * output;
  }
"

stanvars <- stanvar(scode = stan_density, block = "functions")

fit2 <- brm(Y | trials(size) ~ x, data = data, iter = 10, refresh = 1, 
            family = dirichlet_multinomial, stanvars = stanvars)

Thanks in advance - and many thanks for such an amazing R package, brms has become indispensable to me over the past few years!

  • Operating System: Pop!_OS 22.04 LTS (linux)
  • brms Version: 2.17.0