part of plot function, choice

plot_choice.R

library(tidyverse)
library(ggrepel)
library(rstan)
expose_stan_functions("my_stan_functions.stan") # https://gist.github.com/kagaya/c726f16d82b80026bb1924b408a72b5c
source("utility.R") # https://gist.github.com/kagaya/60c4190ac306840daee54115b3c315c3

plot_choice_random_intercept <- function(d, fit){
  # mu[n,2] =  bias_l[ID[n]] + cwl*C_width[n] + ll*Leg_lack[n];
  # mu[n,3] =  bias_no0 + cwno*C_width[n] + lno*Leg_lack[n]; 
  ms <- rstan::extract(fit)
  new_C_width <- seq(min(d$shell_width), max(d$shell_width), length.out=500)
  
  ## 50 percentile prediction
  bias_l0_50 <- quantile(ms$bias_l0, 0.5)
  cw_l_50 <- quantile(ms$cwl, 0.5)
  bias_no0_50 <- quantile(ms$bias_no0, 0.5)
  cw_no_50 <- quantile(ms$cwno, 0.5)
  
  mu2_50 <- c() 
  mu3_50 <- c()
  theta2_50 <- c()
  theta3_50 <- c()
  
  choice_rng <- matrix(nrow=500, ncol=500)
  for (i in 1:500){
    mu2_50[i] <- bias_l0_50 + cw_l_50*new_C_width[i]
    mu3_50[i] <- bias_no0_50 + cw_no_50*new_C_width[i]
    theta2_50[i] <- 1 + exp(mu2_50[i])/(1 + exp(mu2_50[i]) + exp(mu3_50[i]))
    theta3_50[i] <- 1 + 2 * exp(mu3_50[i])/(1 + exp(mu2_50[i]) + exp(mu3_50[i])) # to overlay choice plot
    
    choice_rng[,i] <- my_categorical_logit_rng(500, c(0, mu2_50[i], mu3_50[i]))
  }
  
  colnames(choice_rng) <- new_C_width
  choice_rng <- choice_rng %>% as_tibble() %>% gather(x,y) %>% mutate(x=as.numeric(x))
  
  pred_50 <- tibble(carapace_width=new_C_width, 
                    theta2_50=theta2_50,
                    theta3_50=theta3_50)
  
  
  gid <- levels(d$id)[table(d$id) > 1]
  
  mul <- factor(rep("s", dim(d)[1]), levels=c("g", "s"))
  for (i in 1:dim(d)[1]){
    if (d$id[i] %in% gid){
      mul[i] <- c("g")
    }
  }
  
  d$mul <- mul
  
  dd <- d
  dd$choice <- as.integer(d$choice)
  d_seg <- dd %>%
    group_by(id) %>%
    summarise( choice_min=min(choice),
               choice_max=max(choice),
               shell_width=first(shell_width)) %>%
    drop_na()
  
  d <- add_chosen_col_choice(d)
  
  plt <- ggplot(d, aes(shell_width, choice)) +
    stat_density2d(geom="raster", aes(x=x, y=y, fill=..density..), contour=F, data=choice_rng) +
    geom_line(aes(carapace_width, theta2_50), alpha=0.2, data=pred_50) +
    geom_line(aes(carapace_width, theta3_50), alpha=0.2, data=pred_50) +
    geom_point(aes(size=chosen_num), pch=0, alpha=0.8) +
    geom_segment(
      data = d_seg,
      alpha = 0.5,
      lty=3,
      aes(x = shell_width,
          y = choice_min,
          xend = shell_width,
          yend = choice_max)) +
    geom_text_repel(aes(shell_width, choice, label=id), alpha=0.2, size=3) +
    scale_shape_manual(values=c(1,3)) +
    scale_fill_gradient(low="white", high="gray") +
    ylim(c(0.8, 3.2)) +
    xlab("carapace width (cm)") +
    ylab("behavioral choice") +
    theme_classic()
  return(plt)
}