induced_dirichlet_ordbeta.R

library(tidyverse)
library(magrittr)
library(brms)
#library(betareg) # for example data
#### Code declaring ordered beta family for brms ####
# code from Robert Kubinec

# custom family

ord_beta_reg <- custom_family("ord_beta_reg",
                              dpars=c("mu","phi","cutzero","cutone"),
                              links=c("logit","log",NA,NA),
                              lb=c(NA,0,NA,NA),
                              type="real")

# stan code for density of the model

stan_funs <- "real ord_beta_reg_lpdf(real y, real mu, real phi, real cutzero, real cutone) {

    //auxiliary variables
    real mu_logit = logit(mu);
    vector[2] thresh;
    thresh[1] = cutzero;
    thresh[2] = cutzero + exp(cutone);
    
  if(y==0) {
      return log1m_inv_logit(mu_logit - thresh[1]);
    } else if(y==1) {
      return log_inv_logit(mu_logit  - thresh[2]);
    } else {
      return log(inv_logit(mu_logit  - thresh[1]) - inv_logit(mu_logit - thresh[2])) +
                beta_proportion_lpdf(y|mu,phi);
    }
  }"

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

# For pulling posterior predictions

posterior_predict_ord_beta_reg <- function(i, draws, ...) {
  mu <- draws$dpars$mu[, i]
  if(NCOL(draws$dpars$phi)==1){phi <- draws$dpars$phi}else
  {phi <- draws$dpars$phi[, i]}
  cutzero <- draws$dpars$cutzero
  cutone <- draws$dpars$cutone
  N <- length(phi)
  
  thresh1 <- cutzero
  thresh2 <- cutzero + exp(cutone)
  
  pr_y0 <- 1 - plogis(qlogis(mu) - thresh1)
  pr_y1 <- plogis(qlogis(mu) - thresh2)
  pr_cont <- plogis(qlogis(mu)-thresh1) - plogis(qlogis(mu) - thresh2)
  out_beta <- rbeta(n=N,mu*phi,(1-mu)*phi)
  
  # now determine which one we get for each observation
  outcomes <- sapply(1:N, function(i) {
    sample(1:3,size=1,prob=c(pr_y0[i],pr_cont[i],pr_y1[i]))
  })
  
  final_out <- sapply(1:length(outcomes),function(i) {
    if(outcomes[i]==1) {
      return(0)
    } else if(outcomes[i]==2) {
      return(out_beta[i])
    } else {
      return(1)
    }
  })
  
  final_out
  
}

# for calculating marginal effects/conditional expectations

posterior_epred_ord_beta_reg <- function(draws) {
  
  cutzero <- draws$dpars$cutzero
  cutone <- draws$dpars$cutone
  
  mu <- draws$dpars$mu
  
  thresh1 <- cutzero
  thresh2 <- cutzero + exp(cutone)
  
  low <- 1 - plogis(qlogis(mu) - thresh1)
  middle <- plogis(qlogis(mu)-thresh1) - plogis(qlogis(mu) - thresh2)
  high <- plogis(qlogis(mu) - thresh2)
  
  low*0 + middle*mu + high
}

# for calcuating LOO and Bayes Factors

log_lik_ord_beta_reg <- function(i, draws) {
  
  mu <- draws$dpars$mu[,i]
  if(NCOL(draws$dpars$phi)==1){phi <- draws$dpars$phi}else
  {phi <- draws$dpars$phi[, i]}
  y <- draws$data$Y[i]
  cutzero <- draws$dpars$cutzero
  cutone <- draws$dpars$cutone
  
  thresh1 <- cutzero
  thresh2 <- cutzero + exp(cutone)
  
  if(y==0) {
    out <- log(1 - plogis(qlogis(mu) - thresh1))
  } else if(y==1) {
    out <- log(plogis(qlogis(mu) - thresh2))
  } else {
    out <- log(plogis(qlogis(mu)-thresh1) - plogis(qlogis(mu) - thresh2)) + dbeta(y,mu*phi,(1-mu)*phi,log=T)
  }
  
  out
  
}


###### Code declaring induced dirichlet prior ####
# code from Michael Betancourt
# discussion here: https://discourse.mc-stan.org/t/dirichlet-prior-on-ordinal-regression-cutpoints-in-brms/20640
dirichlet_prior <- "
  real induced_dirichlet_lpdf(vector c, vector alpha, real phi) {
    int K = num_elements(c) + 1;
    vector[K - 1] sigma = inv_logit(phi - c);
    vector[K] p;
    matrix[K, K] J = rep_matrix(0, K, K);
    
    // Induced ordinal probabilities
    p[1] = 1 - sigma[1];
    for (k in 2:(K - 1))
      p[k] = sigma[k - 1] - sigma[k];
    p[K] = sigma[K - 1];
    
    // Baseline column of Jacobian
    for (k in 1:K) J[k, 1] = 1;
    
    // Diagonal entries of Jacobian
    for (k in 2:K) {
      real rho = sigma[k - 1] * (1 - sigma[k - 1]);
      J[k, k] = - rho;
      J[k - 1, k] = rho;
    }
    
    return   dirichlet_lpdf(p | alpha)
           + log_determinant(J);
  }
"
dirichlet_prior_stanvar <- stanvar(scode = dirichlet_prior, block = "functions")

stanvar(scode = "ordered[2] thresh;
              thresh[1] = cutzero;
              thresh[2] = cutzero+exp(cutone);", 
        block = "tparameters") -> # there might be a better way to specify this
  dirichlet_prior_ordbeta_stanvar

dirichlet_prior_ordbeta <- set_prior("target += induced_dirichlet_lpdf(thresh | rep_vector(1, 3), 0)", check=FALSE)

# # example data
# data("MockJurors", package = "betareg")
# 
# MockJurors %<>% as_tibble() %>% mutate(confidence = (round((((confidence*104)+0.5)/103)*100)-1)/100)
# # reverse the transformed beta transformation
# 
# brm(formula = bf(confidence ~ 0 + Intercept + conflict*verdict), data=MockJurors,
#     family=ord_beta_reg,
#     stanvars=stanvars) -> 
#   brms_fit
# 
# brm(formula = bf(confidence ~ 0 + Intercept + conflict*verdict), 
#     data=MockJurors,
#     family=ord_beta_reg,
#     prior = dirichlet_prior_ordbeta,
#     stanvars=stanvars+dirichlet_prior_stanvar+dirichlet_prior_ordbeta_stanvar) -> 
#   brms_fit_dirichlet
# 
# brm(formula = bf(confidence ~ 0 + Intercept + conflict*verdict), 
#     data=MockJurors %>% filter(confidence!=0, confidence!=1),
#     family=ord_beta_reg ,
#     prior = dirichlet_prior_ordbeta,
#     stanvars=stanvars+dirichlet_prior_stanvar+dirichlet_prior_ordbeta_stanvar,
#     control=list(adapt_delta=.99)) -> 
#   brms_fit_dirichlet_no_boundary_data