# Compare HFS and original Harville iterations with lme
# Supporting material for the book 'Practical Smoothing. The Joys of P-splines'
# Paul Eilers and Brian Marx, 2019

library(nlme)
library(spam)
library(JOPS)
library(MASS)

# Simulate data
data(mcycle)
x = mcycle$times
y = mcycle$accel
m = length(y)
xmin = min(x)
xmax = max(x)

# Compute basis functions
nseg = 20
B = bbase(x, xl = xmin, xr = xmax, nseg = nseg, bdeg = 3)
n = ncol(B)
E = diag(n)
d = 2
D = diff(E, diff = d)
P = t(D) %*% D
X = cbind(1, x)
Z = B %*% t(D) %*% solve(D %*% t(D))

# Initialize
sig2 =1
tau2 = 1
P = 0 * diag(n)
jj = 3 :n
nit = 10
mon = T   # Monitor convergence if true

# lme
grp = 0 * x + 1
t0 = Sys.time()
lm1 <- lme(y ~ x, random = list(grp = pdIdent(~Z - 1)))
cfix = fixef(lm1)
cran = unlist(ranef(lm1))
zlm = X %*% cfix + Z %*% cran
sig2b = lm1$sigma ^ 2
tau2b = as.numeric(VarCorr(lm1)[1, 1])
ed = sum(cran ^ 2) / tau2b + 2
cat('Linear mixed model (lme, REML): ED, sig2, tau2, time\n')
cat(ed, sig2b, tau2b, round(Sys.time() - t0, 4), 's\n\n')

# Original Harville iterations
cat('Harville iterations: ED, sig2, tau2\n')
sig2 =1
tau2 = 1
t0 = Sys.time()
for (it in 1:nit) {
  R = diag(m) * sig2
  G = diag(n - 2) * tau2
  Ri = solve(R)
  Gi = solve(G)
  A11 = t(X) %*% Ri %*% X
  A12 = t(X) %*% Ri %*% Z
  A21 = t(A12)
  A22 = t(Z) %*% Ri %*% Z
  S = Ri - Ri %*% X %*% solve(A11) %*% t(X) %*% Ri
  A = rbind(cbind(A11, A12), cbind(A21, A22))
  P[jj, jj] = Gi
  r1 = t(X) %*% Ri %*% y
  r2 = t(Z) %*% Ri %*% y
  r = c(r1, r2)
  a = solve(A + P, r)
  z = cbind(X, Z) %*% a
  T1 = solve(diag(n - 2) + t(Z) %*% S %*% Z %*% G)
  ed = n - sum(diag(T1))
  a2 = a[jj]
  taunew = sum(a2 ^ 2) / (ed - d)
  signew = sum((y - z) ^ 2) / (m - ed)
  sig2 = signew
  tau2 = taunew
  if (mon) cat(it, ed, sig2, tau2, '\n')
}
cat('Original Harville:',  round(Sys.time() - t0, 4), 's\n\n')
 
# Mixed model with SChall
cat('Mixed model with Schall iterations: ED, sig2, tau2\n')
C = cbind(X, Z)
nc = ncol(C)
P = diag(nc)
P[1:2, 1:2] = 0
kappa = 1
t0 = Sys.time()
for (it in 1: nit) {
  a = solve(t(C) %*% C + kappa * P, t(C) %*% y)
  z = C %*% a
  G = solve(t(C) %*% C + kappa * P, t(C) %*% C)
  ed = sum(diag(G)) 
	sig2 = sum((y - z) ^2 )/ (m - ed)
  tau2 = sum(a[3:n]^ 2) / (ed - 2)
  kappa = sig2 / tau2
	if (mon) cat(it, ed, sig2, tau2, kappa, '\n')
}
cat('Mixed model with Schall:', round(Sys.time() - t0, 4), 's\n\n')

# Random model with HFS (Harville-Fellner-Schall) algorithm
cat('Random model HFS iterations: ED, sig2, tau2\n')
lambda = 1
BtB = t(B) %*% B
Bty = t(B) %*% y
n = ncol(B)
D = diff(diag(n), diff = d)
P = t(D) %*% D
t0 = Sys.time()
for (it in 1:nit) {
    a = solve(BtB + lambda * P, Bty)
    G = solve(BtB + lambda * P, BtB)
    ed = sum(diag(G))
    yhat = B %*% a
    sig2 = sum((y - yhat) ^ 2) / (m - ed)
    tau2 = sum((D %*% a) ^2) / (ed - d)
    lambda = sig2 / tau2
  	if (mon) cat(it, ed, sig2, tau2, lambda, '\n')
}   
z = yhat  
cat('Random with HFS:', round(Sys.time() - t0, 4), 's\n\n')