# hdrs-plots.R
#
# Example on growth curve models for Reisby et al. (1977)
#
# content: (1) Read data
#          (2) Linear analysis
#          (3) Fitting quadratic model
#          (3) Centering of variables
#          (4) Check random effects structure
#          (5) Check model assumptions
#
# input:  resiby.txt
# output: hdrs-ind_pred-randomslope.pdf
#         hdrs-ind_pred-quad.pdf
#         hdrs-caterpillar.pdf
#         hdrs_shrinkage_int-week.pdf
#         hdrs_shrinkage_int-weeksq.pdf
#         hdrs_shrinkage_week-weeksq.pdf
#
# last mod: 2026-04-20, NW

library("lme4") 
library("lattice")

#--------------- (1) Read data ------------------------------------------------

dat      <- read.table("../data/reisby.txt", header = TRUE)
dat$id   <- factor(dat$id)
dat$diag <- factor(dat$diag, levels = c("nonen", "endog"))
dat      <- na.omit(dat)    # drop missing values

pdf("hdrs-box.pdf", height = 3.375, width = 3.375, pointsize = 10)

par(mai = c(.6, .6, .1, .1), mgp = c(2, .4, 0))
boxplot(hamd ~ week, dat,
        col = "#3CB4DC",
        ylab = "HDRS score",
        xlab = "Time (week)")

dev.off()

#--------------- (2) Linear analysis ------------------------------------------

# random slope model
lme2 <- lmer(hamd ~ week + (week | id), dat, REML = FALSE)

pdf("hdrs-ind_pred-randomslope.pdf", height = 6, width = 6)

xyplot(hamd + predict(lme2) ~ week | id, data = dat,
       type = c("p", "l", "g"),
       #pch = 16,
       ratio = "xy",
       col = c("#3CB4DC", "#FF6900"),
       distribute.type = TRUE,
       layout = c(11, 6),
       ylab = "HDRS score",
       xlab = "Time (week)")

dev.off()

#--------------- (3) Fitting quadratic model ----------------------------------

# model with quadratic time trend
lme3 <- lmer(hamd ~ week + I(week^2) + (week + I(week^2) | id), dat,
             REML = FALSE)

pdf("hdrs-ind_pred-quad.pdf", height = 6, width = 6)

xyplot(hamd + predict(lme3) ~ week | id, data = dat,
       type = c("p", "l", "g"),
       #pch = 16,
       ratio = "xy",
       col = c("#3CB4DC", "#FF6900"),
       distribute.type = TRUE,
       layout = c(11, 6),
       ylab = "HDRS score",
       xlab = "Time (week)")

dev.off()

#--------------- (4) Check random effects structure ---------------

pdf("hdrs-caterpillar.pdf", height = 8, width = 20)

dotplot(ranef(lme3), col = "#3CB4DC",
        scales = list(x = list(relation = "free")))[[1]]

dev.off()

# Shrinkage plots

df     <- coef(lmList(hamd ~ week_c + I(week_c^2) | id, dat))
cc1    <- as.data.frame(coef(lme3reml)$id)
names(cc1) <- c("A", "B", "C")

df <- cbind(df, cc1)
ff <- fixef(lme3reml)


pdf("hdrs_shrinkage_int-week.pdf", height = 6, width = 6, pointsize = 10)

with(df,
  xyplot(`(Intercept)` ~ week_c, aspect = 1,
    x1 = B, y1 = A,
    panel = function(x, y, x1, y1, subscripts, ...) {
        panel.grid(h = -1, v = -1)
        x1 <- x1[subscripts]
        y1 <- y1[subscripts]
        larrows(x, y, x1, y1, type = "closed", length = 0.1, fill = "black",
                angle = 15, ...)
        lpoints(x, y,
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[2])
        lpoints(x1, y1,
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[1])
        lpoints(ff[2], ff[1], 
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[3])
    },
    xlab = "week_c",
    ylab = "(Intercept)",
    key = list(space = "top", columns = 3,
               text = list(c("Mixed model", "Within-subject", "Population")),
               points = list(col = trellis.par.get("superpose.symbol")$col[1:3],
                             pch = 16))
  )
)

dev.off()


pdf("hdrs_shrinkage_int-weeksq.pdf", height = 6, width = 6, pointsize = 10)

with(df,
  xyplot(`(Intercept)` ~ `I(week_c^2)`, aspect = 1,
    x1 = C, y1 = A,
    panel = function(x, y, x1, y1, subscripts, ...) {
        panel.grid(h = -1, v = -1)
        x1 <- x1[subscripts]
        y1 <- y1[subscripts]
        larrows(x, y, x1, y1, type = "closed", length = 0.1, fill = "black",
                angle = 15, ...)
        lpoints(x, y,
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[2])
        lpoints(x1, y1,
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[1])
        lpoints(ff[3], ff[1], 
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[3])
    },
    xlab = expression(week_c^2),
    ylab = "(Intercept)",
    key = list(space = "top", columns = 3,
               text = list(c("Mixed model", "Within-subject", "Population")),
               points = list(col = trellis.par.get("superpose.symbol")$col[1:3],
                             pch = 16))
  )
)

dev.off()


pdf("hdrs_shrinkage_week-weeksq.pdf", height = 6, width = 6, pointsize = 10)

with(df,
  xyplot(week_c ~ `I(week_c^2)`, aspect = 1,
    x1 = C, y1 = B,
    panel = function(x, y, x1, y1, subscripts, ...) {
        panel.grid(h = -1, v = -1)
        x1 <- x1[subscripts]
        y1 <- y1[subscripts]
        larrows(x, y, x1, y1, type = "closed", length = 0.1, fill = "black",
                angle = 15, ...)
        lpoints(x, y,
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[2])
        lpoints(x1, y1,
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[1])
        lpoints(ff[3], ff[2], 
                pch = 16,
                col = trellis.par.get("superpose.symbol")$col[3])
    },
    xlab = expression(week_c^2),
    ylab = "week_c",
    key = list(space = "top", columns = 3,
               text = list(c("Mixed model", "Within-subject", "Population")),
               points = list(col = trellis.par.get("superpose.symbol")$col[1:3],
                             pch = 16))
  )
)

dev.off()