pgam_glmmtmb.R

library(glmmTMB)
library(mgcv)

#########################################################################
## Function to create prediction matrices from mgcv for glmmTMB
##   from ?mgcv::smooth2random
########################################################################
s2rPred <- function(sm,re,data) {
## Function to aid prediction from smooths represented as type==2
## random effects. re must be the result of smooth2random(sm,...,type=2).
  X <- PredictMat(sm,data)   ## get prediction matrix for new data
  ## transform to r.e. parameterization
  if (!is.null(re$trans.U)) X <- X%*%re$trans.U
  X <- t(t(X)*re$trans.D)
  ## re-order columns according to random effect re-ordering...
  X[,re$rind] <- X[,re$pen.ind!=0] 
  ## re-order penalization index in same way  
  pen.ind <- re$pen.ind; pen.ind[re$rind] <- pen.ind[pen.ind>0]
  ## start return object...
  r <- list(rand=list(),Xf=X[,which(re$pen.ind==0),drop=FALSE])
  for (i in 1:length(re$rand)) { ## loop over random effect matrices
    r$rand[[i]] <- X[,which(pen.ind==i),drop=FALSE]
    attr(r$rand[[i]],"s.label") <- attr(re$rand[[i]],"s.label")
  }
  names(r$rand) <- names(re$rand)
  r
} ## s2rPred


## Continuous covariate
x <- seq(1,10, length=100)
x.pred <-  seq(1,15, length=150)

## Set up penalized thin-plate regression spline for x
sm <- mgcv::smoothCon(s(x), data=as.data.frame(x))[[1]]
re <- mgcv::smooth2random(sm, "", type=2)
Xf <- re$Xf
Xr <- re$rand$Xr

## Simulated data set
set.seed(1234)
b <- rnorm(8,0,3)
beta <- c(0, 2)
lambda <- exp(Xf%*%beta + Xr%*%b)
y <- rpois(length(x),lambda)
plot(x,y)

## Fit with glmmTMB
k <- ncol(Xr)
# make a fake grouping variable
g <- rep(1, length(y)) 
## Set up TMB 'map' argument to have only 1 variance parameter for Xr 
##  This is the TPRS smoothing parameter
tmb_map <- list(theta = factor(c(rep(1, k), rep(NA, k*(k-1)/2))))

ftmb <- glmmTMB(y ~ 0+Xf +(0+Xr|g), family=poisson, map=tmb_map)
pred.obj <- s2rPred(sm, re, data.frame(x=x.pred))
newdata <- list(Xf=pred.obj$Xf, Xr=pred.obj$rand$Xr, g=rep(1, nrow(pred.obj$Xf)))

ptmb <- as.data.frame(predict(ftmb,newdata=newdata, se=T)) 


## Fit with mgcv::gam, method='REML'
fmgcv <- mgcv::gam(y~s(x), family=poisson, method="REML")
pmgcv <- as.data.frame(predict(fmgcv, newdata=data.frame(x=x.pred), se=TRUE, uncond=TRUE)) 

## Compare fits
plot(x, y, xlim=c(0,15))

lines(x.pred, exp(ptmb[,1]), col='blue', lwd=2)
lines(x.pred, exp(ptmb[,1]+1.96*ptmb[,2]), col='blue', lty=3, lwd=2)
lines(x.pred, exp(ptmb[,1]-1.96*ptmb[,2]), col='blue', lty=3, lwd=2)
text(1, 33, labels=c("glmmTMB"), col='blue', pos=4, cex=2)

lines(x.pred, exp(pmgcv[,1]), col='red', lwd=2)
lines(x.pred, exp(pmgcv[,1]+1.96*pmgcv[,2]), col='red', lty=3, lwd=2)
lines(x.pred, exp(pmgcv[,1]-1.96*pmgcv[,2]), col='red', lty=3, lwd=2)
text(1, 30, labels=c("mgcv::gam"), col='red', pos=4, cex=2)

😜 🧹