R/rKfitting.R

In femoerman/PBPGM: Parallellised Bayesian Population Growth Models

#' rKfitting function
#'
#' Function to do the actual parallelised r-K fitting of the models
#'
#' @param rK.code Stan code generated by the GeneraterKcode 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 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
#' rKfitting()
rKfitting <- function(rK.code, dd, K.prior, r0.prior, N0.prior, sdev.prior, cores.to.use, iter, warmup, chains){
  # rK model as ODE
  ode.model_rK = function(t,N,p){
    with(as.list(p),{
      dNdt = p[1] * (1 - (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_rK = stan_model(model_code=rK.code)

  #reset priors to non-log scale
  r0.prior <- exp(r0.prior)
  K.prior <- exp(K.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_rK = 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 rK fitting
      init_rK=rep(list(list(log_r=log(r0.prior),
                            log_K = log(K.prior),
                            log_N0sim=log(N0.prior),
                            sdev=sdev.prior))
                  ,chains)

      # do the rK fit
      fit_rK = sampling(s_model_rK,
                        data=data_rK,
                        iter=iter,
                        warmup=warmup,
                        thin=thin,
                        chains=chains,
                        init=init_rK
      )

      # check model
      print(fit_rK)

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

      # save posterios distributions of parameters of interest
      act_r0_dist <- exp(as.numeric(samples_rK[[1]][,"log_r"]))
      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_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_rK,
                                 parms=c(median(act_r0_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_rK,
                            parms=exp(as.numeric(samples_rK[[1]][act_parms_no,c("log_r","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_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)
}