# Supporting file 2

# Schooler, S.L., N. J. Svoboda, S.P. Finnegan, J. Crye, K.F. Kellner,
# and Jerrold L. Belant. Maternal carryover, winter severity, and 
# brown bear abundance relate to elk demographics. Plos one.

# Programmer: Sarah L. Schooler and Kenneth F. Kellner, 2022
# R version 4.2.0; JagsUI package version 1.5.2

#####################################################################################
# Annotated code for casting the linear model and state-space model using JagsUI    #
#####################################################################################

###################
# Linear model
###################

require(jagsUI)

# y: vector of ratio of calves per 100 adult females (age ratio)
# winterpred: vector of scaled winter severity predictor
# springpred: vector of scaled spring bottleneck predictor
# maternalpred: vector of scaled maternal carryover predictor
# bearpred: vector of scaled brown bear abundance predictor
# timberpred: vector of scaled area of timber harvest predictor
# Note: timberpred only used in timber harvest and island-wide recruitment models

sink("LinearRecruitment.jags")
cat("
model {

# Priors
 # Uniformative prior for variance, tau is 1/variance
 sigma ~ dunif(0, 100)
 tau <- 1/(sigma * sigma)
 
 # Uninformative priors for parameters
 B0 ~ dnorm(0,0.001)
 B1 ~ dnorm(0,0.001)
 B2 ~ dnorm(0, 0.001)
 B3 ~ dnorm(0,0.001)
 B4 ~ dnorm(0,0.001)
 B5 ~ dnorm(0,0.001)


# Likelihood
 for (t in 1:nYears) {
    # Variance
    y[t] ~ dnorm(mu[t], tau)
    
    # Linear model
    mu[t] <- B0 + B1*winterpred[t] + B2*springpred[t] + B3*maternalpred[t] + 
             B4*bearpred[t] + B5*timberpred[t]
 } 


# Assess model fit using a sums-of-squares-type discrepancy
 for (i in 1:nYears) {
    residual[i] <- y[i]-mu[i]		# Residuals for observed data
    predicted[i] <- mu[i]		# Predicted values
    sq[i] <- pow(residual[i], 2)	# Squared residuals for observed data

# Beyesian posterior predictive check; checks perfect data set
    y.new[i] ~ dnorm(mu[i], tau) # One new data set at each MCMC iteration
    sq.new[i] <- pow(y.new[i]-predicted[i], 2)	# Squared residuals for new data
 }
 fit <- sum(sq[])			# Sum of squared residuals for actual data set
 fit.new <- sum(sq.new[])		# Sum of squared residuals for new data set
 test <- step(fit.new - fit)		# Compare datasets
}
",fill=TRUE)
sink()



# Bundle data
jags.data <- list(y = AgeRatios, nYears = length(AgeRatios), winterpred = WinterSeverityPredictor, 
                  maternalpred = MaternalCarryoverPredictor, springpred = SpringBottleneckPredictor,
                  bearpred = BrownBearAbundance, timberpred = AreaofTimberHarvest)  

jags.data <- list(y = AgeRatios, nYear = length(AgeRatios), winterpred = scale(WinterSeverityPredictor), 
                  matpred = scale(MaternalCarryoverPredictor), sprpred = scale(SpringBottleneckPredictor),
                  bearpred = scale(BrownBearAbundance), timberpred = scale(AreaofTimberHarvest))

# Parameters to estimate
params <- c("B0", "B1", "B2",  "B3", "B4", "B5", "sigma", "fit", "fit.new")

# Start Gibbs sampler
out.model <- jagsUI::jags(data = jags.data, parameters.to.save = params, 
                          model.file = "LinearRecruitment.jags", 
                          codaOnly = c("fit", "fit.new"),
                          n.thin = 5, n.chains = 3, n.burnin = 10000, n.iter = 50000)
print(out.model, dig = 3)

# Fit tests
plot(out.model$sims.list$fit, out.model$sims.list$fit.new, main = "Graphical 
     posterior predictive check", las = 1, xlab = "SSQ for actual data set", 
     ylab = "SSQ for ideal (new) data sets")
abline(0, 1)
mean(out.model$sims.list$fit.new > out.model$sims.list$fit)	# Bayesian p-value, should be 0.5
traceplot(out.model, parameters = c("sigma", "B0", "B1",
                                    "B2", "B3", "B4"))


#############################
#############################
#############################
#############################

###################
# State-space model
###################

# y: vector of population size count data
# H: vector of harvest numbers
# winterpred: vector of scaled winter severity predictor
# springpred: vector of scaled spring bottleneck predictor
# maternalpred: vector of scaled maternal carryover predictor

sink("StateSpaceAbundance.jags")
cat("
    model{
       # Priors
       
       # Initial unobserved log population size as a minimally informed prior of 
       # U(log(min(y), log(max(y)))
       logN.init ~ dunif(2, 6) 
       logN.exp[1] <- logN.init
       
       logN.est[1] ~ dnorm(logN.exp[1], tau.p) # Initial estimated population size
       y[1] ~ dnorm(logN.est[1], tau.obs) # Initial log(population size)
       
       # Priors for variance error; tau is 1/variance
       sigma.p ~ dunif(0, 10)            # Process error
       tau.p <- 1/(sigma.p*sigma.p)
       
       sigma.obs ~ dunif(0, 10)          # Observation error
       tau.obs <- 1/(sigma.obs*sigma.obs)
       
       # Uninformative priors for parameters
       B0 ~ dnorm(0,0.001)
       B1 ~ dnorm(0,0.001)
       B2 ~ dnorm(0,0.001)
       B3 ~ dnorm(0,0.001)
       B4 ~ dnorm(0,0.001)

       # Likelihood
       for (t in 1:(nYears-1)){
          # Observation process
          y[t+1] ~ dnorm(logN.est[t+1], tau.obs) 
          
          # State process
          # Solve for logarithmic extraction term
          h[t] <- log(abs(1-H[t]/exp(logN.est[t])))
          # Rate of change of population size as a function of density dependence and predictors
          r[t] <- B0 + B1*(logN.est[t]+h[t]) + B2*springpred[t] + 
                           B3*maternalpred[t] + B4*winterpred[t]
          # Expected value of log(N[t]) including harvest
          logN.exp[t+1] <- logN.est[t] + h[t] + r[t]
          # Estimated value of log(Nt) with process error
          logN.est[t+1] ~ dnorm(logN.exp[t+1], tau.p)
       }
       
       
       # Population sizes on real scale
       for (t in 1:nYears){
          N.est[t] <- exp(logN.est[t])
       }

       # Assess model fit using a sums-of-squares-type discrepancy
       for (i in 1:nYears) {
          residual[i] <- y[i]-logN.est[i]		# Residuals for observed data
          predicted[i] <- logN.est[i]		# Predicted values
          sq[i] <- pow(residual[i], 2)	# Squared residuals for observed data

          # Generate replicate data and compute fit statistics
          logN.est.new[i] ~ dnorm(logN.exp[i], tau.p)
          y.new[i] ~ dnorm(logN.est[i], tau.obs) # One new data set at each MCMC iteration
          sq.new[i] <- pow(y.new[i]-predicted[i], 2)	# Squared residuals for new data
       }
       fit <- sum(sq[])			# Sum of squared residuals for actual data set
       fit.new <- sum(sq.new[])		# Sum of squared residuals for new data set
       test <- step(fit.new - fit)		# Test the difference
}
",fill=TRUE)
sink()


# Bundle data

jags.data <- list(y = log(count), nYears = length(count), H = harvest,
                  springpred = SpringBottleneckPredictor, winterpred = WinterSeverityPredictor, 
                  maternalpred = MaternalCarryoverPredictor)


# To avoid invalid parent value errors, need to initialize the log true abundance X
# at some reasonable values
inits <- function() list(logN.init = log(count[1]))


params <- c("B0", "B1", "B2", "B3", "B4", "sigma.obs", "sigma.p", "N.est", "r",
                 "fit", "fit.new")

out.model <- jagsUI::jags(data = jags.data, parameters.to.save = params, 
                             inits = inits, model.file = "StateSpaceAbundance.jags", 
                         codaOnly = c("fit", "fit.new", "r", "Nt"),
                         n.thin = 5, n.chains = 3, n.burnin = 10000, n.iter = 50000)
print(out.model, dig = 3)

# Fit tests
plot(out.model$sims.list$fit, out.model$sims.list$fit.new, main = "Graphical 
     posterior predictive check", las = 1, xlab = "SSQ for actual data set", 
     ylab = "SSQ for ideal (new) data sets")
abline(0, 1)
mean(out.model$sims.list$fit.new > out.model$sims.list$fit)	# Bayesian p-value, should be 0.5
traceplot(out.model, parameters = c("sigma.obs", "sigma.p", "B0", "B1",
                                   "B2", "B3", "B4"))