dlm-wrapper.R

In TESA-workshops/TESAsamSim: Closed-loop salmon simulator for 2021 TESA workshop

# DLM wrapper and plotting functions ===============================================================================
#
# Description: wrapper and plotting function for Maximum likelihood estimation, Kalman filtering and smoothing of
#              normal linear State Space models, also known as Dynamic Linear Models, using the dlm package by
#              Giovanni Petris (U of Arkansas). Adapted from code generously provided by Cody Szuwalski (NOAA)
#

##require packages
library(dlm)
library(ggpubr)

## fitDLM function
#
# data: a dataframe with three columns for brood years ("byr"), corresponding spawner abundance ("spwn")
#       and resulting recruitment ("rec")
# alpha_vary: should alpha (intercept) be estimated as time varying?
# beta_vary: should beta (slope) be estimated as time varying?

fitDLM <- function(data = bt,
                   alpha_vary = FALSE,
                   beta_vary = FALSE){

  # 0. housekeeping
  lnRS <- log(data$rec/data$spwn)
  alpha <- NULL
  beta <- NULL

  # 1. create a dlm representation of a linear regression model
  mod <- dlmModReg(data$spwn) # this specifies a linear model

  # 2. number of parameters used in the AICc calculation
  dlmPars=3
  if(alpha_vary==TRUE & beta_vary==FALSE){dlmPars=4}
  if(alpha_vary==FALSE & beta_vary==TRUE){dlmPars=4}
  if(alpha_vary==TRUE & beta_vary==TRUE){dlmPars=5}

  # 3. specify the model based on variance structure
  build_mod <- function(parm)
  {
    mod$V <- exp(parm[1]) # observation error variance
    if(alpha_vary==TRUE & beta_vary==FALSE){mod$W[1,1]=exp(parm[2]); mod$W[2,2]=0} # evolution (process error) variances
    if(alpha_vary==FALSE & beta_vary==TRUE){mod$W[1,1]=0; mod$W[2,2]=exp(parm[2])}
    if(alpha_vary==TRUE & beta_vary==TRUE){mod$W[1,1]=exp(parm[2]); mod$W[2,2]=exp(parm[3])}
    return(mod)
  }

  # 4.  maximum likelihood optimization of the variance
  dlm_out<-dlmMLE(y=lnRS, build=build_mod, parm=c(-.1,-.1,-.1), method="Nelder-Mead")

  # 5. log-likelihood
  lls <- dlm_out$value

  # 6. specify the model based on variance structure
  dlmMod <- build_mod(dlm_out$par)

  # 7. apply Kalman filter
  outsFilter <- dlmFilter(y=lnRS,mod=dlmMod)

  # 8. backward recursive smoothing
  outsSmooth	<-dlmSmooth(outsFilter)

  # 9. grab parameters, their SEs and calculate AICc
  alpha<- cbind(alpha,outsSmooth$s[-1,1,drop=FALSE])
  beta<- cbind(beta,outsSmooth$s[-1,2,drop=FALSE])
  alpha_se <- sqrt(array(as.numeric(unlist(dlmSvd2var(outsSmooth$U.S, outsSmooth$D.S))), dim=c( 2, 2,length(lnRS)+1)))[1,1,-1]
  beta_se <- sqrt(array(as.numeric(unlist(dlmSvd2var(outsSmooth$U.S, outsSmooth$D.S))), dim=c( 2, 2,length(lnRS)+1)))[2,2,-1]
  AICc	<- 2*lls + 2*dlmPars +(2*dlmPars*(dlmPars+1)/(length(data$rec)-dlmPars-1))
  BIC <-BIC <- 2*lls + dlmPars*log((length(data$rec)))

  # 10. output results
  results <- cbind(data,alpha, beta,alpha_se,beta_se)
  output <- list(results=results,AICc=AICc, BIC=BIC)

}

## plotDLM function
#
# dlm_model: a list object generated by fitDLM

plotDLM <- function(dlm_model = dlm_model){
  # pre-plotting housekeeping
  spwn_range<-seq(0,max(dlm_model$results$spwn)*1.2,length.out=100)
  r_pred <- matrix(NA,nrow=length(dlm_model$results$byr), ncol=length(spwn_range))

  for(i in 1:length(dlm_model$results$byr)){
    r_pred[i,]<-exp(median(dlm_model$results$alpha[i]))*spwn_range*exp(dlm_model$results$beta[i]*spwn_range)
  }

  rownames(r_pred)<-dlm_model$results$byr
  colnames(r_pred)<-spwn_range

  r_pred<-cbind(dlm_model$results$byr,r_pred)
  colnames(r_pred)[1]<-c("byr")
  sr_pred<-pivot_longer(data =as.data.frame(r_pred), cols=!byr, names_to="spwn",values_to="rec" )
  sr_pred$spwn<-as.numeric(sr_pred$spwn)

  max_spawn <- max(dlm_model$result$spwn)
  max_rec <- max(dlm_model$result$spwn)

  alpha_range<-c(min(dlm_model$results$alpha-(dlm_model$results$alpha_se*2)*2),max(dlm_model$results$alpha+(dlm_model$results$alpha_se*2)*2))
  beta_range<-c(min(dlm_model$results$beta-(dlm_model$results$beta_se*2)*2),max(dlm_model$results$beta+(dlm_model$results$beta_se*2)*2))

  # create each panel for plot starting with spawner-recruitment relationship

  a <- ggplot(data=dlm_model$results, aes(x = spwn, y = rec, colour=byr))+
    geom_point(size = 3) +
    coord_cartesian(xlim=c(0, max_spawn*1.2), ylim=c(0,max_rec*1.2)) +
    scale_colour_viridis_c()+
    xlab("Spawners") +
    ylab("Recruits") +
    geom_line(data = sr_pred, aes(x = spwn, y = rec, group=byr, colour=byr), size = 0.5) +
    theme_bw() +
    labs(color='Year') +
    theme(strip.text.x = element_text(size=8),
          axis.title = element_text(size=10),
          axis.text = element_text(size=7),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          legend.justification = c(0,0),
          legend.position = c(0.1,0.765),
          legend.key.size = unit(6, "pt"),
          legend.background = element_blank(),
          legend.text = element_text(size = 6),
          legend.title = element_text(size = 7),
          plot.margin=unit(c(0.5,0.5,0.5,0.5), units="lines"))


  theme(strip.text = element_text(size=6),
        axis.title = element_text(size=9),
        axis.text = element_text(size=6),

        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank())











  # next plot true and estimate alpha

  b <- ggplot(data=dlm_model$results, aes(x = byr, y = alpha )) +
    geom_line( color="black", size = 1)+
    geom_ribbon(aes(ymin = alpha-alpha_se*2, ymax = alpha+alpha_se*2),
                fill = "grey80", alpha=0.5, linetype=2, colour="gray46") +
    geom_line(aes(x = byr, y = alpha_true), color="red", size = 1)+
    ylab("alpha") +
    xlab("") +
    scale_y_continuous(position = "right") +
    coord_cartesian(ylim=c(alpha_range[1],alpha_range[2])) +
    theme_bw() +
    theme(strip.text.x = element_text(size=8),
          axis.title = element_text(size=10),
          axis.text = element_text(size=7),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          legend.position = "none",
          plot.margin=unit(c(0.5,2.25,0.05,0.5), units="lines"))+
    annotate("text", x = c(30,30),
             y = c(alpha_range[2]*0.95, alpha_range[2]*0.85),
             label = c("True", "Estimated"),
             color= c("red", "dark grey"),
             size=3,
             hjust=0) +
    annotate("segment", x = c(28,28),
             xend=c(29.5,29.5),
             y = c(alpha_range[2]*0.95, alpha_range[2]*0.85),
             yend = c(alpha_range[2]*0.95, alpha_range[2]*0.85),
             lty = c(1,1),
             color=c("red", "dark grey"),
             size=1)

  # next plot true and estimate beta

  c<- ggplot(data=dlm_model$results, aes(x = byr, y = beta)) +
    geom_line( color="black", size = 1)+
    geom_ribbon(aes(ymin = beta-beta_se*2, ymax = beta+beta_se*2),
                fill = "grey80", alpha=0.5, linetype=2, colour="gray46") +
    geom_line(aes(x = byr, y = -beta_true), color="red", size = 1)+
    xlab("Brood year") +
    ylab("beta") +
    scale_y_continuous(position = "right" ) +
    coord_cartesian(ylim=c(beta_range[1],beta_range[2])) +
    theme_bw() +
    theme(strip.text.x = element_text(size=8),
          axis.title = element_text(size=10),
          axis.text = element_text(size=7),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          legend.position = "none",
          plot.margin=unit(c(0.05,0.5,0.5,0.5), units="lines"))

  # combine plots
  g <- ggarrange(a, ggarrange(b,c, nrow =2),
                 ncol=2)

  g
}