Visualizes three dimensional Dirichlet distributions

dirichlet-triangle.R

# devtools::install_github("dkahle/dirichlet")
library('TeX')
library('ggplot2')
library('gridExtra')
library('dirichlet')


plot_dirichlet <- function(alphas = c(.5, .5, .5), add_points = FALSE) {
  f <- function(v) ddirichlet(v, alphas)
  mesh <- simplex_mesh(.0025) %>% as.data.frame %>% tbl_df
  mesh$f <- mesh %>% apply(1, function(v) f(bary2simp(v)))
  
  title <- TeX(
    sprintf('$\\alpha = (%.0f,  %.0f,  %.0f)$', alphas[1], alphas[2], alphas[3])
  )
  
  p <- ggplot(mesh, aes(x, y)) +
    geom_raster(aes(fill = f)) +
    coord_equal(xlim = c(0, 1), ylim = c(0, 1)) +
    # scale_colour_gradient(low = 'white', high = 'black') +
    scale_x_continuous(breaks = scales::pretty_breaks(n = 5)) +
    ggtitle(title) +
    theme(
      axis.line = element_blank(),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      panel.border = element_blank(),
      panel.background = element_blank(),
      axis.title.x = element_blank(),
      axis.title.y = element_blank(),
      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),
      legend.position = 'None',
      plot.title = element_text(hjust = .5, vjust = -10)
    )
  
  if (add_points) {
    points <- rdirichlet(250, alphas) %>% simp2bary %>% 
      as.data.frame %>% tbl_df %>% rename(x = V1, y = V2)
    p <- p + geom_point(data = points, color = 'orange', alpha = .3)
  }
  
  p
}


p1 <- plot_dirichlet(c(1, 1, 1), add_points = TRUE)
p2 <- plot_dirichlet(c(5, 5, 1), add_points = TRUE)
p3 <- plot_dirichlet(c(20, 20, 20), add_points = TRUE)

grid.arrange(p1, p2, p3, nrow = 1)