# 29 March 2004

# Author: Andrew Merton
# Email : merton@stat.colostate.edu

# Software to accompany the paper: "Model Selection for Geostatistical Models",
# Jennifer A. Hoeting, Richard A. Davis, Andrew A. Merton, and
# Sandra E. Thompson,  (2006) "Ecological Applications", 16(1), pp. 87-98.

# ---------------------------------------------------------------------------- #
# The following three functions have been made available for estimating the
# Matern correlation coefficients theta1(range) and theta2(smoothness).  Details
# for each function are listed below.
# -----
# FUNCTION COR.MATERN: Calculates the Matern correlation matix as parameterized
# by Hoeting et al.

# Inputs:
#   theta - Vector (or array) containing theta1, theta2, and theta3.
#             Theta1 = range.
#             Theta2 = smoothness.
#             Theta3 = nugget.
#   D     - Matrix (or array) of distances between sites. Typically D is an
#           (nxn) matrix with zeros along the main diagonal.

# Output:
#   The function returns a matrix (or array) of size D containing the Matern
#   correlations.
# -----
cor.matern <- function(theta,D) {
  if (is.vector(D)) names(D) <- NULL
  if (is.matrix(D)) dimnames(D) <- list(NULL,NULL)
  th1 <- theta[1]                             # Range.
  th2 <- theta[2]                             # Smoothness.
  th3 <- ifelse(length(theta)==3,theta[3],0)  # Nugget effect.

  u <- 2*D*sqrt(th2)/th1
  u <- ifelse(u>0,(1-th3)*u^th2/(2^(th2-1)*gamma(th2))*besselK(u,th2),1)
  return(u)
}
# ---------------------------------------------------------------------------- #
# FUNCTION LOG.PROFILE: Evalutes the negative of the log profile likelihood
# function.  This function is minimized by fitmatern(...) to find the optimal
# values for theta1 and theta2.

# Inputs:
#   theta - Vector (or array) containing theta1 and theta2.
#   D     - Matrix of distances between sites. The matrix must be square (nxn)
#           where n = length of the response vector.  D should be  symmetric
#           with zeros on the main diagonal (which corresponds to the Euclidian
#           distance between site i and site j where i = j).
#   X     - Design matrix. This matrix must be (nxp) where the first column is
#           a column of ones, i.e., there a p-1 explanatory variables. NA's are
#           not allowed.
#   Z     - An (nx1) response vector (or array).

# Output:
#   The function returns the value of the -log profile likelihood. see Hoeting
#   et al for more details.
# -----
log.profile <- function(theta,D,X,Z) {
# Compute the correlation matrix.
  n       <- length(Z)
  psi     <- cor.matern(theta,D)
  L       <- t(chol(psi))
  L.inv   <- forwardsolve(L,diag(nrow=n))
  L.det   <- prod(diag(L))
  psi.inv <- t(L.inv)%*%L.inv

# Compute the MLE for beta.
  beta    <- solve(t(X)%*%psi.inv%*%X)%*%t(X)%*%psi.inv%*%Z

# Compute the MLE for sigma.
  resid   <- Z-X%*%beta
  sigma2  <- t(resid)%*%psi.inv%*%resid/n

# Evaluate -log-profile likelihood.
  prof    <- log(L.det)+n*log(sigma2)/2+n/2

  return(max(prof,-1e+50))   # Return -1e+50 if there is an error.
}
# ---------------------------------------------------------------------------- #
# FUNCTION FITMATERN: Determins the optimal values for theta1 and theta2 that
# maximize the log profile likelihood. This function uses a constrained
# optimization routine and may need to be adjusted for some problems.
# Note: Maximize(log.profile) = Minimize(-log.profile).

# Inputs:
#   X      - (Optional) design matrix. If provided by the user, it must have a
#            column of ones as the first column. If ommited from the function
#            call, the routine automatically creates an (nx1) vector of ones.
#            This is useful for fitting the Matern coefficients to a data set
#            that has already had the mean trend removed.
#   Z      - Vector (or array) of responses (nx1).
#   D      - Matrix (nxn) of Euclidian distances between all pairs of points.
#            It should be symmetrical with zeros along the main diagonal.
#   theta0 - Vector (or array) containing initial guesses for theta1 and theta2.

# Outputs:
#   A list containing the maximum likelihood estimates (MLEs) for beta,
#   sigma2, theta (theta1,theta2), and the AICC. For further details, see
#   Hoeting et al.
# -----
fitmatern <- function(X=NULL,Z,D,theta0) {
# Preliminaries:
  n <- length(Z)
  if (is.null(X))   {X<-matrix(1,nrow=n)}
  if (is.vector(X)) {names(X)<-NULL;p<-1}
  if (is.matrix(X)) {dimnames(X)<-list(NULL,NULL);p<-length(X[1,])}
  if (is.vector(Z)) {names(Z)<-NULL}
  if (is.matrix(Z)) {dimnames(Z)<-list(NULL,NULL)}
  D.max    <- max(D)

# Maximize the log profile likelihood, i.e. minimize the -log prof likelihood.
  prof.max <- optim(theta0,log.profile,method="L-BFGS-B",
                    lower=c(0.0001,0.0001),upper=c(D.max,5),D=D,X=X,Z=Z)

# Return beta, sigma2, theta, and AICC.
  k       <- 2
  theta   <- prof.max$par
  psi     <- cor.matern(theta,D)  # Make correlation matrix.
  L       <- t(chol(psi))         # Cholesky decomposition.
  L.inv   <- forwardsolve(L,diag(nrow=n))
  psi.inv <- t(L.inv)%*%L.inv
  beta    <- solve(t(X)%*%psi.inv%*%X)%*%t(X)%*%psi.inv%*%Z
  resid   <- Z-X%*%beta
  sigma2  <- t(resid)%*%psi.inv%*%resid/n
  AICC    <- 2*prof.max$value + 2*n*(p+k+1)/(n-p-k-2)

  output  <- list(beta,sigma2,theta,AICC)
  names(output) <- list("beta","sigma2","theta","AICC")
  return(output)
}
# ---------------------------------------------------------------------------- #