# depression.R
#
# content: (1) Read data
#          (2) Visualize data
#          (3) Fit growth curve model
#
# input: data/reisby.txt
# output: --
#
# 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

boxplot(hamd ~ week, dat,
        col = "#3CB4DC",
        ylab = "HDRS score",
        xlab = "Time (week)")

#----- (2) Fitting random slope model -----------------------------------------

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

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

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

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

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

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

# Catterpillar plots

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

# 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)


## shrinkage intercept and week
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))
  )
)

## shrinkage intercept and week^2
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))
  )
)

## shrinkage week and week^2
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))
  )
)