stan model using log_exp_sum()

size_comm_example_stan.R

rm(list = ls())

# Stan message board example
library(tidyverse)
library(magrittr)
library(rstan)

d <- data.frame(ln_rate = c(1,1.5,3,5), temp = c(4, 5, 6, 7), group = c(0,0,1,1))

b <- data.frame(temp = c(4,4,4,5,5,5,5,5,6,6,6, 7,7,7,7,7)) %>%
  mutate(., mass = round(rnorm(n(), mean = 3, sd = 0.5), digits = 2))

# build model in stan
stan_model <- '
data{
  int N;                              // number of samples
  int K;                              // number of mass samples
  int mass_n[N];                      // number of masses per sample
  vector[N] ln_rate;                  // rate
  vector[N] temp;                     // temperature
  vector[K] mass;                     // masses
}
parameters{  
  real<lower=0> sigma;      // standard deviation
  real ln_r0;               // community normalisation constant
  real bT;                  // effect of temperature
  real a;                   // size scaling exponent
}
model{
  int pos;
  vector[N] mu;
  vector[N] ln_biomass;
  ln_r0 ~ normal(0,10);
  bT ~ normal(0,10);
  a ~ normal(0,10);
  sigma ~ cauchy(0,2);

  pos = 1;
  for(n in 1:N){
    ln_biomass[n] = log_sum_exp(a * segment(mass, pos, mass_n[n]));
    pos = pos + mass_n[n]; 
  }  

  # likelihood regression for size dependence of rate
  mu = ln_r0 + ln_biomass + bT*temp;
  ln_rate ~  normal(mu, sigma);

}'
  
# make list
m_list <- b %>% 
  group_by(temp) %>% 
  summarise(n = n()) %>% 
  data.frame()
m = m_list$n
m_list <- as.list(m_list$n)

# create data ####
stan_datalist = list(N = nrow(d),
                     K = nrow(b),
                     mass_n = m,
                     temp = d$temp, 
                     ln_rate = d$ln_rate,
                     mass = log(b$mass))

# initial values function ####
stan_inits <- function() list(a = 0.75, bT = 0, ln_r0 = 1, tau_R = 1)

# run stan ####
model_stan <- rstan::stan(model_code = stan_model,
                   data = stan_datalist, 
                   init = stan_inits, 
                   iter = 1e4, 
                   chains = 1)
model_stan