nimble_issj.R

# setting up a Nimble model for the island scrub jay data
# David Lawrence Miller 2020

# data provided as part of the unmarked package
# Sillett, T. S., Chandler, R. B., Royle, J. A., Kéry, M., & Morrison, S. A. (2012). Hierarchical distance-sampling models to estimate population size and habitat-specific abundance of an island endemic. Ecological Applications, 22(7), 1997–2006. https://doi.org/10.1890/11-1400.1

# this model uses mgcv::jagam to set up a spatial smooth and a custom
# point transect detection function coded in Nimble

# Thanks to Christopher Paciorek for help setting-up user functions in Nimble

library(dsm)
library(mgcv)
library(Distance)
library(nimble)

# load scrub jay data
# this has been slightly re-formatted
load("issj.RData")

## 1. setup smoothers and related data structures using dsm and mgcv::jagam

# fit detection function to fall data
df_fit <- ds(obs_fall, transect="point", formula=~chaparral)

# fit spatial model, need this to setup the data for the next call
prell <- dsm(count~s(x,y, k=40, bs="ts"),
             observation.data=obs_fall, segment.data=segs, ddf.obj=df_fit,
             transect="point", family=nb())

# use jagam to generate template JAGS code for smoothers
# precomputed model components from jagssetup.R
# includes penalties, data, design matrices etc
ll <- jagam(count~s(x,y, k=40, bs="ts"), data=prell$data, family=poisson,
            file="jagstemplate.jags")

library(ggplot2)
ggplot(prell$data) +
  geom_point(aes(x=x, y=y, size=count)) +
  coord_equal() +
  theme_minimal()


## 2. specify the Nimble model

# inline the model code
# smooth parts of the code generated using mgcv::jagam
# pasted from jagstemplate.jags, generated above
model_code <- nimbleCode({
  # priors on parameters
  # intercept for scale
  beta0 ~ dnorm(0, 0.0001)
  # chaparral coef
  beta1 ~ dnorm(0, 0.0001)

  # detection function
  for(j in 1:n_obs){
    obs_distance[j] ~ dhnpt(beta0, beta1, obs_chaparral[j], width)
  }

  # prior on negbin par
  r ~ dunif(0.00001, 4)

  eta[] <- X[,] %*% b[] ## linear predictor

  # build the mean
  for (i in 1:n) {
    # calculate effective area for this segment
    # calculate sigma
    sigma_seg[i] <- exp(beta0 + beta1*seg_chaparral[i])
    # (Intro distance book eqn 3.45)
    nu[i] <- 2 * pi * sigma_seg[i]^2 *
              (1 - exp(-(width^2)/(2*sigma_seg[i]^2)))

    y[i] ~ dnegbin(prob=p[i], size=r)
    p[i] <- r/(r+mu[i])
    mu[i] <- nu[i]*exp(eta[i])
  }

  ## intercept prior CHECK tau=1/24^2 is appropriate!
  b[1] ~ dnorm(0,0.00018)
  ## prior for s(x,y)...
  K1[1:39,1:39] <- S1[1:39,1:39] * lambda[1]
  b[2:40] ~ dmnorm(zero[2:40],K1[1:39,1:39])
  ## smoothing parameter priors CHECK...
  for (i in 1:1) {
    lambda[i] ~ dgamma(.05,.005)
    rho[i] <- log(lambda[i])
  }

  # calculate abundance over the prediction grid
  Nhat[] <- areas[]*exp(Lp[,] %*% b[])
  Nhat_total <- sum(Nhat[])
})

# define half-normal point transect detection function pdf
dhnpt <<- nimbleFunction(
    run = function(x = double(0), b0 = double(0), b1 = double(0),
                   covar = double(0), width = double(0),
                   log = integer(0, default = 0)) {
        returnType(double(0))

      # calculate scale parameter
      sigma <- exp(b0 + b1*covar)

      # analytic expression for integral of 2*r*g(r)/width^2 when
      #  g(r) is half-normal (Intro distance book eqn 3.45)
      nu <- sigma^2 * (1 - exp(-(width^2)/(2*sigma^2)))

      # evaluate the detection function at distance/chaparral combination
      g <- exp(-(x^2)/(2*sigma^2))

      # pdf
      L <- (x * g)/nu

      if(log) return(log(L))
      else return(L)
    }
)
rhnpt <<- nimbleFunction(
    run = function(n = integer(0), b0 = double(0), b1 = double(0),
                   covar = double(0), width = double(0),
                   log = integer(0, default = 0)) {
        returnType(double())
        return(0)
    }
)
nimble::registerDistributions(list(
  dhnpt = list(BUGSdist="dhnpt(b0, b1, covar, width)",
               pqAvail = FALSE,
               range = c(0, 1000))))


## 3. MCMC setup

# parameters to be monitored
varsToMonitor <- c("rho", "Nhat_total", "beta0", "beta1", "r")

# MCMC settings
nitt <- 500000
burnin <- 10000
thin <- 10

# data for jags is generated by mgcv::jagam()
# add extra bits here
nimble_data <- ll$jags.data
nimble_data$Lp <- predict(prell, cruz, type="lpmatrix")
nimble_data$width <- df_fit$ddf$meta.data$width
nimble_data$obs_distance  <- obs_fall$distance
nimble_data$obs_chaparral <- obs_fall$chaparral
nimble_data$seg_chaparral <- segs$chaparral
nimble_data$y             <- prell$data$count


# again, initial values are found by mgcv::jagam
# but add some extra ones here
# here using the values from mgcv/dsm
nimble_inits <- ll$jags.ini
nimble_inits$beta0 <- df_fit$ddf$par[1]
nimble_inits$beta1 <- df_fit$ddf$par[1]
nimble_inits$r <- prell$family$getTheta(TRUE)
nimble_constants <- list()
nimble_constants$n_obs <- nrow(obs_fall)
nimble_constants$n <- nrow(segs)
nimble_constants$areas <- cruz$off.set
nimble_constants$pi <- pi

# declare variable dimensions when nimble complains
nimble_dims <- list(b    = 40,
                    eta  = nrow(nimble_data$X),
                    Nhat = nrow(nimble_data$Lp))


## 4. run the model

# build and compile the model
nm <- nimbleModel(model_code, data=nimble_data,
                  constants=nimble_constants,
                  dimensions=nimble_dims)
# build the MCMC samplers for the compiled model
my_mcmc <- buildMCMC(nm,
                     monitors=c("rho", "lambda", "Nhat_total", "r"),
                     thin=thin)
# compile our model
nm_cc <- compileNimble(nm)

# compile the MCMC, using the nm "project"
cc <- compileNimble(my_mcmc, project=nm)

# actually do some sampling
samples <- runMCMC(cc, inits=nimble_inits,
                   niter=nitt, nburnin=burnin)


# look at MCMC output
library(coda)
plot(as.mcmc(samples))

# calculate some simple summary stats

# estimate of Nhat
median(samples[,1])
# CV
sd(samples[,1])/median(samples[,1])