R/BHfitting.R

In femoerman/PBPGM: Parallellised Bayesian Population Growth Models

#' BHfitting function
#'
#' Function to do the actual parallelised Beverton holt fitting of the models
#'
#' @param BH.code Stan code generated by the GenerateBHcode functions
#' @param dd Dataframe with columns for population size data (colname=popsize), time data (colname=time) and unique identifiers for each population (colname=ident)
#' @param K.prior Prior value for the mean carrying capacity (K). Must be on log scale and numeric
#' @param Ksd.prior Prior value for the standard deviation on carrying capacity (K). Must be on log scale and numeric
#' @param r0.prior Prior value for the mean intrinsic rate of growth (r0). Must be on log scale and numeric
#' @param r0sd.prior Prior value for the standard deviation on intrinsic rate of growth (r0). Must be on log scale and numeric
#' @param d.prior Prior value for the mean death rate (d). Must be on log scale and numeric
#' @param d.prior Prior value for the mean death rate (d). Must be on log scale and numeric
#' @param N0.prior Prior value for the mean starting population size (N0). Must be on log scale and numeric
#' @param N0sd.prior Prior value for the standard deviation on starting population size (N0). Must be on log scale and numeric
#' @param sdev.prior Prior value for the standard deviation for the model fitting. Must be numeric. Defaults to 1.
#' @param cores.to.use Available cores to do the bayesian models
#' @param iter Number of iterations to run for the model, defaults to 1e4. Must be integer
#' @param warmup Number of iterations to run warmup. Defaults to 1e3. Must be integer and smaller than warmup
#' @param chains Number of chains to run for each fit. Must be integer. Defaults to 1
#' @keywords Beverton Holt Bayesian parallel model fitting
#' @import tidyverse
#' @import rstan
#' @import loo
#' @import deSolve
#' @import coda
#' @import parallel
#' @export
#' @examples
#' BHfitting()
BHfitting <- function(BH.code, dd, K.prior, r0.prior, d.prior, N0.prior, sdev.prior, cores.to.use, iter, warmup, chains){
  # rK model as ODE
  ode.model_BH = function(t,N,p){
    with(as.list(p),{
      dNdt = ((p[1] + p[2])/(1 + ((p[1]/(p[3]*p[2])) * N)) - p[2])*N
      return(list(dNdt))
    })
  }

  ###############################################################################
  # stan options
  chains = chains
  rstan_options(auto_write = TRUE)
  options(mc.cores = 1)
  iter   =  iter
  warmup =  warmup
  thin   =     1

  # compile models
  s_model_BH = stan_model(model_code=BH.code)

  #reset priors to non-log scale
  r0.prior <- exp(r0.prior)
  K.prior <- exp(K.prior)
  d.prior <- exp(d.prior)
  N0.prior <- exp(N0.prior)


  ###############################################################################
  # global counter var
  cnt <- 0

  ###############################################################################
  # loop over blocks and replicates
  parbayes <- function(act_rep, dd){
    cnt <- act_rep

    ########################################
    # get the subset
    act_name <- unique(dd$ident)[act_rep]
    act_subset <- which(dd$ident == unique(dd$ident)[act_rep])

    # get the densities and current time points
    act_densities <- dd$popsize[act_subset]
    act_times <- dd$time[act_subset]

    act_times <- act_times[which(act_densities != 0)]
    act_densities <- act_densities[which(act_densities != 0)]

    #Get output dataframe
    output <- data.frame(row.names = c(act_name))

    if(length(act_densities) >= 3){

      #plot(act_densities ~ act_times)

      ###############################################################################
      # rK fitting

      # create data object for rstan
      data_BH = list(n = length(act_times)-1,
                     log_N0 = log(act_densities[1]),
                     log_N  = log(act_densities[2:length(act_times)]),
                     t0    = act_times[1],
                     t     = act_times[2:length(act_times)])

      # initial values for BH fitting
      init_BH=rep(list(list(log_r=log(r0.prior),
                            log_d=log(d.prior),
                            log_K = log(K.prior),
                            log_N0sim=log(N0.prior),
                            sdev=sdev.prior))
                  ,chains)

      # do the BH fit
      fit_BH = sampling(s_model_BH,
                        data=data_BH,
                        iter=iter,
                        warmup=warmup,
                        thin=thin,
                        chains=chains,
                        init=init_BH
      )

      # check model
      print(fit_BH)

      # convert to mcmc object and plot posterior distributions
      samples_rK <- mcmc.list(lapply(1:ncol(fit_BH), function(x) mcmc(as.array(fit_BH)[,1,][,names(init_BH[[1]])])))
      #get summary statistics
      sumstats <- summary(fit_BH)$summary

      # save posterios distributions of parameters of interest
      act_r0_dist <- exp(as.numeric(samples_rK[[1]][,"log_r"]))
      act_d_dist <- exp(as.numeric(samples_rK[[1]][,"log_d"]))
      act_K_dist <- exp(as.numeric(samples_rK[[1]][,"log_K"]))
      act_alpha_dist <- exp(as.numeric(samples_rK[[1]][,"log_r"]))/exp(as.numeric(samples_rK[[1]][,"log_K"]))
      act_N0sim_dist <- exp(as.numeric(samples_rK[[1]][,"log_N0sim"]))

      # save in list
      output$all_r0_dist <- list(act_r0_dist)
      output$logr0_N_eff <- sumstats["log_r", "n_eff"]
      output$logr0_Rhat <- sumstats["log_r", "Rhat"]
      output$all_d_dist <- list(act_d_dist)
      output$logd_N_eff <- sumstats["log_d", "n_eff"]
      output$logd_Rhat <- sumstats["log_d", "Rhat"]
      output$all_K_dist <- list(act_K_dist)
      output$logK_N_eff <- sumstats["log_K", "n_eff"]
      output$logK_Rhat <- sumstats["log_K", "Rhat"]
      output$all_alpha_dist <- list(act_alpha_dist)
      ###############################################################################
      # get predictions
      # set time interval for prediction
      timessim <- seq(min(act_times),max(act_times),len=100)

      # get median model
      Nsim = as.data.frame(lsoda(y=median(act_N0sim_dist),
                                 times=timessim,
                                 func=ode.model_BH,
                                 parms=c(median(act_r0_dist), median(act_d_dist), median(act_K_dist))))[,2]

      # help function for prediction of CI
      hlp_pred <- function(act_parms_no){
        as.data.frame(lsoda(y=act_N0sim_dist[act_parms_no],
                            times=timessim,
                            func=ode.model_BH,
                            parms=exp(as.numeric(samples_rK[[1]][act_parms_no,c("log_r", "log_d","log_K")]))))[,2]
      }

      # predict distribution at every time point
      pred_pop_dyn <- sapply(1:dim(samples_rK[[1]])[1],hlp_pred)

      # function for extracting CI and median of prediction
      get_quantile <- function(act_t_no){
        quantile(pred_pop_dyn[act_t_no,],p=c(0.025,0.5,0.975),na.rm=T)
      }

      # get prediction
      model_predictions <- sapply(1:length(timessim),get_quantile)

      ###############################################################################
      output$act_densities <- list(act_densities)
      output$act_times <- list(act_times)
      output$act_2.5 <- list(model_predictions[1, ])
      output$act_97.5 <- list(model_predictions[3, ])
      output$act_50 <- list(model_predictions[2, ])
      output$timesim <- list(timessim)
      output$Nsim <- list(Nsim)
    }else{
      output$all_r0_dist <- NA
      output$logr0_N_eff <- NA
      output$logr0_Rhat <- NA
      output$all_d_dist <- NA
      output$logd_N_eff <- NA
      output$logd_Rhat <- NA
      output$all_K_dist <- NA
      output$logK_N_eff <- NA
      output$logK_Rhat <- NA
      output$all_alpha_dist <- NA
      output$act_densities <- list(act_densities)
      output$act_times <- list(act_times)
      output$act_2.5 <- NA
      output$act_97.5 <- NA
      output$act_50 <- NA
      output$timesim <- NA
      output$Nsim <- NA
    }

    output$ident <- act_name
    output$cnt <- cnt

    return(output)
  }

  #Run parallelised Bayesian model
  cl <- makeCluster(cores.to.use)
  tempoutput <- mclapply(1:length(unique(dd$ident)), parbayes, dd, mc.cores = cores.to.use)
  stopCluster(cl)
  fulloutput <- data.frame(tempoutput[1])
  for (elem in 2:length(tempoutput)){
    fulloutput <- rbind(fulloutput, data.frame(tempoutput[elem]))

  }
  fulloutput <- arrange(fulloutput, ident)
  return(fulloutput)
}