########## R script: curves with random intercepts
#
## Last modified December 1, 2003 by John Staudenmayer
## email: jstauden@math.umass.edu
##

library("nlme") 

# Set parameters
n <- 100  # number of data points
sig <- 1 # std dev of regression error
num.add.comps <- 2 # number of additive components

# number of knots in splines
num.knots <- c(min(30,floor(n/3)),min(30,floor(n/3)))

# Generate data (i.e. scatterplot of x1 versus y)
x1 <- runif(n)
x2 <- runif(n)
eps <- rnorm(n)
f1 <- function(x)
  {
    val <- (1.5*dnorm(x,.3,.1)-dnorm(x,.75,0.07))
    val <- val - mean(val)
    return(val)
  }
f2 <- function(x)
  {
    val <- (20*(x-.5)^2)+ 2*x
    val <- val - mean(val)
    return(val)
  }

y <- f1(x1) + f2(x2) + sig*eps

# Set up design matrices and random effects block structure
knots <- list()
knots[[1]] <- quantile(unique(x1),
                       seq(0,1,
                           length=(num.knots[1]+2))[-c(1,(num.knots[1]+2))])
knots[[2]] <- quantile(unique(x2),
                       seq(0,1,
                           length=(num.knots[2]+2))[-c(1,(num.knots[2]+2))])

Z.inds <- list()
Z.inds[[1]] <-
  (num.add.comps+2):(num.knots[1]+(num.add.comps+1))
Z.inds[[2]] <-
  (Z.inds[[1]][num.knots[1]]+1):(sum(num.knots)+(num.add.comps+1))


X <- cbind(rep(1,n),x1,x2)

Z.1 <- outer(x1,knots[[1]],"-")
Z.1 <- Z.1*(Z.1>0)
Z.1 <- cbind(Z.1)

Z.2 <- outer(x2,knots[[2]],"-")
Z.2 <- Z.2*(Z.2>0)
Z.2 <- cbind(Z.2)

C.mat <- cbind(X,Z.1,Z.2)
cols <- dim(C.mat)[2]
rescale.mean <- apply(C.mat[,2:cols],2,mean)

C.mat[,2:cols] <- scale(C.mat[,2:cols],scale=F)
X <- cbind(rep(1,length=n),C.mat[,2:(num.add.comps+1)])

Z.1 <- C.mat[,Z.inds[[1]]]
Z.2 <- C.mat[,Z.inds[[2]]]

re.block.val <- list(1:num.knots[1],1:num.knots[2])

Z.block  <-  list()
Z.block[[1]] <-
  as.formula(paste("~Z.1[,c(",
                   paste(re.block.val[[1]],collapse=","),")]-1"))
Z.block[[2]] <-
  as.formula(paste("~Z.2[,c(",
                   paste(re.block.val[[1]],collapse=","),")]-1"))

# Fit model using lme() and extract coefficient estimates
groups.1 <- rep(1,length=n)

data.fr <- groupedData( y ~ X[,-1] | groups.1,
                       data = data.frame( y,X,Z.1,Z.2))

lme.fit <- lme(y~X[,-1],
               data=data.fr,
               random=
               list(groups.1 = pdBlocked(Z.block,
                      pdClass=c("pdIdent","pdIdent"))))

u1.hat <- as.vector(unlist(lme.fit$coef$ran))[Z.inds[[1]]-(num.add.comps+1)]
u2.hat <- as.vector(unlist(lme.fit$coef$ran))[Z.inds[[2]]-(num.add.comps+1)]

beta.hat <- as.vector(unlist(lme.fit$coef$fix))
sigusq1.hat <-
  (as.numeric(exp(attributes(summary(lme.fit)$apVar)$Pars[1])))^2
sigusq2.hat <-
  (as.numeric(exp(attributes(summary(lme.fit)$apVar)$Pars[2])))^2
sigesq.hat <- (lme.fit$sigma^2)

# Draw fits 
grid.size <- 101
x1.grid <- seq(0,1,length=grid.size)
x2.grid <- seq(0,1,length=grid.size)
ones.grid <- rep(1,grid.size)

X.grid <- cbind(ones.grid,x1.grid,x2.grid)
Z1.grid <- outer(x1.grid,knots[[1]],"-")
Z1.grid <- Z1.grid*(Z1.grid>0)
Z2.grid <- outer(x2.grid,knots[[2]],"-")
Z2.grid <- Z2.grid*(Z2.grid>0)

C.grid <- cbind(X.grid,Z1.grid,Z2.grid)
C.grid[,2:cols] <- scale(C.grid[,2:cols],scale=F,center=rescale.mean)

X.grid <- cbind(rep(1,length=grid.size),
           C.grid[,2:(num.add.comps+1)])
Z1.grid <- C.grid[,Z.inds[[1]]]
Z2.grid <- C.grid[,Z.inds[[2]]]

f1.hat.grid <- as.vector(X.grid[,2]*beta.hat[2]+Z1.grid%*%u1.hat)
f1.grid <- f1(x1.grid)

f2.hat.grid <- as.vector(X.grid[,3]*beta.hat[3]+Z2.grid%*%u2.hat)
f2.grid <- f2(x2.grid)

x11()
par(bty="l",mfrow=c(2,1))

plot(c(x1.grid),c(f1.hat.grid),bty="l",xlab="x1",
     ylab="f1(x)",type="n",ylim=range(c(f1.grid,f1.hat.grid)),
     main="Additive Model Component One and Fit")
lines(x1.grid,f1.grid) 
lines(x1.grid,f1.hat.grid,lwd=2,col="red")

plot(c(x2.grid),c(f2.hat.grid),bty="l",xlab="x2",
     ylab="f2(x)",type="n",ylim=range(c(f2.grid,f2.hat.grid)),
     main="Additive Model Component Two and Fit")
lines(x2.grid,f2.grid) 
lines(x2.grid,f2.hat.grid,lwd=2,col="red")