# R/core.R In BayesNSGP: Bayesian Analysis of Non-Stationary Gaussian Process Models

#### Documented in determineNeighborsnsCrossdistnsCrossdist3dnsDistnsDist3dnsgpModelnsgpPredictorderCoordinatesMMD

#================================================
# Bayesian nonstationary Gaussian process
# modeling in NIMBLE
# Mark Risser and Daniel Turek
# Lawrence Berkeley National Laboratory
# January, 2019
#================================================

#================================================
# Functions for ordering coordinates and finding
# nearest neighbors
#================================================

## Script #2: nsgpOrderingNN.R (functions for ordering and finding nearest neighbors)
##
## - orderCoordinatesMMD: order coordinates by maxmin distance
## - determineNeighbors: identify k nearest neighbors

#==============================================================================
# Maximum-minimum distance (MMD) coordinate ordering
#==============================================================================

# ROxygen comments ----
#' Order coordinates according to a maximum-minimum distance criterion.
#'
#' \code{orderCoordinatesMMD} orders an array of (x,y) spatial coordinates
#' according to the "maximum minimum distance" (MMD), as described in Guinness,
#' 2018. (Points are selected to maximize their minimum distance to already-
#' selected points).
#'
#' @param coords N x 2 array of N 2-dimensional (x,y) spatial coordinates.
#' @param exact Logical; \code{FALSE} uses a fast approximation to MMD ordering
#' (and is almost always recommended), while \code{TRUE} uses exact MMD
#' ordering but is infeasible for large number of locations.
#'
#' @return A list of distances matrices, with the following components:
#' \item{orderedCoords}{N x 2 matrix; contains the ordered spatial coordinates
#' as \code{coords}.}
#' \item{orderedIndicesNoNA}{N-vector; contains the ordered indices with any
#' NA values removed.}
#'
#' @examples
#' coords <- cbind(runif(100), runif(100))
#' orderCoordinatesMMD(coords)
#'
#' @export
#'
orderCoordinatesMMD <- function(coords, exact = FALSE) {
## input coords: an Nx2 array of spatial coordinates
N <- dim(coords)[1]
if(N < 3) return(coords)
if(!exact) {       ## approximate MMD ordering
initialOrdering <- sample(1:N)
orderedIndices <- c(initialOrdering, rep(NA, 3*N))
indexLookupVector <- order(initialOrdering)
maxNeighbors <- floor(sqrt(N))
NN <- FNN::get.knn(coords, k = maxNeighbors)$nn.index nextSpot <- N+1 cycleCheckIndex <- -1 for(i in 2:(3*N)) { (targetIndex <- orderedIndices[i]) if(cycleCheckIndex == targetIndex) break if(cycleCheckIndex == -1) cycleCheckIndex <- targetIndex targetNeighbors <- NN[targetIndex, 1:min(maxNeighbors, round(N/(i+N-nextSpot)))] targetNeighborLocs <- indexLookupVector[targetNeighbors] if(min(targetNeighborLocs) < i) { ## relocate this index to the back orderedIndices[nextSpot] <- targetIndex orderedIndices[i] <- NA indexLookupVector[targetIndex] <- nextSpot nextSpot <- nextSpot + 1 } else cycleCheckIndex <- -1 } orderedIndicesNoNA <- orderedIndices[!is.na(orderedIndices)] orderedCoords <- coords[orderedIndicesNoNA,] } else { ## exact MMD ordering availableIndices <- 1:N orderedCoords <- array(NA, c(N,2)) sbar <- apply(coords, 2, mean) ## group centroid iNext <- which.min(sapply(1:N, function(i) sum((coords[i,] - sbar)^2))) orderedCoords[1,] <- coords[iNext,] availableIndices <- setdiff(availableIndices, iNext) for(i in 2:N) { aIndNext <- which.max( ## this indexes the availableIndices vector sapply(1:(N-i+1), function(j) { min(sapply(1:(i-1), function(k) sum((coords[availableIndices[j],] - orderedCoords[k,])^2))) })) iNext <- availableIndices[aIndNext] ## this indexes rows of the original s[] array orderedCoords[i,] <- coords[iNext,] availableIndices <- setdiff(availableIndices, iNext) } orderedIndicesNoNA <- NULL } return(list(orderedCoords = orderedCoords, orderedIndicesNoNA = orderedIndicesNoNA)) } #============================================================================== # Determine the k-nearest neighbors #============================================================================== # ROxygen comments ---- #' Determine the k-nearest neighbors for each spatial coordinate. #' #' \code{determineNeighbors} returns an N x k matrix of the nearest neighbors #' for spatial locations coords, with the ith row giving indices of the k nearest #' neighbors to the ith location, which are selected from among the 1,...(i-1) #' other spatial locations. The first row is -1's, since the first location has #' no neighbors. The i=2 through i=(k+1) rows each necessarily contain 1:i. #' #' @param coords N x 2 array of N 2-dimensional (x,y) spatial coordinates. #' @param k Scalar; number of neighbors #' #' @return An N x k matrix of nearest neighbor indices #' #' @examples #' coords <- cbind(runif(100), runif(100)) #' determineNeighbors(coords, 20) #' #' @export #' determineNeighbors <- function(coords, k) { N <- dim(coords)[1] d <- dim(coords)[2] if(k+2 > N) stop() nID <- array(-1, c(N,k)) ## populate unused values with -1, to prevent a warning from NIMBLE for(i in 2:(k+1)) nID[i, 1:(i-1)] <- as.numeric(1:(i-1)) if(d == 2){ for(i in (k+2):N) nID[i, 1:k] <- as.numeric(order((coords[1:(i-1),1] - coords[i,1])^2 + (coords[1:(i-1),2] - coords[i,2])^2)[1:k]) } else{ for(i in (k+2):N){ disti <- 0 for(j in 1:d){ disti <- disti + (coords[1:(i-1),j] - coords[i,j])^2 } nID[i, 1:k] <- as.numeric(order(disti)[1:k]) } } return(nID) } #================================================ # Core package functionality #================================================ ## Script #1: nsgpCore.R (core bayesNSGP functionality) ## ## - inverseEigen: calculate covariance elements based on eigendecomposition components ## - nsCorr: calculate a nonstationary Matern correlation matrix ## - nsCrosscorr: calculate a nonstationary Matern cross-correlation matrix ## - nsDist: calculate coordinate-specific distance matrices ## - nsCrossdist: calculate coordinate-specific cross-distance matrices ## - nsDist3d (formerly ns_dist_3d) ## - nsCrossdist3d ## - nsgpModel: NIMBLE code for a generic nonstationary GP model ## - nsgpPredict: posterior prediction for the NSGP #============================================================================== # Inverse eigendecomposition #============================================================================== # ROxygen comments ---- #' Calculate covariance elements based on eigendecomposition components #' #' \code{inverseEigen} calculates the inverse eigendecomposition -- in other #' words, the covariance elements based on the eigenvalues and vectors. For a #' 2x2 anisotropy (covariance) matrix, we parameterize the three unique values #' in terms of the two log eigenvalues and a rotation parameter on the #' rescaled logit. The function is coded as a \code{nimbleFunction} (see the #' \code{nimble} package) but can also be used as a regular R function. #' #' @param eigen_comp1 N-vector; contains values of the log of the first #' anisotropy eigenvalue for a set of locations. #' @param eigen_comp2 N-vector; contains values of the log of the second #' anisotropy eigenvalue for a set of locations. #' @param eigen_comp3 N-vector; contains values of the rescaled logit of #' the anisotropy rotation for a set of locations. #' @param which_Sigma Scalar; one of \code{(1,2,3)}, corresponding to which #' covariance component should be calculated (Sigma11, Sigma22, or Sigma12, #' respectively). #' #' @return A vector of anisotropy values (Sigma11, Sigma22, or Sigma12; depends #' on \code{which_Sigma}) for the corresponding set of locations. #' #' @examples #' # Generate some eigendecomposition elements (all three are real-valued) #' eigen_comp1 <- rnorm(10) #' eigen_comp2 <- rnorm(10) #' eigen_comp3 <- rnorm(10) #' inverseEigen( eigen_comp1, eigen_comp2, eigen_comp3, 2) # Return the Sigma22 values #' #' @export #' inverseEigen <- nimble::nimbleFunction( run = function( eigen_comp1 = double(1), eigen_comp2 = double(1), eigen_comp3 = double(1), which_Sigma = double(0) ) { returnType(double(1)) rotAngle <- (3.141592653/2)*exp(eigen_comp3)/(1 + exp(eigen_comp3)) # pi = 3.141592653 Gam11 <- cos(rotAngle) Gam22 <- cos(rotAngle) Gam12 <- -sin(rotAngle) Gam21 <- sin(rotAngle) Lam1 <- exp(eigen_comp1) Lam2 <- exp(eigen_comp2) if( which_Sigma == 1 ){ # Return Sigma11 return( Gam11^2*Lam1 + Gam12^2*Lam2 ) } if( which_Sigma == 2 ){ # Return Sigma22 return( Gam21^2*Lam1 + Gam22^2*Lam2 ) } if( which_Sigma == 3 ){ # Return Sigma12 return( Gam11*Gam21*Lam1 + Gam12*Gam22*Lam2 ) } stop('Error in inverseEigen function') ## prevent compiler warning return(numeric(10)) ## prevent compiler warning } ) #============================================================================== # Compiled besselK function #============================================================================== # ROxygen comments ---- # Compiled besselK function # # \code{RbesselK} and \code{CbesselK} calculates the modified Bessel function # of the third kind. # # @param dst Matrix; contains distances for the besselK function # @param nu Scalar; smoothness. # # @return A matrix with values of the corresponding Bessel function. # # RbesselK <- nimbleFunction( # run = function(dst = double(2), nu = double(0)) { # xVector <- besselK(dst, nu) # xMatrix <- matrix(xVector, dim(dst)[1], dim(dst)[2]) # returnType(double(2)) # return(xMatrix) # } # ) #============================================================================== # Calculate a nonstationary Matern correlation matrix #============================================================================== # ROxygen comments ---- #' Calculate a nonstationary Matern correlation matrix #' #' \code{nsCorr} calculates a nonstationary correlation matrix for a #' fixed set of locations, based on vectors of the unique anisotropy #' parameters for each station. Since the correlation function uses a #' spatially-varying Mahalanobis distance, this function requires coordinate- #' specific distance matrices (see below). The function is coded as a #' \code{nimbleFunction} (see the \code{nimble} package) but can also be #' used as a regular R function. #' #' @param dist1_sq N x N matrix; contains values of pairwise squared distances #' in the x-coordinate. #' @param dist2_sq N x N matrix; contains values of pairwise squared distances #' in the y-coordinate. #' @param dist12 N x N matrix; contains values of pairwise signed cross- #' distances between the x- and y-coordinates. The sign of each element is #' important; see \code{nsDist} function for the details of this calculation. #' in the x-coordinate. #' @param Sigma11 Vector of length N; contains the 1-1 element of the #' anisotropy process for each station. #' @param Sigma22 Vector of length N; contains the 2-2 element of the #' anisotropy process for each station. #' @param Sigma12 Vector of length N; contains the 1-2 element of the #' anisotropy process for each station. #' @param nu Scalar; Matern smoothness parameter. \code{nu = 0.5} corresponds #' to the Exponential correlation; \code{nu = Inf} corresponds to the Gaussian #' correlation function. #' @param d Scalar; dimension of the spatial coordinates. #' #' @return A correlation matrix for a fixed set of stations and fixed #' parameter values. #' #' @examples #' # Generate some coordinates and parameters #' coords <- cbind(runif(100),runif(100)) #' Sigma11 <- rep(1, 100) # Identity anisotropy process #' Sigma22 <- rep(1, 100) #' Sigma12 <- rep(0, 100) #' nu <- 2 #' # Calculate distances #' dist_list <- nsDist(coords) #' # Calculate the correlation matrix #' corMat <- nsCorr(dist_list$dist1_sq, dist_list$dist2_sq, dist_list$dist12,
#'                  Sigma11, Sigma22, Sigma12, nu, ncol(coords))
#'
#' @export
#'
nsCorr <- nimble::nimbleFunction(
run = function( dist1_sq = double(2), dist2_sq = double(2), dist12 = double(2),
Sigma11 = double(1), Sigma22 = double(1), Sigma12 = double(1),
nu = double(0), d = double(0) ) {

returnType(double(2))
N <- length(Sigma11)
if( dist2_sq[1,1] == -1 ){ # Isotropic case
# Calculate the scale matrix
if(N == 1){
det1 <- Sigma11^d
diagSqrtSqrtDet1 <- matrix(sqrt(sqrt(det1)), N, N)
} else{
det1 <- Sigma11^d
diagSqrtSqrtDet1 <- diag(sqrt(sqrt(det1)))
}
mat11_a <- matrix(Sigma11, nrow = N, ncol = N)
mat11 <- 0.5*(mat11_a + t(mat11_a))
det12 <- mat11^d
oneOverDet12 <- 1/det12
Scale.mat <- diagSqrtSqrtDet1 %*% sqrt(oneOverDet12) %*% diagSqrtSqrtDet1
# Calculate the distance matrix
Dist.mat <- sqrt( dist1_sq/mat11 )
} else{
# Calculate the scale matrix
if(N == 1){
det1 <- Sigma11*Sigma22 - Sigma12^2
diagSqrtSqrtDet1 <- matrix(sqrt(sqrt(det1)), N, N)
} else{
det1 <- Sigma11*Sigma22 - Sigma12^2
diagSqrtSqrtDet1 <- diag(sqrt(sqrt(det1)))
}
mat11_a <- matrix(Sigma11, nrow = N, ncol = N)
mat22_a <- matrix(Sigma22, nrow = N, ncol = N)
mat12_a <- matrix(Sigma12, nrow = N, ncol = N)
mat11 <- 0.5*(mat11_a + t(mat11_a))
mat22 <- 0.5*(mat22_a + t(mat22_a))
mat12 <- 0.5*(mat12_a + t(mat12_a))
det12 <- mat11*mat22 - mat12^2
oneOverDet12 <- 1/det12
Scale.mat <- diagSqrtSqrtDet1 %*% sqrt(oneOverDet12) %*% diagSqrtSqrtDet1
# Calculate the distance matrix
inv11 <-  mat22 * oneOverDet12
inv22 <-  mat11 * oneOverDet12
inv12 <- -mat12 * oneOverDet12
Dist.mat <- sqrt( inv11*dist1_sq + 2*inv12*dist12 + inv22*dist2_sq )
}

# Combine
if( nu == 0.5 ){ # Exponential correlation
Unscl.corr <- exp(-Dist.mat)
} else{
if( nu == Inf ){ # Gaussian (squared exponential) correlation
Unscl.corr <- exp(-(Dist.mat^2))
} else{ # Else: Matern with smoothness nu
##Unscl.corr <- (exp(lgamma(nu)) * 2^(nu - 1))^(-1) * (Dist.mat)^nu * besselK( x = Dist.mat, nu = nu )
xVector <- besselK(Dist.mat, nu)
xMatrix <- matrix(xVector, dim(Dist.mat)[1], dim(Dist.mat)[2])
Unscl.corr <- (exp(lgamma(nu)) * 2^(nu - 1))^(-1) * (Dist.mat)^nu * xMatrix
diag(Unscl.corr) <- 1
}
}
nsCorr <- Scale.mat*Unscl.corr
return(nsCorr)
}
)

#==============================================================================
# Calculate a stationary Matern correlation matrix
#==============================================================================

# ROxygen comments ----
#' Calculate a stationary Matern correlation matrix
#'
#' \code{matern_corr} calculates a stationary Matern correlation matrix for a
#' fixed set of locations, based on a range and smoothness parameter. This
#' function is primarily used for the "npGP" and "approxGP" models. The
#' function is coded as a \code{nimbleFunction} (see the \code{nimble} package)
#' but can also be used as a regular R function.
#'
#' @param dist N x N matrix; contains values of pairwise Euclidean distances in
#' the x-y plane.
#' @param rho Scalar; "range" parameter used to rescale distances
#' @param nu Scalar; Matern smoothness parameter. \code{nu = 0.5} corresponds
#' to the Exponential correlation; \code{nu = Inf} corresponds to the Gaussian
#' correlation function.
#'
#' @return A correlation matrix for a fixed set of stations and fixed
#' parameter values.
#'
#' @examples
#' # Generate some coordinates
#' coords <- cbind(runif(100),runif(100))
#' nu <- 2
#' # Calculate distances -- can use nsDist to calculate Euclidean distances
#' dist_list <- nsDist(coords, isotropic = TRUE)
#' # Calculate the correlation matrix
#' corMat <- matern_corr(sqrt(dist_list$dist1_sq), 1, nu) #' #' @export #' matern_corr <- nimble::nimbleFunction( run = function( dist = double(2), rho = double(0), nu = double(0) ) { returnType(double(2)) Nr <- dim(dist)[1] Nc <- dim(dist)[2] if( nu == 0.5 ){ # Exponential correlation return(exp(-dist/rho)) } if( nu == Inf ){ # Gaussian (squared exponential) correlation return(exp(-(dist/rho)^2)) } # Else: Matern with smoothness nu xVector <- besselK(dist/rho, nu) xMatrix <- matrix(xVector, dim(dist)[1], dim(dist)[2]) temp <- (exp(lgamma(nu)) * 2^(nu - 1))^(-1) * (dist/rho)^nu * xMatrix # Check for zeros in dist if(min(dist) == 0) { for(i in 1:Nc){ if( min(dist[1:Nr,i]) == 0 ) { for(j in 1:Nr) { if(dist[j,i] == 0) { temp[j,i] <- 1 } } } } } return(temp) } ) #============================================================================== # Calculate a nonstationary Matern cross-correlation matrix #============================================================================== # ROxygen comments ---- #' Calculate a nonstationary Matern cross-correlation matrix #' #' \code{nsCrosscorr} calculates a nonstationary cross-correlation matrix #' between two fixed sets of locations (a prediction set with M locations, and #' the observed set with N locations), based on vectors of the unique anisotropy #' parameters for each station. Since the correlation function uses a #' spatially-varying Mahalanobis distance, this function requires coordinate- #' specific distance matrices (see below). The function is coded as a #' \code{nimbleFunction} (see the \code{nimble} package) but can also be #' used as a regular R function. #' #' @param Xdist1_sq M x N matrix; contains values of pairwise squared cross-distances #' in the x-coordinate. #' @param Xdist2_sq M x N matrix; contains values of pairwise squared cross-distances #' in the y-coordinate. #' @param Xdist12 M x N matrix; contains values of pairwise signed cross/cross- #' distances between the x- and y-coordinates. The sign of each element is #' important; see \code{nsDist} function for the details of this calculation. #' in the x-coordinate. #' @param Sigma11 Vector of length N; contains the 1-1 element of the #' anisotropy process for each observed location. #' @param Sigma22 Vector of length N; contains the 2-2 element of the #' anisotropy process for each observed location. #' @param Sigma12 Vector of length N; contains the 1-2 element of the #' anisotropy process for each observed location. #' @param PSigma11 Vector of length N; contains the 1-1 element of the #' anisotropy process for each prediction location. #' @param PSigma22 Vector of length N; contains the 2-2 element of the #' anisotropy process for each prediction location. #' @param PSigma12 Vector of length N; contains the 1-2 element of the #' anisotropy process for each prediction location. #' @param nu Scalar; Matern smoothness parameter. \code{nu = 0.5} corresponds #' to the Exponential correlation; \code{nu = Inf} corresponds to the Gaussian #' correlation function. #' @param d Scalar; dimension of the spatial domain. #' #' @return A M x N cross-correlation matrix for two fixed sets of stations and #' fixed parameter values. #' #' @examples #' # Generate some coordinates and parameters #' coords <- cbind(runif(100),runif(100)) #' Sigma11 <- rep(1, 100) # Identity anisotropy process #' Sigma22 <- rep(1, 100) #' Sigma12 <- rep(0, 100) #' Pcoords <- cbind(runif(200),runif(200)) #' PSigma11 <- rep(1, 200) # Identity anisotropy process #' PSigma22 <- rep(1, 200) #' PSigma12 <- rep(0, 200) #' nu <- 2 #' # Calculate distances #' Xdist_list <- nsCrossdist(coords, Pcoords) #' # Calculate the correlation matrix #' XcorMat <- nsCrosscorr(Xdist_list$dist1_sq, Xdist_list$dist2_sq, Xdist_list$dist12,
#'    Sigma11, Sigma22, Sigma12, PSigma11, PSigma22, PSigma12, nu, ncol(coords))
#'
#' @export
#'
nsCrosscorr <- nimble::nimbleFunction(
run = function( Xdist1_sq = double(2), Xdist2_sq = double(2), Xdist12 = double(2),
Sigma11 = double(1), Sigma22 = double(1), Sigma12 = double(1),
PSigma11 = double(1), PSigma22 = double(1), PSigma12 = double(1),
nu = double(0), d = double(0) ) {

returnType(double(2))
N <- length(Sigma11)
M <- length(PSigma11)

if( Xdist2_sq[1,1] == -1 ){ # Isotropic case
# Calculate the scale matrix
if(N == 1){
det1 <- Sigma11^d
diagSqrtSqrtDet1 <- matrix(sqrt(sqrt(det1)), N, N)
} else{
det1 <- Sigma11^d
diagSqrtSqrtDet1 <- diag(sqrt(sqrt(det1)))
}
if(M == 1){
Pdet1 <- PSigma11^d
diagSqrtSqrtPDet1 <- matrix(sqrt(sqrt(Pdet1)), M, M)
} else{
Pdet1 <- PSigma11^d
diagSqrtSqrtPDet1 <- diag(sqrt(sqrt(Pdet1)))
}
mat11_1 <- t(matrix(Sigma11, nrow = N, ncol = M))
mat11_2 <- matrix(PSigma11, nrow = M, ncol = N)
mat11 <- 0.5*(mat11_1 + mat11_2)
det12 <- mat11^d
oneOverDet12 <- 1/det12
Scale.mat <- diagSqrtSqrtPDet1 %*% sqrt(oneOverDet12) %*% diagSqrtSqrtDet1

# Calculate the distance matrix
Dist.mat <- sqrt( Xdist1_sq/mat11 )
} else{
# Calculate the scale matrix
if(N == 1){
det1 <- Sigma11*Sigma22 - Sigma12^2
diagSqrtSqrtDet1 <- matrix(sqrt(sqrt(det1)), N, N)
} else{
det1 <- Sigma11*Sigma22 - Sigma12^2
diagSqrtSqrtDet1 <- diag(sqrt(sqrt(det1)))
}
if(M == 1){
Pdet1 <- PSigma11*PSigma22 - PSigma12^2
diagSqrtSqrtPDet1 <- matrix(sqrt(sqrt(Pdet1)), M, M)
} else{
Pdet1 <- PSigma11*PSigma22 - PSigma12^2
diagSqrtSqrtPDet1 <- diag(sqrt(sqrt(Pdet1)))
}
mat11_1 <- t(matrix(Sigma11, nrow = N, ncol = M))
mat11_2 <- matrix(PSigma11, nrow = M, ncol = N)
mat22_1 <- t(matrix(Sigma22, nrow = N, ncol = M))
mat22_2 <- matrix(PSigma22, nrow = M, ncol = N)
mat12_1 <- t(matrix(Sigma12, nrow = N, ncol = M))
mat12_2 <- matrix(PSigma12, nrow = M, ncol = N)
mat11 <- 0.5*(mat11_1 + mat11_2)
mat22 <- 0.5*(mat22_1 + mat22_2)
mat12 <- 0.5*(mat12_1 + mat12_2)
det12 <- mat11*mat22 - mat12^2
oneOverDet12 <- 1/det12
Scale.mat <- diagSqrtSqrtPDet1 %*% sqrt(oneOverDet12) %*% diagSqrtSqrtDet1

# Calculate the distance matrix
inv11 <-  mat22 * oneOverDet12
inv22 <-  mat11 * oneOverDet12
inv12 <- -mat12 * oneOverDet12
Dist.mat <- sqrt( inv11*Xdist1_sq + 2*inv12*Xdist12 + inv22*Xdist2_sq )
}

# Combine
if( nu == 0.5 ){ # Exponential correlation
Unscl.corr <- exp(-Dist.mat)
} else{
if( nu == Inf ){ # Gaussian (squared exponential) correlation
Unscl.corr <- exp(-(Dist.mat^2))
} else{ # Else: Matern with smoothness nu
##Unscl.corr <- (exp(lgamma(nu)) * 2^(nu - 1))^(-1) * (Dist.mat)^nu * besselK( x = Dist.mat, nu = nu )
xVector <- besselK(Dist.mat, nu)
xMatrix <- matrix(xVector, dim(Dist.mat)[1], dim(Dist.mat)[2])
Unscl.corr <- (exp(lgamma(nu)) * 2^(nu - 1))^(-1) * (Dist.mat)^nu * xMatrix
## this line will not fly:  Unscl.corr[Unscl.corr == Inf] <- 1
## inelegant, but accomplishes the same:
if(min(Dist.mat) == 0) {
for(i in 1:dim(Unscl.corr)[1]) {
for(j in 1:dim(Unscl.corr)[2]) {
if(Dist.mat[i,j] == 0)   Unscl.corr[i,j] <- 1
}
}
}
# diag(Unscl.corr) <- 1
}
}
nsCrosscorr <- Scale.mat*Unscl.corr
return(nsCrosscorr)
}
)

#==============================================================================
# Calculate coordinate-specific distance matrices
#==============================================================================

# ROxygen comments ----
#' Calculate coordinate-specific distance matrices
#'
#' \code{nsDist} calculates x, y, and x-y distances for use in the
#' nonstationary correlation calculation. The sign of the cross-distance
#' is important. The function contains an optional argument for re-scaling
#' the distances such that the coordinates lie in a square.
#'
#' @param coords N x 2 matrix; contains the x-y coordinates of stations
#' @param scale_factor Scalar; optional argument for re-scaling the distances.
#' @param isotropic Logical; indicates whether distances should be calculated
#' separately for each coordinate dimension (FALSE) or simultaneously for all
#' coordinate dimensions (TRUE). \code{isotropic = TRUE} can only be used for
#' two-dimensional coordinate systems.
#'
#' @return A list of distances matrices, with the following components:
#' \item{dist1_sq}{N x N matrix; contains values of pairwise squared distances
#' in the x-coordinate.}
#' \item{dist2_sq}{N x N matrix; contains values of pairwise squared distances
#' in the y-coordinate.}
#' \item{dist12}{N x N matrix; contains values of pairwise signed cross-
#' distances between the x- and y-coordinates.}
#' \item{scale_factor}{Value of the scale factor used to rescale distances.}
#'
#' @examples
#' # Generate some coordinates
#' coords <- cbind(runif(100),runif(100))
#' # Calculate distances
#' dist_list <- nsDist(coords)
#' # Use nsDist to calculate Euclidean distances
#' dist_Euclidean <- sqrt(nsDist(coords, isotropic = TRUE)$dist1_sq) #' #' @export #' nsDist <- function( coords, scale_factor = NULL, isotropic = FALSE ){ N <- nrow(coords) d <- ncol(coords) if( !isotropic & d != 2 ) stop("nsDist: Anisotropy (isotropic = FALSE) only available for 2-dimensional coordinate systems.") if(!isotropic){ # Calculate distances dists1 <- as.matrix(dist(coords[,1], upper = T, diag = T)) dists2 <- as.matrix(dist(coords[,2], upper = T, diag = T)) temp1 <- matrix(coords[,1], nrow = N, ncol = N) temp2 <- matrix(coords[,2], nrow = N, ncol = N) sgn_mat1 <- ( temp1 - t(temp1) >= 0 ) sgn_mat1[sgn_mat1 == FALSE] <- -1 sgn_mat2 <- ( temp2 - t(temp2) >= 0 ) sgn_mat2[sgn_mat2 == FALSE] <- -1 dist1_sq <- dists1^2 dist2_sq <- dists2^2 dist12 <- sgn_mat1*dists1*sgn_mat2*dists2 } else{ dist1_sq <- as.matrix(dist(coords, upper = T, diag = T))^2 dist2_sq <- matrix(-1, N, N) dist12 <- matrix(0, N, N) } # Rescale if needed if( !is.null(scale_factor) ){ dist1_sq <- dist1_sq/scale_factor dist2_sq <- dist2_sq/scale_factor dist12 <- dist12/scale_factor } return(list( dist1_sq = dist1_sq, dist2_sq = dist2_sq, dist12 = dist12, scale_factor = scale_factor )) } #============================================================================== # Coordinate-specific distance matrices, only for NN #============================================================================== # ROxygen comments ---- #' Calculate coordinate-specific distance matrices, only for nearest neighbors #' and store in an array #' #' \code{nsDist3d} generates and returns new 3-dimensional arrays containing #' the former dist1_sq, dist2_sq, and dist12 matrices, but #' only as needed for the k nearest-neighbors of each location. #' these 3D matrices (dist1_3d, dist2_3d, and dist12_3d) #' are used in the new implementation of calculateAD_ns(). #' #' @param coords N x 2 matrix; contains the x-y coordinates of stations. #' @param nID N x k matrix; contains indices of nearest neighbors. #' @param scale_factor Scalar; optional argument for re-scaling the distances. #' @param isotropic Logical; indicates whether distances should be calculated #' separately for each coordinate dimension (FALSE) or simultaneously for all #' coordinate dimensions (TRUE). \code{isotropic = TRUE} can only be used for #' two-dimensional coordinate systems. #' #' @return Arrays with nearest neighbor distances in each coordinate #' direction. #' #' @examples #' # Generate some coordinates and neighbors #' coords <- cbind(runif(100),runif(100)) #' nID <- determineNeighbors(coords, 10) #' # Calculate distances #' nsDist3d(coords, nID) #' #' @export #' nsDist3d <- function(coords, nID, scale_factor = NULL, isotropic = FALSE) { N <- nrow(coords) d <- ncol(coords) if( !isotropic & d != 2 ) stop("Anisotropy (isotropic = FALSE) only available for 2-dimensional coordinate systems.") k <- ncol(nID) dist1_3d <- array(0, c(N, k+1, k+1)) dist2_3d <- array(0, c(N, k+1, k+1)) dist12_3d <- array(0, c(N, k+1, k+1)) if(!isotropic){ for(i in 2:N) { if(i<=k) nNei <- i-1 else nNei <- k ind <- c( nID[i,1:nNei], i ) thisN <- nNei + 1 theseCoords <- coords[ind, ] dists1 <- as.matrix(dist(theseCoords[,1])) dists2 <- as.matrix(dist(theseCoords[,2])) temp1 <- matrix(theseCoords[,1], nrow = thisN, ncol = thisN) temp2 <- matrix(theseCoords[,2], nrow = thisN, ncol = thisN) sgn_mat1 <- ( temp1 - t(temp1) >= 0 ) sgn_mat1[sgn_mat1 == FALSE] <- -1 sgn_mat2 <- ( temp2 - t(temp2) >= 0 ) sgn_mat2[sgn_mat2 == FALSE] <- -1 dist1_3d[i, 1:thisN, 1:thisN] <- dists1^2 dist2_3d[i, 1:thisN, 1:thisN] <- dists2^2 dist12_3d[i, 1:thisN, 1:thisN] <- sgn_mat1*dists1*sgn_mat2*dists2 } } else{ dist2_3d[1,,] <- -1 for(i in 2:N) { if(i<=k) nNei <- i-1 else nNei <- k ind <- c( nID[i,1:nNei], i ) thisN <- nNei + 1 theseCoords <- coords[ind, ] dists1 <- as.matrix(dist(theseCoords)) dist1_3d[i, 1:thisN, 1:thisN] <- dists1^2 dist2_3d[i,,] <- -1 } } if(!is.null(scale_factor)) { dist1_3d <- dist1_3d / scale_factor dist2_3d <- dist2_3d / scale_factor dist12_3d <- dist12_3d / scale_factor } return(list(dist1_3d = dist1_3d, dist2_3d = dist2_3d, dist12_3d = dist12_3d, scale_factor = scale_factor)) } #============================================================================== # Calculate coordinate-specific cross-distance matrices #============================================================================== # ROxygen comments ---- #' Calculate coordinate-specific cross-distance matrices #' #' \code{nsCrossdist} calculates coordinate-specific cross distances in x, y, #' and x-y for use in the nonstationary cross-correlation calculation. This #' function is useful for calculating posterior predictions. #' #' @param coords N x 2 matrix; contains x-y coordinates of station (observed) #' locations. #' @param Pcoords M x 2 matrix; contains x-y coordinates of prediction #' locations. #' @param scale_factor Scalar; optional argument for re-scaling the distances. #' @param isotropic Logical; indicates whether distances should be calculated #' using Euclidean distance (\code{isotropic = TRUE}) or using the anisotropic #' formulation (\code{isotropic = FALSE}). #' #' @return A list of distances matrices, with the following components: #' \item{dist1_sq}{M x N matrix; contains values of pairwise squared cross- #' distances in the x-coordinate.} #' \item{dist2_sq}{M x N matrix; contains values of pairwise squared cross- #' distances in the y-coordinate.} #' \item{dist12}{M x N matrix; contains values of pairwise signed cross- #' distances between the x- and y-coordinates.} #' \item{scale_factor}{Value of the scale factor used to rescale distances.} #' #' @examples #' # Generate some coordinates #' coords <- cbind(runif(100),runif(100)) #' Pcoords <- cbind(runif(200),runif(200)) #' # Calculate distances #' Xdist_list <- nsCrossdist(coords, Pcoords) #' #' @export #' nsCrossdist <- function(coords, Pcoords, scale_factor = NULL, isotropic = FALSE ){ N <- nrow(coords) M <- nrow(Pcoords) d <- ncol(Pcoords) if( !isotropic & d != 2 ) stop("nsCrossdist: Anisotropy (isotropic = FALSE) only available for 2-dimensional coordinate systems.") if(!isotropic){ # Calculate distances dists1 <- mahalanobis.dist(data.x = Pcoords[,1], data.y = coords[,1], vc = diag(1)) dists2 <- mahalanobis.dist(data.x = Pcoords[,2], data.y = coords[,2], vc = diag(1)) temp1a <- matrix(coords[,1], nrow = M, ncol = N, byrow = TRUE) temp1b <- matrix(Pcoords[,1], nrow = M, ncol = N) temp2a <- matrix(coords[,2], nrow = M, ncol = N, byrow = TRUE) temp2b <- matrix(Pcoords[,2], nrow = M, ncol = N) sgn_mat1 <- ( temp1a - temp1b >= 0 ) sgn_mat1[sgn_mat1 == FALSE] <- -1 sgn_mat2 <- ( temp2a - temp2b >= 0 ) sgn_mat2[sgn_mat2 == FALSE] <- -1 dist1_sq <- dists1^2 dist2_sq <- dists2^2 dist12 <- sgn_mat1*dists1*sgn_mat2*dists2 } else{ dist1_sq <- mahalanobis.dist(data.x = Pcoords, data.y = coords, vc = diag(d))^2 dist2_sq <- matrix(-1, M, N) dist12 <- matrix(0, M, N) } # Rescale if needed if( !is.null(scale_factor) ){ dist1_sq <- dist1_sq/scale_factor dist2_sq <- dist2_sq/scale_factor dist12 <- dist12/scale_factor } return(list( dist1_sq = dist1_sq, dist2_sq = dist2_sq, dist12 = dist12, scale_factor = scale_factor )) } #============================================================================== # Coordinate-specific cross-distance matrices, only for NN #============================================================================== # ROxygen comments ---- #' Calculate coordinate-specific cross-distance matrices, only for nearest neighbors #' and store in an array #' #' \code{nsCrossdist3d} generates and returns new 3-dimensional arrays containing #' the former dist1_sq, dist2_s1, and dist12 matrices, but #' only as needed for the k nearest-neighbors of each location. #' these 3D matrices (dist1_3d, dist2_3d, and dist12_3d) #' are used in the new implementation of calculateAD_ns(). #' #' @param coords N x d matrix; contains the x-y coordinates of stations. #' @param predCoords M x d matrix #' @param P_nID N x k matrix; contains indices of nearest neighbors. #' @param scale_factor Scalar; optional argument for re-scaling the distances. #' @param isotropic Logical; indicates whether distances should be calculated #' separately for each coordinate dimension (FALSE) or simultaneously for all #' coordinate dimensions (TRUE). \code{isotropic = FALSE} can only be used for #' two-dimensional coordinate systems. #' #' @return Arrays with nearest neighbor distances in each coordinate #' direction. When the spatial dimension d > 2, dist1_3d contains squared #' Euclidean distances, and dist2_3d and dist12_3d are empty. #' #' @examples #' # Generate some coordinates and neighbors #' coords <- cbind(runif(100),runif(100)) #' predCoords <- cbind(runif(200),runif(200)) #' P_nID <- FNN::get.knnx(coords, predCoords, k = 10)$nn.index # Prediction NN
#' # Calculate distances
#' Pdist <- nsCrossdist3d(coords, predCoords, P_nID)
#'
#' @export
#'
nsCrossdist3d <- function(coords, predCoords, P_nID, scale_factor = NULL, isotropic = FALSE) {
N <- nrow(coords)
M <- nrow(predCoords)

d <- ncol(coords)
if( !isotropic & d != 2 ) stop("Anisotropy (isotropic = FALSE) only available for 2-dimensional coordinate systems.")

k <- ncol(P_nID)
thisN <- k+1
dist1_3d <- array(0, c(M, k+1, k+1))
dist2_3d <- array(0, c(M, k+1, k+1))
dist12_3d <- array(0, c(M, k+1, k+1))

if(!isotropic){
for(i in 1:M){
theseCoords <- rbind(coords[P_nID[i,], ], predCoords[i,])
dists1 <- as.matrix(dist(theseCoords[,1]))
dists2 <- as.matrix(dist(theseCoords[,2]))
temp1 <- matrix(theseCoords[,1], nrow = thisN, ncol = thisN)
temp2 <- matrix(theseCoords[,2], nrow = thisN, ncol = thisN)
sgn_mat1 <- ( temp1 - t(temp1) >= 0 )
sgn_mat1[sgn_mat1 == FALSE] <- -1
sgn_mat2 <- ( temp2 - t(temp2) >= 0 )
sgn_mat2[sgn_mat2 == FALSE] <- -1
dist1_3d[i, 1:thisN, 1:thisN] <- dists1^2
dist2_3d[i, 1:thisN, 1:thisN] <- dists2^2
dist12_3d[i, 1:thisN, 1:thisN] <- sgn_mat1*dists1*sgn_mat2*dists2
}
} else{
for(i in 1:M){
theseCoords <- rbind(coords[P_nID[i,], ], predCoords[i,])
dists1 <- as.matrix(dist(theseCoords))
dist1_3d[i, 1:thisN, 1:thisN] <- dists1^2
dist2_3d[i,,] <- -1
}
}

if(!is.null(scale_factor)) {
dist1_3d  <- dist1_3d  / scale_factor
dist2_3d  <- dist2_3d  / scale_factor
dist12_3d <- dist12_3d / scale_factor
}
return(list(dist1_3d  = dist1_3d,
dist2_3d  = dist2_3d,
dist12_3d = dist12_3d,
scale_factor = scale_factor))
}

#==============================================================================
# NIMBLE code for a generic nonstationary GP model
#==============================================================================

# ROxygen comments ----
#' NIMBLE code for a generic nonstationary GP model
#'
#' This function sets up and compiles a nimble model for a general
#' nonstationary Gaussian process.
#'
#' @param tau_model Character; specifies the model to be used for the log(tau)
#' process. Options are \code{"constant"} (spatially-constant),
#' \code{"logLinReg"} (log-linear regression), and \code{"approxGP"}
#' (approximation to a Gaussian process).
#' @param sigma_model Character; specifies the model to be used for the
#' log(sigma) process. See \code{tau_model} for options.
#' @param Sigma_model Character; specifies the model to be used for the
#' Sigma anisotropy process. Options are \code{"constant"} (spatially-constant),
#' \code{"constantIso"} (spatially-constant and isotropic), \code{"covReg"}
#' (covariance regression), \code{"compReg"} (componentwise regression),
#' \code{"compRegIso"} (isotropic componentwise regression), \code{"npApproxGP"}
#' (nonparameteric regression via an approximation to a stationary Gaussian
#' process), and \code{"npApproxGPIso"} (isotropic nonparameteric regression
#' via an approximation to a stationary Gaussian process)
#' @param mu_model Character; specifies the model to be used for the mu mean
#' process. Options are \code{"constant"} (spatially-constant), \code{"linReg"}
#' (linear regression), and \code{"zero"} (a fixed zero-mean).
#' @param likelihood Character; specifies the likelihood model. Options are
#' \code{"fullGP"} (the exact Gaussian process likelihood), \code{"NNGP"} (the
#' nearest-neighbor GP for the response approximate likelihood), and \code{"SGV"}
#' (the sparse general Vecchia approximate likelihood).
#' @param coords N x d matrix of spatial coordinates.
#' @param data N-vector; observed vector of the spatial process of interest
#' @param constants A list of constants required to build the model; depends on
#' the specific parameter process models chosen.
#' @param monitorAllSampledNodes Logical; indicates whether all sampled nodes
#' should be stored (\code{TRUE}) or not (\code{FALSE}).
#' @param ... Additional arguments can be passed to the function; for example,
#' as an alternative to the \code{constants} list, items can be passed directly
#' via this argument.
#'
#' @return A \code{nimbleCode} object.
#'
#' @examples
#' # Generate some data: stationary/isotropic
#' N <- 100
#' coords <- matrix(runif(2*N), ncol = 2)
#' alpha_vec <- rep(log(sqrt(1)), N) # Log process SD
#' delta_vec <- rep(log(sqrt(0.05)), N) # Log nugget SD
#' Sigma11_vec <- rep(0.4, N) # Kernel matrix element 1,1
#' Sigma22_vec <- rep(0.4, N) # Kernel matrix element 2,2
#' Sigma12_vec <- rep(0, N) # Kernel matrix element 1,2
#' mu_vec <- rep(0, N) # Mean
#' nu <- 0.5 # Smoothness
#' dist_list <- nsDist(coords)
#' Cor_mat <- nsCorr( dist1_sq = dist_list$dist1_sq, dist2_sq = dist_list$dist2_sq,
#'                    dist12 = dist_list$dist12, Sigma11 = Sigma11_vec, #' Sigma22 = Sigma22_vec, Sigma12 = Sigma12_vec, nu = nu ) #' Cov_mat <- diag(exp(alpha_vec)) %*% Cor_mat %*% diag(exp(alpha_vec)) #' D_mat <- diag(exp(delta_vec)^2) #' set.seed(110) #' data <- as.numeric(mu_vec + t(chol(Cov_mat + D_mat)) %*% rnorm(N)) #' # Set up constants #' constants <- list( nu = 0.5, Sigma_HP1 = 2 ) #' # Defaults: tau_model = "constant", sigma_model = "constant", mu_model = "constant", #' # and Sigma_model = "constant" #' Rmodel <- nsgpModel(likelihood = "fullGP", constants = constants, coords = coords, data = data ) #' #' @export #' nsgpModel <- function( tau_model = "constant", sigma_model = "constant", Sigma_model = "constant", mu_model = "constant", likelihood = "fullGP", coords, data, constants = list(), monitorAllSampledNodes = TRUE, ... ) { ##============================================ ## Models for tau ##============================================ tau_model_list <- list( constant = list( ## 1. tau_HP1 Standard deviation for the log-linear standard deviation ## 2. delta Scalar; represents the standard deviation (constant over the domain) ## 3. ones N-vector of 1's code = quote({ log_tau_vec[1:N] <- log(sqrt(delta))*ones[1:N] delta ~ dunif(0, tau_HP1) }), constants_needed = c("ones", "tau_HP1"), inits = list(delta = quote(tau_HP1/10)) ), logLinReg = list( ## 1. X_tau N x p_tau design matrix; leading column of 1's with (p_tau - 1) other covariates ## 2. tau_HP1 Standard deviation for the log-linear regression coefficients ## 3. p_tau Number of design columns ## 4. delta Vector of length p_tau; represents log-linear regression coefficients code = quote({ log_tau_vec[1:N] <- X_tau[1:N,1:p_tau] %*% delta[1:p_tau] for(l in 1:p_tau){ delta[l] ~ dnorm(0, sd = tau_HP1) } tau_constraint1 ~ dconstraint( max(abs(log_tau_vec[1:N])) < maxAbsLogSD ) }), constants_needed = c("X_tau", "p_tau", "tau_HP1", "maxAbsLogSD"), inits = list(delta = quote(rep(0, p_tau))), constraints_needed = c('tau_constraint1') ), approxGP = list( ## 1. tau_HP1 Gaussian process standard deviation ## 2. tau_HP2 Gaussian process mean ## 3. tau_HP3 Gaussian process range ## 4. tau_HP4 Gaussian process smoothness ## 5. ones N-vector of 1's ## 6. tau_cross_dist N x p_tau matrix of inter-point Euclidean distances, obs. coords vs. knot locations ## 7. tau_knot_dist p_tau x p_tau matrix of inter-point Euclidean distances, knot locations ## 8. p_tau Number of knot locations code = quote({ log_tau_vec[1:N] <- tauGP_mu*ones[1:N] + tauGP_sigma*Pmat_tau[1:N,1:p_tau] %*% w_tau[1:p_tau] Pmat_tau[1:N,1:p_tau] <- matern_corr(tau_cross_dist[1:N,1:p_tau], tauGP_phi, tau_HP2) Vmat_tau[1:p_tau,1:p_tau] <- matern_corr(tau_knot_dist[1:p_tau,1:p_tau], tauGP_phi, tau_HP2) w_tau_mean[1:p_tau] <- 0*ones[1:p_tau] w_tau[1:p_tau] ~ dmnorm( mean = w_tau_mean[1:p_tau], prec = Vmat_tau[1:p_tau,1:p_tau] ) # Hyperparameters tauGP_mu ~ dnorm(0, sd = tau_HP1) tauGP_phi ~ dunif(0, tau_HP3) # Range parameter, GP tauGP_sigma ~ dunif(0, tau_HP4) # SD parameter, GP # Constraint tau_constraint1 ~ dconstraint( max(abs(log_tau_vec[1:N])) < maxAbsLogSD ) }), constants_needed = c("ones", "tau_knot_coords", "tau_cross_dist", "tau_knot_dist", "p_tau", "tau_HP1", "tau_HP2", "tau_HP3", "tau_HP4", "maxAbsLogSD"), inits = list( w_tau = quote(rep(0, p_tau)), tauGP_mu = quote(0), tauGP_phi = quote(tau_HP3/100), tauGP_sigma = quote(tau_HP4/100) ), constraints_needed = c('tau_constraint1') ) ) ##============================================ ## Models for sigma ##============================================ sigma_model_list <- list( constant = list( ## 1. sigma_HP1 Standard deviation for the log-linear standard deviation ## 2. alpha Scalar; represents the standard deviation (constant over the domain) ## 3. ones N-vector of 1's code = quote({ log_sigma_vec[1:N] <- log(sqrt(alpha))*ones[1:N] alpha ~ dunif(0, sigma_HP1) }), constants_needed = c("ones", "sigma_HP1"), inits = list(alpha = quote(sigma_HP1/10)) ), logLinReg = list( ## 1. X_sigma N x p_sigma design matrix; leading column of 1's with (p_sigma - 1) other covariates ## 2. sigma_HP1 Standard deviation for the log-linear regression coefficients ## 3. p_sigma Number of design columns ## 4. alpha Vector of length p_sigma; represents log-linear regression coefficients code = quote({ log_sigma_vec[1:N] <- X_sigma[1:N,1:p_sigma] %*% alpha[1:p_sigma] for(l in 1:p_sigma){ alpha[l] ~ dnorm(0, sd = sigma_HP1) } # Constraint sigma_constraint1 ~ dconstraint( max(abs(log_sigma_vec[1:N])) < maxAbsLogSD ) }), constants_needed = c("X_sigma", "p_sigma", "sigma_HP1", "maxAbsLogSD"), inits = list(alpha = quote(rep(0, p_sigma))), constraints_needed = c('sigma_constraint1') ), approxGP = list( ## 1. sigma_HP1 Gaussian process standard deviation ## 2. sigma_HP2 Gaussian process mean ## 3. sigma_HP3 Gaussian process range ## 4. sigma_HP4 Gaussian process smoothness ## 5. ones N-vector of 1's ## 6. sigma_cross_dist N x p_sigma matrix of inter-point Euclidean distances, obs. coords vs. knot locations ## 7. sigma_knot_dist p_sigma x p_sigma matrix of inter-point Euclidean distances, knot locations ## 8. p_sigma Number of knot locations code = quote({ log_sigma_vec[1:N] <- sigmaGP_mu*ones[1:N] + sigmaGP_sigma*Pmat_sigma[1:N,1:p_sigma] %*% w_sigma[1:p_sigma] Pmat_sigma[1:N,1:p_sigma] <- matern_corr(sigma_cross_dist[1:N,1:p_sigma], sigmaGP_phi, sigma_HP2) Vmat_sigma[1:p_sigma,1:p_sigma] <- matern_corr(sigma_knot_dist[1:p_sigma,1:p_sigma], sigmaGP_phi, sigma_HP2) w_sigma_mean[1:p_sigma] <- 0*ones[1:p_sigma] w_sigma[1:p_sigma] ~ dmnorm( mean = w_sigma_mean[1:p_sigma], prec = Vmat_sigma[1:p_sigma,1:p_sigma] ) # Hyperparameters sigmaGP_mu ~ dnorm(0, sd = sigma_HP1) sigmaGP_phi ~ dunif(0, sigma_HP3) # Range parameter, GP sigmaGP_sigma ~ dunif(0, sigma_HP4) # SD parameter, GP # Constraint sigma_constraint1 ~ dconstraint( max(abs(log_sigma_vec[1:N])) < maxAbsLogSD ) }), constants_needed = c("ones", "sigma_knot_coords", "sigma_cross_dist", "sigma_knot_dist", "p_sigma", "sigma_HP1", "sigma_HP2", "sigma_HP3", "sigma_HP4", "maxAbsLogSD"), inits = list( w_sigma = quote(rep(0, p_sigma)), sigmaGP_mu = quote(0), sigmaGP_phi = quote(sigma_HP3/100), sigmaGP_sigma = quote(sigma_HP4/100)), constraints_needed = c('sigma_constraint1') ) ) ##============================================ ## Models for Sigma ##============================================ Sigma_model_list <- list( constant = list( ## 1. ones N-vector of 1's ## 2. Sigma_HP1 Upper bound for the eigenvalues ## 3. Sigma_coef{1,2,3} Vectors of length p_Sigma; represents the anisotropy components code = quote({ Sigma11[1:N] <- ones[1:N]*(Sigma_coef1*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef2*sin(Sigma_coef3)*sin(Sigma_coef3)) Sigma22[1:N] <- ones[1:N]*(Sigma_coef2*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef1*sin(Sigma_coef3)*sin(Sigma_coef3)) Sigma12[1:N] <- ones[1:N]*(Sigma_coef1*cos(Sigma_coef3)*sin(Sigma_coef3) - Sigma_coef2*cos(Sigma_coef3)*sin(Sigma_coef3)) Sigma_coef1 ~ dunif(0, Sigma_HP1[1]) # phi1 Sigma_coef2 ~ dunif(0, Sigma_HP1[1]) # phi2 Sigma_coef3 ~ dunif(0, 1.570796) # eta --> 1.570796 = pi/2 }), constants_needed = c("ones", "Sigma_HP1"), inits = list( Sigma_coef1 = quote(Sigma_HP1[1]/4), Sigma_coef2 = quote(Sigma_HP1[1]/4), Sigma_coef3 = 0.7853982 # pi/4 ) ), constantIso = list( # Isotropic version of ## 1. ones N-vector of 1's ## 2. Sigma_HP1 Standard deviation for the anisotropy components ## 3. Sigma_coef{1,2,3} Vectors of length p_Sigma; represents the anisotropy components code = quote({ Sigma11[1:N] <- ones[1:N]*Sigma_coef1 Sigma22[1:N] <- ones[1:N]*Sigma_coef1 Sigma12[1:N] <- ones[1:N]*0 Sigma_coef1 ~ dunif(0, Sigma_HP1[1]) # phi1 }), constants_needed = c("ones", "Sigma_HP1"), inits = list( Sigma_coef1 = quote(Sigma_HP1[1]/4) ) ), covReg = list( code = quote({ ## 1. X_Sigma N x p_Sigma design matrix; leading column of 1's with (p_Sigma - 1) other covariates ## 2. Sigma_HP1 Standard deviation for the covariance regression coefficients ## 3. p_Sigma Number of design columns ## 4. gamma1, gamma2 Vectors of length p_Sigma; represents covariance regression coefficients ## 5. psi11, psi22, rho Baseline covariance regression parameters ## 6. Sigma_HP2 Upper bound for the baseline covariance regression variances Sigma11[1:N] <- psi11*ones[1:N] + (X_Sigma[1:N,1:p_Sigma] %*% gamma1[1:p_Sigma])^2 Sigma12[1:N] <- rho*sqrt(psi11*psi22)*ones[1:N] + (X_Sigma[1:N,1:p_Sigma]%*%gamma1[1:p_Sigma])*(X_Sigma[1:N,1:p_Sigma]%*%gamma2[1:p_Sigma]) Sigma22[1:N] <- psi22*ones[1:N] + (X_Sigma[1:N,1:p_Sigma] %*% gamma2[1:p_Sigma])^2 psi11 ~ dunif(0, Sigma_HP2[1]) psi22 ~ dunif(0, Sigma_HP2[2]) rho ~ dunif(-1, 1) for(j in 1:p_Sigma){ gamma1[j] ~ dnorm(0, sd = Sigma_HP1[1]) gamma2[j] ~ dnorm(0, sd = Sigma_HP1[2]) } # Constraints: upper limits on eigen_comp1 and eigen_comp2 Sigma_constraint1 ~ dconstraint( max(Sigma11[1:N]) < maxAnisoRange ) Sigma_constraint2 ~ dconstraint( max(Sigma22[1:N]) < maxAnisoRange ) Sigma_constraint3 ~ dconstraint( min(Sigma11[1:N]*Sigma22[1:N] - Sigma12[1:N]*Sigma12[1:N]) > minAnisoDet ) }), constants_needed = c("ones", "X_Sigma", "p_Sigma", "Sigma_HP1", "Sigma_HP2", "maxAnisoRange", "minAnisoDet"), inits = list( psi11 = quote(Sigma_HP2[1]/4), psi22 = quote(Sigma_HP2[2]/4), rho = 0, gamma1 = quote(rep(0, p_Sigma)), gamma2 = quote(rep(0, p_Sigma)) ), constraints_needed = c('Sigma_constraint1', 'Sigma_constraint2', 'Sigma_constraint3') ), compReg = list( code = quote({ ## 1. X_Sigma N x p_Sigma design matrix; leading column of 1's with (p_Sigma - 1) other covariates ## 2. Sigma_HP1 Standard deviation for the component regression coefficients ## 3. p_Sigma Number of design columns ## 4. Sigma_coef{1,2,3} Vectors of length p_Sigma; represents component regression coefficients eigen_comp1[1:N] <- X_Sigma[1:N,1:p_Sigma] %*% Sigma_coef1[1:p_Sigma] eigen_comp2[1:N] <- X_Sigma[1:N,1:p_Sigma] %*% Sigma_coef2[1:p_Sigma] eigen_comp3[1:N] <- X_Sigma[1:N,1:p_Sigma] %*% Sigma_coef3[1:p_Sigma] Sigma11[1:N] <- inverseEigen(eigen_comp1[1:N], eigen_comp2[1:N], eigen_comp3[1:N], 1) Sigma12[1:N] <- inverseEigen(eigen_comp1[1:N], eigen_comp2[1:N], eigen_comp3[1:N], 3) Sigma22[1:N] <- inverseEigen(eigen_comp1[1:N], eigen_comp2[1:N], eigen_comp3[1:N], 2) for(j in 1:p_Sigma){ Sigma_coef1[j] ~ dnorm(0, sd = Sigma_HP1[1]) Sigma_coef2[j] ~ dnorm(0, sd = Sigma_HP1[1]) Sigma_coef3[j] ~ dnorm(0, sd = Sigma_HP1[1]) } # Constraints: upper limits on eigen_comp1 and eigen_comp2 Sigma_constraint1 ~ dconstraint( max(Sigma11[1:N]) < maxAnisoRange ) Sigma_constraint2 ~ dconstraint( max(Sigma22[1:N]) < maxAnisoRange ) Sigma_constraint3 ~ dconstraint( min(Sigma11[1:N]*Sigma22[1:N] - Sigma12[1:N]*Sigma12[1:N]) > minAnisoDet ) }), constants_needed = c("X_Sigma", "p_Sigma", "Sigma_HP1", "maxAnisoRange", "minAnisoDet"), inits = list( Sigma_coef1 = quote(c(log(maxAnisoRange/100), rep(0, p_Sigma-1))), Sigma_coef2 = quote(c(log(maxAnisoRange/100), rep(0, p_Sigma-1))), Sigma_coef3 = quote(rep(0, p_Sigma)) ), constraints_needed = c('Sigma_constraint1', 'Sigma_constraint2', 'Sigma_constraint3') ), compRegIso = list( # Isotropic version of compReg code = quote({ ## 1. X_Sigma N x p_Sigma design matrix; leading column of 1's with (p_Sigma - 1) other covariates ## 2. Sigma_HP1 Standard deviation for the component regression coefficients ## 3. p_Sigma Number of design columns ## 4. Sigma_coef{1,2,3} Vectors of length p_Sigma; represents component regression coefficients eigen_comp1[1:N] <- X_Sigma[1:N,1:p_Sigma] %*% Sigma_coef1[1:p_Sigma] Sigma11[1:N] <- exp(eigen_comp1[1:N]) Sigma22[1:N] <- exp(eigen_comp1[1:N]) Sigma12[1:N] <- ones[1:N]*0 for(j in 1:p_Sigma){ Sigma_coef1[j] ~ dnorm(0, sd = Sigma_HP1[1]) } # Constraints: upper limits on eigen_comp1 Sigma_constraint1 ~ dconstraint( max(Sigma11[1:N]) < maxAnisoRange ) }), constants_needed = c("ones", "X_Sigma", "p_Sigma", "Sigma_HP1", "maxAnisoRange"), inits = list( Sigma_coef1 = quote(c(log(maxAnisoRange/100), rep(0, p_Sigma-1))) ), constraints_needed = c('Sigma_constraint1') ), npApproxGP = list( code = quote({ ## 1. Sigma_HP1 3-vector; Gaussian process mean ## 2. Sigma_HP2 3-vector; Gaussian process smoothness ## 5. ones N-vector of 1's ## 6. dist N x N matrix of inter-point Euclidean distances ## 7. Sigma_cross_dist N x p_Sigma matrix of inter-point Euclidean distances, obs. coords vs. knot locations ## 8. Sigma_knot_dist p_Sigma x p_Sigma matrix of inter-point Euclidean distances, knot locations ## 9. p_Sigma Number of knot locations Sigma11[1:N] <- inverseEigen(eigen_comp1[1:N], eigen_comp2[1:N], eigen_comp3[1:N], 1) Sigma12[1:N] <- inverseEigen(eigen_comp1[1:N], eigen_comp2[1:N], eigen_comp3[1:N], 3) Sigma22[1:N] <- inverseEigen(eigen_comp1[1:N], eigen_comp2[1:N], eigen_comp3[1:N], 2) # approxGP1, approxGP2 eigen_comp1[1:N] <- SigmaGP_mu[1]*ones[1:N] + SigmaGP_sigma[1] * Pmat12_Sigma[1:N,1:p_Sigma] %*% w1_Sigma[1:p_Sigma] eigen_comp2[1:N] <- SigmaGP_mu[1]*ones[1:N] + SigmaGP_sigma[1] * Pmat12_Sigma[1:N,1:p_Sigma] %*% w2_Sigma[1:p_Sigma] Pmat12_Sigma[1:N,1:p_Sigma] <- matern_corr(Sigma_cross_dist[1:N,1:p_Sigma], SigmaGP_phi[1], Sigma_HP2[1]) Vmat12_Sigma[1:p_Sigma,1:p_Sigma] <- matern_corr(Sigma_knot_dist[1:p_Sigma,1:p_Sigma], SigmaGP_phi[1], Sigma_HP2[1]) w12_Sigma_mean[1:p_Sigma] <- 0*ones[1:p_Sigma] w1_Sigma[1:p_Sigma] ~ dmnorm( mean = w12_Sigma_mean[1:p_Sigma], prec = Vmat12_Sigma[1:p_Sigma,1:p_Sigma] ) w2_Sigma[1:p_Sigma] ~ dmnorm( mean = w12_Sigma_mean[1:p_Sigma], prec = Vmat12_Sigma[1:p_Sigma,1:p_Sigma] ) # approxGP3 eigen_comp3[1:N] <- SigmaGP_mu[2]*ones[1:N] + SigmaGP_sigma[2] * Pmat3_Sigma[1:N,1:p_Sigma] %*% w3_Sigma[1:p_Sigma] Pmat3_Sigma[1:N,1:p_Sigma] <- matern_corr(Sigma_cross_dist[1:N,1:p_Sigma], SigmaGP_phi[2], Sigma_HP2[2]) Vmat3_Sigma[1:p_Sigma,1:p_Sigma] <- matern_corr(Sigma_knot_dist[1:p_Sigma,1:p_Sigma], SigmaGP_phi[2], Sigma_HP2[2]) w3_Sigma_mean[1:p_Sigma] <- 0*ones[1:p_Sigma] w3_Sigma[1:p_Sigma] ~ dmnorm( mean = w3_Sigma_mean[1:p_Sigma], prec = Vmat3_Sigma[1:p_Sigma,1:p_Sigma] ) # Hyperparameters for(w in 1:2){ SigmaGP_mu[w] ~ dnorm(0, sd = Sigma_HP1[w]) SigmaGP_phi[w] ~ dunif(0, Sigma_HP3[w]) # Range parameter, GP SigmaGP_sigma[w] ~ dunif(0, Sigma_HP4[w]) # SD parameter, GP } # Constraints: upper limits on eigen_comp1 and eigen_comp2 Sigma_constraint1 ~ dconstraint( max(Sigma11[1:N]) < maxAnisoRange ) Sigma_constraint2 ~ dconstraint( max(Sigma22[1:N]) < maxAnisoRange ) Sigma_constraint3 ~ dconstraint( min(Sigma11[1:N]*Sigma22[1:N] - Sigma12[1:N]*Sigma12[1:N]) > minAnisoDet ) }), constants_needed = c("ones", "Sigma_HP1", "Sigma_HP2", "Sigma_HP3", "Sigma_HP4", "maxAnisoRange", "minAnisoDet", "Sigma_knot_coords", "Sigma_cross_dist", "Sigma_knot_dist", "p_Sigma"), inits = list( w1_Sigma = quote(rep(0,p_Sigma)), w2_Sigma = quote(rep(0,p_Sigma)), w3_Sigma = quote(rep(0,p_Sigma)), SigmaGP_mu = quote(rep(log(maxAnisoRange/100),2)), SigmaGP_phi = quote(rep(Sigma_HP3[1]/100,2)), SigmaGP_sigma = quote(rep(Sigma_HP4[1]/100,2)) ), constraints_needed = c('Sigma_constraint1', 'Sigma_constraint2', 'Sigma_constraint3') ), npApproxGPIso = list( code = quote({ ## 1. Sigma_HP1 3-vector; Gaussian process mean ## 2. Sigma_HP2 3-vector; Gaussian process smoothness ## 5. ones N-vector of 1's ## 6. dist N x N matrix of inter-point Euclidean distances ## 7. Sigma_cross_dist N x p_Sigma matrix of inter-point Euclidean distances, obs. coords vs. knot locations ## 8. Sigma_knot_dist p_Sigma x p_Sigma matrix of inter-point Euclidean distances, knot locations ## 9. p_Sigma Number of knot locations Sigma11[1:N] <- exp(eigen_comp1[1:N]) Sigma22[1:N] <- exp(eigen_comp1[1:N]) Sigma12[1:N] <- ones[1:N]*0 # approxGP1 eigen_comp1[1:N] <- SigmaGP_mu[1]*ones[1:N] + SigmaGP_sigma[1] * Pmat12_Sigma[1:N,1:p_Sigma] %*% w1_Sigma[1:p_Sigma] Pmat12_Sigma[1:N,1:p_Sigma] <- matern_corr(Sigma_cross_dist[1:N,1:p_Sigma], SigmaGP_phi[1], Sigma_HP2[1]) Vmat12_Sigma[1:p_Sigma,1:p_Sigma] <- matern_corr(Sigma_knot_dist[1:p_Sigma,1:p_Sigma], SigmaGP_phi[1], Sigma_HP2[1]) w12_Sigma_mean[1:p_Sigma] <- 0*ones[1:p_Sigma] w1_Sigma[1:p_Sigma] ~ dmnorm( mean = w12_Sigma_mean[1:p_Sigma], prec = Vmat12_Sigma[1:p_Sigma,1:p_Sigma] ) # Hyperparameters for(w in 1){ SigmaGP_mu[w] ~ dnorm(0, sd = Sigma_HP1[w]) SigmaGP_phi[w] ~ dunif(0, Sigma_HP3[w]) # Range parameter, GP SigmaGP_sigma[w] ~ dunif(0, Sigma_HP4[w]) # SD parameter, GP } # Constraints: upper limits on eigen_comp1 and eigen_comp2 Sigma_constraint1 ~ dconstraint( max(Sigma11[1:N]) < maxAnisoRange ) }), constants_needed = c("ones", "Sigma_HP1", "Sigma_HP2", "Sigma_HP3", "Sigma_HP4", "maxAnisoRange", "Sigma_knot_coords", "Sigma_cross_dist", "Sigma_knot_dist", "p_Sigma"), inits = list( w1_Sigma = quote(rep(0,p_Sigma)), SigmaGP_mu = quote(rep(log(maxAnisoRange/100),1)), SigmaGP_phi = quote(rep(Sigma_HP3[1]/100,1)), SigmaGP_sigma = quote(rep(Sigma_HP4[1]/100,1)) ), constraints_needed = c('Sigma_constraint1') ) ) ##============================================ ## Models for mu ##============================================ mu_model_list <- list( constant = list( ## 1. sigma_HP1 Standard deviation for the log-linear standard deviation ## 2. alpha Scalar; represents log-linear standard deviation (constant over the domain) ## 3. ones N-vector of 1's code = quote({ mu[1:N] <-beta*ones[1:N] beta ~ dnorm(0, sd = mu_HP1) }), constants_needed = c("ones", "mu_HP1"), inits = list(beta = 0)), linReg = list( ## 1. X_mu N x p_mu design matrix; leading column of 1's with (p_mu - 1) other covariates ## 2. p_mu Number of design columns ## 3. beta Vector of length p_mu; represents regression coefficients code = quote({ mu[1:N] <- X_mu[1:N,1:p_mu] %*% beta[1:p_mu] for(l in 1:p_mu){ beta[l] ~ dnorm(0, sd = mu_HP1) } }), constants_needed = c("X_mu", "p_mu", "mu_HP1"), inits = list(beta = quote(rep(0, p_mu)))), zero = list( ## 1. zeros N-vector of 0's code = quote({ mu[1:N] <- zeros[1:N] }), constants_needed = c("zeros"), inits = list() ) ) ##============================================ ## Models for likelihood ##============================================ likelihood_list <- list( fullGP = list( code = quote({ Cor[1:N,1:N] <- nsCorr(dist1_sq[1:N,1:N], dist2_sq[1:N,1:N], dist12[1:N,1:N], Sigma11[1:N], Sigma22[1:N], Sigma12[1:N], nu, d) sigmaMat[1:N,1:N] <- diag(exp(log_sigma_vec[1:N])) Cov[1:N, 1:N] <- sigmaMat[1:N,1:N] %*% Cor[1:N,1:N] %*% sigmaMat[1:N,1:N] C[1:N,1:N] <- Cov[1:N, 1:N] + diag(exp(log_tau_vec[1:N])^2) z[1:N] ~ dmnorm(mean = mu[1:N], cov = C[1:N,1:N]) }), constants_needed = c("N", "coords", "d", "dist1_sq", "dist2_sq", "dist12", "nu"), ## keep N, coords, d here inits = list() ), NNGP = list( code = quote({ AD[1:N,1:(k+1)] <- calculateAD_ns(dist1_3d[1:N,1:(k+1),1:(k+1)], dist2_3d[1:N,1:(k+1),1:(k+1)], dist12_3d[1:N,1:(k+1),1:(k+1)], Sigma11[1:N], Sigma22[1:N], Sigma12[1:N], log_sigma_vec[1:N], log_tau_vec[1:N], nID[1:N,1:k], N, k, nu, d) z[1:N] ~ dmnorm_nngp(mu[1:N], AD[1:N,1:(k+1)], nID[1:N,1:k], N, k) }), constants_needed = c("N", "coords", "d", "dist1_3d", "dist2_3d", "dist12_3d", "nID", "k", "nu"), ## keep N, coords, d here inits = list() ), SGV = list( code = quote({ U[1:num_NZ,1:3] <- calculateU_ns( dist1_3d[1:N,1:(k+1),1:(k+1)], dist2_3d[1:N,1:(k+1),1:(k+1)], dist12_3d[1:N,1:(k+1),1:(k+1)], Sigma11[1:N], Sigma22[1:N], Sigma12[1:N], log_sigma_vec[1:N], log_tau_vec[1:N], nu, nID[1:N,1:k], cond_on_y[1:N,1:k], N, k, d ) z[1:N] ~ dmnorm_sgv(mu[1:N], U[1:num_NZ,1:3], N, k) }), constants_needed = c("N", "coords", "d", "dist1_3d", "dist2_3d", "dist12_3d", "nID", "k", "nu", "cond_on_y", "num_NZ"), ## keep N, coords, d here inits = list() ) ) ##============================================ ## Setup ##============================================ if(is.null( tau_model_list[[ tau_model]])) stop("unknown specification for tau_model") if(is.null(sigma_model_list[[sigma_model]])) stop("unknown specification for sigma_model") if(is.null(Sigma_model_list[[Sigma_model]])) stop("unknown specification for Sigma_model") if(is.null( mu_model_list[[ mu_model]])) stop("unknown specification for mu_model") if(is.null( likelihood_list[[ likelihood]])) stop("unknown specification for likelihood") model_selections_list <- list( tau = tau_model_list [[tau_model]], sigma = sigma_model_list[[sigma_model]], Sigma = Sigma_model_list[[Sigma_model]], mu = mu_model_list [[mu_model]], likelihood = likelihood_list [[likelihood]] ) ## code code_template <- quote({ SIGMA_MODEL ## Log variance TAU_MODEL ## Log nugget -- gotta respect the nugget CAP_SIGMA_MODEL ## Anisotropy MU_MODEL ## Mean LIKELIHOOD_MODEL ## Likelihood }) code <- eval(substitute(substitute( CODE, list(TAU_MODEL = model_selections_list$tau$code, SIGMA_MODEL = model_selections_list$sigma$code, CAP_SIGMA_MODEL = model_selections_list$Sigma$code, CAP_SIGMA_MODEL = model_selections_list$Sigma$code, MU_MODEL = model_selections_list$mu$code, LIKELIHOOD_MODEL = model_selections_list$likelihood$code)), list(CODE = code_template))) if(missing(data)) stop("must provide data as 'data' argument") N <- length(data) if(missing(coords)) stop("must provide 'coords' argument, array of spatial coordinates") d <- ncol(coords) coords <- as.matrix(coords) sd_default <- 100 mu_default <- 0 matern_rho_default <- 1 matern_nu_default <- 5 maxDist <- 0 for(j in 1:d){ maxDist <- maxDist + (max(coords[,j]) - min(coords[,j]))^2 } maxDist <- sqrt(maxDist) # max(dist(coords)) if(N < 200){ ones_set <- rep(1,200); zeros_set <- rep(1,200) } else{ ones_set <- rep(1,N); zeros_set <- rep(1,N) } constants_defaults_list <- list( N = N, coords = coords, d = d, zeros = zeros_set, ones = ones_set, mu_HP1 = sd_default, ## standard deviation tau_HP1 = sd_default, ## standard deviation/upper bound for constant nugget tau_HP2 = matern_nu_default, ## approxGP smoothness tau_HP3 = maxDist, ## upper bound for approxGP range tau_HP4 = sd_default, ## upper bound for approxGP sd sigma_HP1 = sd_default, ## standard deviation sigma_HP2 = matern_nu_default, ## approxGP smoothness sigma_HP3 = maxDist, ## upper bound for approxGP range sigma_HP4 = sd_default, ## upper bound for approxGP sd Sigma_HP1 = rep(10,2), ## standard deviation/upper bound Sigma_HP2 = rep(10,2), ## uniform upper bound for covReg 'psi' parameters / latent approxGP smoothness Sigma_HP3 = rep(maxDist,2), ## upper bound for approxGP range Sigma_HP4 = rep(sd_default, 2), ## upper bound for approxGP sd maxAbsLogSD = 10, ## logSD must live between +/- maxAbsLogSD maxAnisoRange = maxDist, ## maximum value for the diagonal elements of the anisotropy process minAnisoDet = 1e-5, ## lower bound for the determinant of the anisotropy process nu = matern_nu_default ## Process smoothness parameter ) ## use the isotropic model? useIsotropic <- (Sigma_model %in% c("constantIso", "compRegIso", "npApproxGPIso")) ## initialize constants_to_use with constants_defaults_list constants_to_use <- constants_defaults_list ## update constants_to_use with those arguments provided via ... dotdotdot <- list(...) ## make sure all ... arguments were provided with names if(length(dotdotdot) > 0 && (is.null(names(dotdotdot)) || any(names(dotdotdot) == ""))) stop("Only named arguemnts should be provided through ... argument") constants_to_use[names(dotdotdot)] <- dotdotdot ## add 'constants' argument to constants_to_use list: ## if provided, make sure 'constants' argument is a named list if(!missing(constants)) { if(length(constants) > 0 && (is.null(names(constants)) || any(names(constants) == ""))) stop("All elements in constants list argument must be named") constants_to_use[names(constants)] <- constants } ## generate and add dist1_sq, dist2_sq, and dist12 arrays to constants_to_use list if(likelihood == 'fullGP') { dist_list <- nsDist(coords = coords, isotropic = useIsotropic) } else { ## likelihood is NNGP, or SGV: if(is.null(constants_to_use$k)) stop(paste0('missing k constants argument for ', likelihood, ' likelihood'))
mmd.seed <- sample(1e5, 1) # Set seed for reproducibility (randomness in orderCoordinatesMMD function)
if(likelihood == 'NNGP') {
# cat("\nOrdering the prediction locations and determining neighbors for NNGP (this may take a minute).\n")
# Re-order the coordinates/data
coords_mmd <- orderCoordinatesMMD(coords)
ord <- coords_mmd$orderedIndicesNoNA coords <- coords[ord,] data <- data[ord] # Set neighbors and calculate distances nID <- determineNeighbors(coords, constants_to_use$k)
constants_to_use$nID <- nID dist_list <- nsDist3d(coords = coords, nID = nID, isotropic = useIsotropic) } if(likelihood == 'SGV') { # cat("\nOrdering the prediction locations and determining neighbors/conditioning sets for SGV (this may take a minute).\n") setupSGV <- sgvSetup(coords = coords, k = constants_to_use$k, seed = mmd.seed)
constants_to_use$nID <- setupSGV$nID_ord
constants_to_use$cond_on_y <- setupSGV$condition_on_y_ord
constants_to_use$num_NZ <- setupSGV$num_NZ
ord <- setupSGV$ord # Re-order the coordinates/data coords <- coords[ord,] data <- data[ord] dist_list <- nsDist3d(coords = coords, nID = setupSGV$nID_ord, isotropic = useIsotropic)
}
# Re-order any design matrices
if(!is.null(constants_to_use$X_tau)) constants_to_use$X_tau <- constants_to_use$X_tau[ord,] if(!is.null(constants_to_use$X_sigma)) constants_to_use$X_sigma <- constants_to_use$X_sigma[ord,]
if(!is.null(constants_to_use$X_Sigma)) constants_to_use$X_Sigma <- constants_to_use$X_Sigma[ord,] if(!is.null(constants_to_use$X_mu)) constants_to_use$X_mu <- constants_to_use$X_mu[ord,]
}
constants_to_use$coords <- coords constants_to_use[names(dist_list)] <- dist_list ## if any models use approxGP: calculate XX_knot_dist and XX_cross_dist (coords already re-ordered for NNGP/SGV) if( tau_model == 'approxGP' ) { if(is.null(constants_to_use$tau_knot_coords)) stop(paste0('missing tau_knot_coords for tau_model: approxGP'))
constants_to_use$tau_knot_dist <- sqrt(nsDist(coords = constants_to_use$tau_knot_coords, isotropic = TRUE)$dist1_sq) constants_to_use$tau_cross_dist <- sqrt(nsCrossdist(Pcoords = coords, coords = constants_to_use$tau_knot_coords, isotropic = TRUE)$dist1_sq)
}
if( sigma_model == 'approxGP' ) {
if(is.null(constants_to_use$sigma_knot_coords)) stop(paste0('missing sigma_knot_coords for sigma_model: approxGP')) constants_to_use$sigma_knot_dist <- sqrt(nsDist(coords = constants_to_use$sigma_knot_coords, isotropic = TRUE)$dist1_sq)
constants_to_use$sigma_cross_dist <- sqrt(nsCrossdist(Pcoords = coords, coords = constants_to_use$sigma_knot_coords, isotropic = TRUE)$dist1_sq) } if( Sigma_model %in% c('npApproxGP', 'npApproxGPIso') ) { if(is.null(constants_to_use$Sigma_knot_coords)) stop(paste0('missing Sigma_knot_coords for Sigma_model: ', Sigma_model))
constants_to_use$Sigma_knot_dist <- sqrt(nsDist(coords = constants_to_use$Sigma_knot_coords, isotropic = TRUE)$dist1_sq) constants_to_use$Sigma_cross_dist <- sqrt(nsCrossdist(Pcoords = coords, coords = constants_to_use$Sigma_knot_coords, isotropic = TRUE)$dist1_sq)
}

## add the following (derived numbers of columns) to constants_to_use:
## p_tau:
if(!is.null(constants_to_use$X_tau)) constants_to_use$p_tau <- ncol(constants_to_use$X_tau) if(!is.null(constants_to_use$tau_cross_dist)) constants_to_use$p_tau <- ncol(constants_to_use$tau_cross_dist)
## p_sigma:
if(!is.null(constants_to_use$X_sigma)) constants_to_use$p_sigma <- ncol(constants_to_use$X_sigma) if(!is.null(constants_to_use$sigma_cross_dist)) constants_to_use$p_sigma <- ncol(constants_to_use$sigma_cross_dist)
## p_Sigma:
if(!is.null(constants_to_use$X_Sigma)) constants_to_use$p_Sigma <- ncol(constants_to_use$X_Sigma) if(!is.null(constants_to_use$Sigma_cross_dist)) constants_to_use$p_Sigma <- ncol(constants_to_use$Sigma_cross_dist)
## p_mu:
if(!is.null(constants_to_use$X_mu)) constants_to_use$p_mu <- ncol(constants_to_use$X_mu) ## get a vector of all the constants we need for this model constants_needed <- unique(unlist(lapply(model_selections_list, function(x) x$constants_needed), use.names = FALSE))
## check if we're missing any constants we need, and throw an error if any are missing
constants_missing <- setdiff(constants_needed, names(constants_to_use))
if(length(constants_missing) > 0) {
stop(paste0("Missing values for the following model constants: ",
paste0(constants_missing, collapse = ", "),
".\nThese values should be provided as named arguments, or named elements in the constants list argument"))
}

## generate the constants list
constants <- constants_to_use[constants_needed]

## append the mmd.seed and ord for SGV/NNGP
if(likelihood != 'fullGP') {
constants$mmd.seed <- mmd.seed constants$ord <- ord
}

## ensure Sigma_HPX parameters are vectors of length 2
if(!is.null(constants$Sigma_HP1)){ if(length(constants$Sigma_HP1) == 1) constants$Sigma_HP1 <- rep(constants$Sigma_HP1, 2)
}
if(!is.null(constants$Sigma_HP2)){ if(length(constants$Sigma_HP2) == 1) constants$Sigma_HP2 <- rep(constants$Sigma_HP2, 2)
}
if(!is.null(constants$Sigma_HP3)){ if(length(constants$Sigma_HP3) == 1) constants$Sigma_HP3 <- rep(constants$Sigma_HP3, 2)
}
if(!is.null(constants$Sigma_HP4)){ if(length(constants$Sigma_HP4) == 1) constants$Sigma_HP41 <- rep(constants$Sigma_HP4, 2)
}

## gather constraints_needed data
constraints_needed <- unique(unlist(lapply(model_selections_list, function(x) x$constraints_needed), use.names = FALSE)) constraints_data <- as.list(rep(1, length(constraints_needed))) names(constraints_data) <- constraints_needed ## data data <- c(list(z = data), constraints_data) ## inits inits_uneval <- do.call("c", unname(lapply(model_selections_list, function(x) x$inits)))
inits <- lapply(inits_uneval, function(x) eval(x, envir = constants))

# if(returnModelComponents) return(list(code=code, constants=constants, data=data, inits=inits))

## generate the "name" for the nimble model object, containing which submodels were used
thisName <- paste0(
'tau='       , tau_model  , '_',
'sigma='     , sigma_model, '_',
'Sigma='     , Sigma_model, '_',
'mu='        , mu_model   , '_',
'likelihood=', likelihood
)

## register custom NNGP or SGV distributions (if necessary),
## importantly, using mixedSizes = TRUE to avoid warnings
if(likelihood == 'NNGP') {
registerDistributions(list(
dmnorm_nngp = list(
BUGSdist = 'dmnorm_nngp(mean, AD, nID, N, k)',
types = c('value = double(1)', 'mean = double(1)', 'AD = double(2)', 'nID = double(2)', 'N = double()', 'k = double()'),
mixedSizes = TRUE)
), verbose = FALSE)
}
if(likelihood == 'SGV') {
registerDistributions(list(
dmnorm_sgv = list(
BUGSdist = 'dmnorm_sgv(mean, U, N, k)',
types = c('value = double(1)', 'mean = double(1)', 'U = double(2)', 'N = double()', 'k = double()'),
mixedSizes = TRUE)
), verbose = FALSE)
}

## NIMBLE model object
Rmodel <- nimbleModel(code, constants, data, inits, name = thisName)
lp <- Rmodel$getLogProb() if(is(lp, 'try-error') || is.nan(lp) || is.na(lp) || abs(lp) == Inf) stop('model not properly initialized') ## store 'constants' list into Rmodel$isDataEnv
Rmodel$isDataEnv$.BayesNSGP_constants_list <- constants

## using the nsgpModel() argument monitorAllSampledNodes,
## set nimble package option: MCMCmonitorAllSampledNodes,
## so that latent process values are monitored by default, for use in predicition
nimbleOptions(MCMCmonitorAllSampledNodes = monitorAllSampledNodes)

return(Rmodel)
}

#==============================================================================
# Posterior prediction for the NSGP
#==============================================================================

# ROxygen comments ----
#' Posterior prediction for the NSGP
#'
#' \code{nsgpPredict} conducts posterior prediction for MCMC samples generated
#' using nimble and nsgpModel.
#'
#' @param model A NSGP nimble object; the output of \code{nsgpModel}.
#' @param samples A matrix of \code{J} rows, each is an MCMC sample of the
#' parameters corresponding to the specification in \code{nsgpModel}.
#' @param coords.predict M x d matrix of prediction coordinates.
#' @param predict.process Logical; determines whether the prediction corresponds to
#' the y(Â·) process (\code{TRUE}) or z(Â·) (\code{FALSE}; this would likely
#' only be used for, e.g., cross-validation).
#' @param constants  An optional list of contants to use for prediction;
#' alternatively, additional arguments can be passed to the function via the
#' ... argument.
#' @param seed An optional random seed argument for reproducibility.
#' @param ... Additional arguments can be passed to the function; for example,
#' as an alternative to the \code{constants} list, items can be passed directly
#' via this argument.
#'
#' @return The output of the function is a list with two elements: \code{obs},
#' a matrix of \code{J} posterior predictive samples for the N observed
#' locations (only for \code{likelihood = "SGV"}, which produces predictions
#' for the observed locations by default; this element is \code{NULL}
#' otherwise); and \code{pred}, a corresponding matrix of posterior predictive
#' samples for the prediction locations. Ordering and neighbor selection
#' for the prediction coordinates in the SGV likelihood are conducted
#' internally, as with \code{nsgpModel}.
#'
#' @examples
#' \donttest{
#' # Generate some data: stationary/isotropic
#' N <- 100
#' coords <- matrix(runif(2*N), ncol = 2)
#' alpha_vec <- rep(log(sqrt(1)), N) # Log process SD
#' delta_vec <- rep(log(sqrt(0.05)), N) # Log nugget SD
#' Sigma11_vec <- rep(0.4, N) # Kernel matrix element 1,1
#' Sigma22_vec <- rep(0.4, N) # Kernel matrix element 2,2
#' Sigma12_vec <- rep(0, N) # Kernel matrix element 1,2
#' mu_vec <- rep(0, N) # Mean
#' nu <- 0.5 # Smoothness
#' dist_list <- nsDist(coords)
#' Cor_mat <- nsCorr( dist1_sq = dist_list$dist1_sq, dist2_sq = dist_list$dist2_sq,
#'                    dist12 = dist_list$dist12, Sigma11 = Sigma11_vec, #' Sigma22 = Sigma22_vec, Sigma12 = Sigma12_vec, nu = nu ) #' Cov_mat <- diag(exp(alpha_vec)) %*% Cor_mat %*% diag(exp(alpha_vec)) #' D_mat <- diag(exp(delta_vec)^2) #' set.seed(110) #' data <- as.numeric(mu_vec + t(chol(Cov_mat + D_mat)) %*% rnorm(N)) #' # Set up constants #' constants <- list( nu = 0.5, Sigma_HP1 = 2 ) #' # Defaults: tau_model = "constant", sigma_model = "constant", mu_model = "constant", #' # and Sigma_model = "constant" #' Rmodel <- nsgpModel(likelihood = "fullGP", constants = constants, coords = coords, data = data ) #' conf <- configureMCMC(Rmodel) #' Rmcmc <- buildMCMC(conf) #' Cmodel <- compileNimble(Rmodel) #' Cmcmc <- compileNimble(Rmcmc, project = Rmodel) #' samples <- runMCMC(Cmcmc, niter = 200, nburnin = 100) #' # Prediction #' predCoords <- as.matrix(expand.grid(seq(0,1,l=10),seq(0,1,l=10))) #' postpred <- nsgpPredict( model = Rmodel, samples = samples, coords.predict = predCoords ) #' } #' #' @export #' nsgpPredict <- function(model, samples, coords.predict, predict.process = TRUE, constants, seed = 0, ... ) { if(!nimble::is.model(model)) stop('first argument must be NSGP NIMBLE model object') Rmodel <- if(nimble::is.Rmodel(model)) model else model$Rmodel
if(!nimble::is.Rmodel(Rmodel)) stop('something went wrong')

model_constants <- Rmodel$isDataEnv$.BayesNSGP_constants_list
coords <- model_constants$coords mcmc_samples <- samples predCoords <- as.matrix(coords.predict) ## extract the "submodel" information from the nimble model object "name" thisName <- Rmodel$getModelDef()$name modelsList <- lapply(strsplit(thisName, '_')[[1]], function(x) strsplit(x, '=')[[1]][2]) names(modelsList) <- sapply(strsplit(thisName, '_')[[1]], function(x) strsplit(x, '=')[[1]][1]) ## avaialble for use: ## modelsList$tau
## modelsList$sigma ## modelsList$Sigma
## modelsList$mu ## modelsList$likelihood

## order predCoords for SGV
if( modelsList$likelihood == "SGV" ) { message("\nOrdering the prediction locations and determining neighbors/conditioning sets for SGV (this may take a minute).\n") pred.mmd.seed <- sample(1e5, 1) predSGV_setup <- sgvSetup(coords = coords, coords_pred = predCoords, k = model_constants$k,
pred.seed = pred.mmd.seed, order_coords = FALSE)
prednID_SGV <- predSGV_setup$nID_ord obs_ord <- predSGV_setup$ord
pred_ord <- predSGV_setup$ord_pred predCoords <- predCoords[pred_ord,] } constants_to_use <- list() ## update constants_to_use with those arguments provided via ... dotdotdot <- list(...) ## make sure all ... arguments were provided with names if(length(dotdotdot) > 0 && (is.null(names(dotdotdot)) || any(names(dotdotdot) == ""))) stop("Only named arguemnts should be provided through ... argument") constants_to_use[names(dotdotdot)] <- dotdotdot ## add 'constants' argument to constants_to_use list: ## if provided, make sure 'constants' argument is a named list if(!missing(constants)) { if(length(constants) > 0 && (is.null(names(constants)) || any(names(constants) == ""))) stop("All elements in constants list argument must be named") constants_to_use[names(constants)] <- constants } ## calculate XX_cross_dist_pred if necessary if( modelsList$tau == 'approxGP' ) {
if(is.null(model_constants$tau_knot_coords)) stop(paste0('missing tau_knot_coords for tau_model: approxGP')) constants_to_use$tau_cross_dist_pred <- sqrt(nsCrossdist(Pcoords = predCoords, coords = model_constants$tau_knot_coords, isotropic = TRUE)$dist1_sq)
}
if( modelsList$sigma == 'approxGP' ) { if(is.null(model_constants$sigma_knot_coords)) stop(paste0('missing sigma_knot_coords for sigma_model: approxGP'))
constants_to_use$sigma_cross_dist_pred <- sqrt(nsCrossdist(Pcoords = predCoords, coords = model_constants$sigma_knot_coords, isotropic = TRUE)$dist1_sq) } if( modelsList$Sigma %in% c('npApproxGP', 'npApproxGPIso') ) {
if(is.null(model_constants$Sigma_knot_coords)) stop(paste0('missing Sigma_knot_coords for Sigma_model: ', modelsList$Sigma))
constants_to_use$Sigma_cross_dist_pred <- sqrt(nsCrossdist(Pcoords = predCoords, coords = model_constants$Sigma_knot_coords, isotropic = TRUE)$dist1_sq) } ## check for discrepancies in any duplicates, between ## constants_to_use provided here, and model_constants from Rmodel duplicatedNames <- intersect(names(constants_to_use), names(model_constants)) discrepancies <- character() if(length(duplicatedNames) > 0) { for(name in duplicatedNames) { if(!identical(constants_to_use[[name]], model_constants[[name]])) discrepancies <- c(discrepancies, name) } } if(length(discrepancies) > 0) stop(paste0('Inconsistent values were provided for the following constants: ', paste0(discrepancies, collapse = ', '))) ## move original model_constants from nsgpModel into constants_to_use: constants_to_use[names(model_constants)] <- model_constants ## determine the constants needed predictConstantsNeeded <- list( tau = list( constant = character(), logLinReg = c('X_tau', 'PX_tau'), approxGP = c('p_tau', 'tau_cross_dist', 'tau_cross_dist_pred', 'tau_HP2')), sigma = list( constant = character(), logLinReg = c('X_sigma', 'PX_sigma'), approxGP = c('p_sigma', 'sigma_cross_dist', 'sigma_cross_dist_pred', 'sigma_HP2')), Sigma = list( constant = character(), constantIso = character(), covReg = c('X_Sigma', 'PX_Sigma'), compReg = c('X_Sigma', 'PX_Sigma'), compRegIso = c('X_Sigma', 'PX_Sigma'), npApproxGP = c('p_Sigma', 'Sigma_cross_dist', 'Sigma_cross_dist_pred', 'Sigma_HP2'), npApproxGPIso = c('p_Sigma', 'Sigma_cross_dist', 'Sigma_cross_dist_pred', 'Sigma_HP2')), mu = list( constant = character(), linReg = c('X_mu', 'PX_mu'), zero = character()), likelihood = list( fullGP = character(), NNGP = character(), SGV = character()) ) constants_needed <- unique(unlist(lapply(1:length(modelsList), function(i) predictConstantsNeeded[[names(modelsList)[i]]][[modelsList[[i]]]] ))) ## check if we're missing any constants we need, and throw an error if any are missing constants_missing <- setdiff(constants_needed, names(constants_to_use)) if(length(constants_missing) > 0) { stop(paste0("Missing values for the following model constants: ", paste0(constants_missing, collapse = ", "), ".\nThese values should be provided as named arguments, or named elements in the constants list argument")) } ## generate the constants list ## do NOT truncate constants list like this: ##constants <- constants_to_use[constants_needed] constants <- constants_to_use d <- constants$d
z <- Rmodel$z if( modelsList$likelihood == "fullGP" ){ # Predictions for the full GP likelihood
# Extract needed variables from constants
dist1_sq <- constants$dist1_sq dist2_sq <- constants$dist2_sq
dist12 <- constants$dist12 N <- constants$N # number of observed locations
nu <- constants$nu # Prediction distances if(dist2_sq[1,1] == -1){ # Isotropic Pdist <- nsDist(predCoords, isotropic = TRUE) Xdist <- nsCrossdist(coords, predCoords, isotropic = TRUE) } else{ Pdist <- nsDist(predCoords) Xdist <- nsCrossdist(coords, predCoords) } Pdist1_sq <- Pdist$dist1_sq
Pdist2_sq <- Pdist$dist2_sq Pdist12 <- Pdist$dist12
M <- nrow(Pdist1_sq) # number of prediction locations
# Cross distances
Xdist1_sq <- Xdist$dist1_sq Xdist2_sq <- Xdist$dist2_sq
Xdist12 <- Xdist$dist12 postPredDrawsCols <- M } else if( modelsList$likelihood == "NNGP" ){ # Predictions for the NNGP likelihood
# "Local kriging" only possible for NNGP
# Extract needed variables from constants
dist1_3d <- constants$dist1_3d dist2_3d <- constants$dist2_3d
dist12_3d <- constants$dist12_3d N <- constants$N # number of observed locations
k <- constants$k # number of neighbors nu <- constants$nu
# Prediction/cross distances
P_nID <- get.knnx(coords, predCoords, k = k)$nn.index # Prediction NN if(dist2_3d[1,1,1] == -1){ Pdist <- nsCrossdist3d(coords, predCoords, P_nID, isotropic = TRUE) } else{ Pdist <- nsCrossdist3d(coords, predCoords, P_nID) } Pdist1_3d <- Pdist$dist1_3d
Pdist2_3d <- Pdist$dist2_3d Pdist12_3d <- Pdist$dist12_3d
M <- dim(Pdist1_3d)[1] # number of prediction locations
postPredDrawsCols <- M
} else if( modelsList$likelihood == "SGV" ){ # Predictions for the SGV likelihood if(predict.process == FALSE){ stop("Prediction for Z(.) not available with SGV.") } # Extract needed variables from constants dist1_3d <- constants$dist1_3d
dist2_3d <- constants$dist2_3d dist12_3d <- constants$dist12_3d
N <- constants$N # number of observed locations k <- constants$k # number of neighbors
nu <- constants$nu # Prediction setup if(dist2_3d[1,1,1] == -1){ preddist_SGV <- nsDist3d( predSGV_setup$coords_ord, prednID_SGV, isotropic = TRUE )
} else{
preddist_SGV <- nsDist3d( predSGV_setup$coords_ord, prednID_SGV ) } Alldist1_3d <- preddist_SGV$dist1_3d
Alldist2_3d <- preddist_SGV$dist2_3d Alldist12_3d <- preddist_SGV$dist12_3d
M <- dim(predCoords)[1] # number of prediction locations
# REORDER INPUTS (for model != "constant") =================
if( modelsList$tau == "logLinReg" ){ constants$PX_tau <- constants$PX_tau[pred_ord,] } if( modelsList$sigma == "logLinReg" ){
constants$PX_sigma <- constants$PX_sigma[pred_ord,]
}
if( modelsList$Sigma == "covReg" | modelsList$Sigma == "compReg" ){
constants$PX_Sigma <- constants$PX_Sigma[pred_ord,]
}
if( modelsList$mu == "linReg" ){ constants$PX_mu <- constants$PX_mu[pred_ord,] } postPredDrawsCols <- M+N } else stop('') J <- nrow(mcmc_samples) # Posterior draws - storage postPredDraws <- matrix(NA, nrow = J, ncol = postPredDrawsCols) nimCat("|-------------|-------------|-------------|-------------|") nimCat("\n|") for(j in 1:J){ # Loop over MCMC samples samp_j <- mcmc_samples[j,] # Calculate log_tau_vec and Plog_tau_vec ====================== if( modelsList$tau == "constant" ){
# Required constants: none
log_tau_vec_j <- log(sqrt(samp_j["delta"]))*rep(1,N)
Plog_tau_vec_j <- log(sqrt(samp_j["delta"]))*rep(1,M)
}
if( modelsList$tau == "logLinReg" ){ # Required constants: X_tau, PX_tau X_tau <- constants$X_tau
PX_tau <- constants$PX_tau log_tau_vec_j <- as.numeric(X_tau %*% samp_j[paste("delta[",1:ncol(X_tau),"]",sep = "")]) Plog_tau_vec_j <- as.numeric(PX_tau %*% samp_j[paste("delta[",1:ncol(PX_tau),"]",sep = "")]) } if( modelsList$tau == "approxGP" ){
# Required constants: p_tau, tau_cross_dist, tau_cross_dist_pred, tau_HP2
p_tau <- constants$p_tau tau_cross_dist_obs <- constants$tau_cross_dist
tau_cross_dist_pred <- constants$tau_cross_dist_pred tau_HP2 <- constants$tau_HP2
w_tau_j <- as.numeric(samp_j[paste("w_tau[",1:p_tau,"]",sep = "")])

# Obs locations
Pmat_tau_obs_j <- matern_corr(tau_cross_dist_obs, samp_j["tauGP_phi"], tau_HP2)
log_tau_vec_j <- as.numeric(samp_j["tauGP_mu"]*rep(1,N) + samp_j["tauGP_sigma"] * Pmat_tau_obs_j %*% w_tau_j)

# Pred locations
Pmat_tau_pred_j <- matern_corr(tau_cross_dist_pred, samp_j["tauGP_phi"], tau_HP2)
Plog_tau_vec_j <- as.numeric(samp_j["tauGP_mu"]*rep(1,M) + samp_j["tauGP_sigma"] * Pmat_tau_pred_j %*% w_tau_j)
}

# Calculate log_sigma_vec and Plog_sigma_vec ==================
if( modelsList$sigma == "constant" ){ # Required constants: none log_sigma_vec_j <- log(sqrt(samp_j["alpha"]))*rep(1,N) Plog_sigma_vec_j <- log(sqrt(samp_j["alpha"]))*rep(1,M) } if( modelsList$sigma == "logLinReg" ){
# Required constants: X_sigma, PX_sigma
X_sigma <- constants$X_sigma PX_sigma <- constants$PX_sigma
log_sigma_vec_j <- as.numeric(X_sigma %*% samp_j[paste("alpha[",1:ncol(X_sigma),"]",sep = "")])
Plog_sigma_vec_j <- as.numeric(PX_sigma %*% samp_j[paste("alpha[",1:ncol(PX_sigma),"]",sep = "")])
}
if( modelsList$sigma == "approxGP" ){ # Required constants: p_sigma, sigma_cross_dist, sigma_cross_dist_pred, sigma_HP2 p_sigma <- constants$p_sigma
sigma_cross_dist_obs <- constants$sigma_cross_dist sigma_cross_dist_pred <- constants$sigma_cross_dist_pred
sigma_HP2 <- constants$sigma_HP2 w_sigma_j <- as.numeric(samp_j[paste("w_sigma[",1:p_sigma,"]",sep = "")]) # Obs locations Pmat_sigma_obs_j <- matern_corr(sigma_cross_dist_obs, samp_j["sigmaGP_phi"], sigma_HP2) log_sigma_vec_j <- as.numeric(samp_j["sigmaGP_mu"]*rep(1,N) + samp_j["sigmaGP_sigma"] * Pmat_sigma_obs_j %*% w_sigma_j) # Pred locations Pmat_sigma_pred_j <- matern_corr(sigma_cross_dist_pred, samp_j["sigmaGP_phi"], sigma_HP2) Plog_sigma_vec_j <- as.numeric(samp_j["sigmaGP_mu"]*rep(1,M) + samp_j["sigmaGP_sigma"] * Pmat_sigma_pred_j %*% w_sigma_j) } # Calculate SigmaXX and PSigmaXX ============================== if( modelsList$Sigma == "constant" ){
# Required constants: none
Sigma_coef1 <- samp_j["Sigma_coef1"]
Sigma_coef2 <- samp_j["Sigma_coef2"]
Sigma_coef3 <- samp_j["Sigma_coef3"]
Sigma11_j <- rep(1,N)*(Sigma_coef1*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef2*sin(Sigma_coef3)*sin(Sigma_coef3))
Sigma22_j <- rep(1,N)*(Sigma_coef2*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef1*sin(Sigma_coef3)*sin(Sigma_coef3))
Sigma12_j <- rep(1,N)*(Sigma_coef1*cos(Sigma_coef3)*sin(Sigma_coef3) - Sigma_coef2*cos(Sigma_coef3)*sin(Sigma_coef3))
PSigma11_j <- rep(1,M)*(Sigma_coef1*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef2*sin(Sigma_coef3)*sin(Sigma_coef3))
PSigma22_j <- rep(1,M)*(Sigma_coef2*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef1*sin(Sigma_coef3)*sin(Sigma_coef3))
PSigma12_j <- rep(1,M)*(Sigma_coef1*cos(Sigma_coef3)*sin(Sigma_coef3) - Sigma_coef2*cos(Sigma_coef3)*sin(Sigma_coef3))
}
if( modelsList$Sigma == "constantIso" ){ # Required constants: none Sigma_coef1 <- samp_j["Sigma_coef1"] Sigma_coef2 <- samp_j["Sigma_coef1"] Sigma_coef3 <- 0 Sigma11_j <- rep(1,N)*(Sigma_coef1*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef2*sin(Sigma_coef3)*sin(Sigma_coef3)) Sigma22_j <- rep(1,N)*(Sigma_coef2*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef1*sin(Sigma_coef3)*sin(Sigma_coef3)) Sigma12_j <- rep(1,N)*(Sigma_coef1*cos(Sigma_coef3)*sin(Sigma_coef3) - Sigma_coef2*cos(Sigma_coef3)*sin(Sigma_coef3)) PSigma11_j <- rep(1,M)*(Sigma_coef1*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef2*sin(Sigma_coef3)*sin(Sigma_coef3)) PSigma22_j <- rep(1,M)*(Sigma_coef2*cos(Sigma_coef3)*cos(Sigma_coef3) + Sigma_coef1*sin(Sigma_coef3)*sin(Sigma_coef3)) PSigma12_j <- rep(1,M)*(Sigma_coef1*cos(Sigma_coef3)*sin(Sigma_coef3) - Sigma_coef2*cos(Sigma_coef3)*sin(Sigma_coef3)) } if( modelsList$Sigma == "covReg" ){
# Required constants: X_Sigma, PX_Sigma
X_Sigma <- constants$X_Sigma PX_Sigma <- constants$PX_Sigma
Sigma11_j <- as.numeric(samp_j["psi11"]*rep(1,N) + (X_Sigma %*% samp_j[paste("gamma1[",1:ncol(X_Sigma),"]",sep = "")])^2)
Sigma12_j <- as.numeric(samp_j["rho"]*sqrt(samp_j["psi11"]*samp_j["psi22"])*rep(1,N) + (X_Sigma %*% samp_j[paste("gamma1[",1:ncol(X_Sigma),"]",sep = "")])*(X_Sigma %*% samp_j[paste("gamma2[",1:ncol(X_Sigma),"]",sep = "")]))
Sigma22_j <- as.numeric(samp_j["psi22"]*rep(1,N) + (X_Sigma %*% samp_j[paste("gamma2[",1:ncol(X_Sigma),"]",sep = "")])^2)
PSigma11_j <- as.numeric(samp_j["psi11"]*rep(1,M) + (PX_Sigma %*% samp_j[paste("gamma1[",1:ncol(PX_Sigma),"]",sep = "")])^2)
PSigma12_j <- as.numeric(samp_j["rho"]*sqrt(samp_j["psi11"]*samp_j["psi22"])*rep(1,M) + (PX_Sigma %*% samp_j[paste("gamma1[",1:ncol(PX_Sigma),"]",sep = "")])*(PX_Sigma %*% samp_j[paste("gamma2[",1:ncol(PX_Sigma),"]",sep = "")]))
PSigma22_j <- as.numeric(samp_j["psi22"]*rep(1,M) + (PX_Sigma %*% samp_j[paste("gamma2[",1:ncol(PX_Sigma),"]",sep = "")])^2)
}
if( modelsList$Sigma == "compReg" ){ # Required constants: X_Sigma, PX_Sigma X_Sigma <- constants$X_Sigma
PX_Sigma <- constants$PX_Sigma eigen_comp1_j <- X_Sigma %*% samp_j[paste("Sigma_coef1[",1:ncol(X_Sigma),"]",sep = "")] eigen_comp2_j <- X_Sigma %*% samp_j[paste("Sigma_coef2[",1:ncol(X_Sigma),"]",sep = "")] eigen_comp3_j <- X_Sigma %*% samp_j[paste("Sigma_coef3[",1:ncol(X_Sigma),"]",sep = "")] Sigma11_j <- as.numeric(inverseEigen(eigen_comp1_j, eigen_comp2_j, eigen_comp3_j, 1)) Sigma12_j <- as.numeric(inverseEigen(eigen_comp1_j, eigen_comp2_j, eigen_comp3_j, 3)) Sigma22_j <- as.numeric(inverseEigen(eigen_comp1_j, eigen_comp2_j, eigen_comp3_j, 2)) Peigen_comp1_j <- PX_Sigma %*% samp_j[paste("Sigma_coef1[",1:ncol(X_Sigma),"]",sep = "")] Peigen_comp2_j <- PX_Sigma %*% samp_j[paste("Sigma_coef2[",1:ncol(X_Sigma),"]",sep = "")] Peigen_comp3_j <- PX_Sigma %*% samp_j[paste("Sigma_coef3[",1:ncol(X_Sigma),"]",sep = "")] PSigma11_j <- as.numeric(inverseEigen(Peigen_comp1_j, Peigen_comp2_j, Peigen_comp3_j, 1)) PSigma12_j <- as.numeric(inverseEigen(Peigen_comp1_j, Peigen_comp2_j, Peigen_comp3_j, 3)) PSigma22_j <- as.numeric(inverseEigen(Peigen_comp1_j, Peigen_comp2_j, Peigen_comp3_j, 2)) } if( modelsList$Sigma == "compRegIso" ){
# Required constants: X_Sigma, PX_Sigma
X_Sigma <- constants$X_Sigma PX_Sigma <- constants$PX_Sigma
eigen_comp1_j <- X_Sigma %*% samp_j[paste("Sigma_coef1[",1:ncol(X_Sigma),"]",sep = "")]
Sigma11_j <- as.numeric(exp(eigen_comp1_j))
Sigma12_j <- as.numeric(exp(eigen_comp1_j))
Sigma22_j <- rep(0,N)

Peigen_comp1_j <- PX_Sigma %*% samp_j[paste("Sigma_coef1[",1:ncol(X_Sigma),"]",sep = "")]
PSigma11_j <- as.numeric(exp(Peigen_comp1_j))
PSigma12_j <- as.numeric(exp(Peigen_comp1_j))
PSigma22_j <- rep(0,N)
}
if( modelsList$Sigma == "npApproxGP" ){ # Required constants: p_Sigma, Sigma_cross_dist, Sigma_cross_dist_pred, Sigma_HP2 p_Sigma <- constants$p_Sigma
Sigma_cross_dist_obs <- constants$Sigma_cross_dist Sigma_cross_dist_pred <- constants$Sigma_cross_dist_pred
Sigma_HP2 <- constants$Sigma_HP2 w1_Sigma_j <- samp_j[paste("w1_Sigma[",1:p_Sigma,"]",sep = "")] w2_Sigma_j <- samp_j[paste("w2_Sigma[",1:p_Sigma,"]",sep = "")] w3_Sigma_j <- samp_j[paste("w3_Sigma[",1:p_Sigma,"]",sep = "")] # Obs locations Pmat12_Sigma_obs_j <- matern_corr(Sigma_cross_dist_obs, samp_j["SigmaGP_phi[1]"], Sigma_HP2[1]) Pmat3_Sigma_obs_j <- matern_corr(Sigma_cross_dist_obs, samp_j["SigmaGP_phi[2]"], Sigma_HP2[2]) eigen_comp1_j <- samp_j["SigmaGP_mu[1]"]*rep(1,N) + samp_j["SigmaGP_sigma[1]"] * Pmat12_Sigma_obs_j %*% w1_Sigma_j eigen_comp2_j <- samp_j["SigmaGP_mu[1]"]*rep(1,N) + samp_j["SigmaGP_sigma[1]"] * Pmat12_Sigma_obs_j %*% w2_Sigma_j eigen_comp3_j <- samp_j["SigmaGP_mu[2]"]*rep(1,N) + samp_j["SigmaGP_sigma[2]"] * Pmat3_Sigma_obs_j %*% w3_Sigma_j Sigma11_j <- as.numeric(inverseEigen(eigen_comp1_j, eigen_comp2_j, eigen_comp3_j, 1)) Sigma12_j <- as.numeric(inverseEigen(eigen_comp1_j, eigen_comp2_j, eigen_comp3_j, 3)) Sigma22_j <- as.numeric(inverseEigen(eigen_comp1_j, eigen_comp2_j, eigen_comp3_j, 2)) # Pred locations Pmat12_Sigma_pred_j <- matern_corr(Sigma_cross_dist_pred, samp_j["SigmaGP_phi[1]"], Sigma_HP2[1]) Pmat3_Sigma_pred_j <- matern_corr(Sigma_cross_dist_pred, samp_j["SigmaGP_phi[2]"], Sigma_HP2[2]) Peigen_comp1_j <- samp_j["SigmaGP_mu[1]"]*rep(1,M) + samp_j["SigmaGP_sigma[1]"] * Pmat12_Sigma_pred_j %*% w1_Sigma_j Peigen_comp2_j <- samp_j["SigmaGP_mu[1]"]*rep(1,M) + samp_j["SigmaGP_sigma[1]"] * Pmat12_Sigma_pred_j %*% w2_Sigma_j Peigen_comp3_j <- samp_j["SigmaGP_mu[2]"]*rep(1,M) + samp_j["SigmaGP_sigma[2]"] * Pmat3_Sigma_pred_j %*% w3_Sigma_j PSigma11_j <- as.numeric(inverseEigen(Peigen_comp1_j, Peigen_comp2_j, Peigen_comp3_j, 1)) PSigma12_j <- as.numeric(inverseEigen(Peigen_comp1_j, Peigen_comp2_j, Peigen_comp3_j, 3)) PSigma22_j <- as.numeric(inverseEigen(Peigen_comp1_j, Peigen_comp2_j, Peigen_comp3_j, 2)) } if( modelsList$Sigma == "npApproxGPIso" ){
# Required constants: p_Sigma, Sigma_cross_dist, Sigma_cross_dist_pred, Sigma_HP2
p_Sigma <- constants$p_Sigma Sigma_cross_dist_obs <- constants$Sigma_cross_dist
Sigma_cross_dist_pred <- constants$Sigma_cross_dist_pred Sigma_HP2 <- constants$Sigma_HP2
w1_Sigma_j <- samp_j[paste("w1_Sigma[",1:p_Sigma,"]",sep = "")]

# Obs locations
Pmat12_Sigma_obs_j <- matern_corr(Sigma_cross_dist_obs, samp_j["SigmaGP_phi[1]"], Sigma_HP2[1])
eigen_comp1_j <- samp_j["SigmaGP_mu[1]"]*rep(1,N) + samp_j["SigmaGP_sigma[1]"] * Pmat12_Sigma_obs_j %*% w1_Sigma_j
Sigma11_j <- as.numeric(exp(eigen_comp1_j))
Sigma12_j <- as.numeric(exp(eigen_comp1_j))
Sigma22_j <- rep(0,N)

# Pred locations
Pmat12_Sigma_pred_j <- matern_corr(Sigma_cross_dist_pred, samp_j["SigmaGP_phi[1]"], Sigma_HP2[1])
Peigen_comp1_j <- samp_j["SigmaGP_mu[1]"]*rep(1,M) + samp_j["SigmaGP_sigma[1]"] * Pmat12_Sigma_pred_j %*% w1_Sigma_j
PSigma11_j <- as.numeric(exp(Peigen_comp1_j))
PSigma12_j <- as.numeric(exp(Peigen_comp1_j))
PSigma22_j <- rep(0,N)
}

# Calculate mu and Pmu ========================================
if( modelsList$mu == "constant" ){ mu <- samp_j[paste("beta")]*rep(1,N) Pmu <- samp_j[paste("beta")]*rep(1,M) } if( modelsList$mu == "linReg" ){
X_mu <- constants$X_mu PX_mu <- constants$PX_mu
mu <- as.numeric(X_mu %*% samp_j[paste("beta[",1:ncol(X_mu),"]",sep = "")])
Pmu <- as.numeric(PX_mu %*% samp_j[paste("beta[",1:ncol(X_mu),"]",sep = "")])
}
if( modelsList$mu == "zero" ){ mu <- 0*rep(1,N) Pmu <- 0*rep(1,M) } if( modelsList$likelihood == "fullGP" ){ # Predictions for the full GP likelihood
# Posterior predictive draw ===================================
# Obs covariance
Cor <- nsCorr(dist1_sq, dist2_sq, dist12, Sigma11_j, Sigma22_j, Sigma12_j, nu, d)
sigmaMat <- diag(exp(log_sigma_vec_j))
Cov <- sigmaMat %*% Cor %*% sigmaMat
C <- Cov + diag(exp(log_tau_vec_j)^2)
C_chol <- chol(C)
# Prediction covariance
PCor <- nsCorr(Pdist1_sq, Pdist2_sq, Pdist12, PSigma11_j, PSigma22_j, PSigma12_j, nu, d)
PsigmaMat <- diag(exp(Plog_sigma_vec_j))
PCov <- PsigmaMat %*% PCor %*% PsigmaMat
if(predict.process){ # Do not include the nugget variance
PC <- PCov
} else{ # Include nugget variance
PC <- PCov + diag(exp(Plog_tau_vec_j)^2)
}
# Cross-covariance
XCor <- nsCrosscorr(Xdist1_sq, Xdist2_sq, Xdist12,
Sigma11_j, Sigma22_j, Sigma12_j,
PSigma11_j, PSigma22_j, PSigma12_j, nu, d)
XCov <- PsigmaMat %*% XCor %*% sigmaMat
# Conditional mean/covariance
crscov_covinv <- t(backsolve(C_chol, backsolve(C_chol, t(XCov), transpose = TRUE)))
condMean <- Pmu + crscov_covinv %*% (z - mu)
condCov <- PC - crscov_covinv %*% t(XCov)
condCov_chol <- chol(condCov)
# Store
postPredDraws[j,] <- condMean + t(condCov_chol) %*% rnorm(M)
} else if( modelsList$likelihood == "NNGP" ){ # Predictions for the NNGP likelihood # Local kriging =============================================== for(m in 1:M){ # Obs covariance -- for nearest neighbors Cor <- nsCorr(Pdist1_3d[m,1:k,1:k], Pdist2_3d[m,1:k,1:k], Pdist12_3d[m,1:k,1:k], Sigma11_j[P_nID[m,]], Sigma22_j[P_nID[m,]], Sigma12_j[P_nID[m,]], nu, d) sigmaMat <- diag(exp(log_sigma_vec_j[P_nID[m,]])) Cov <- sigmaMat %*% Cor %*% sigmaMat C <- Cov + diag(exp(log_tau_vec_j[P_nID[m,]])^2) C_chol <- chol(C) # Prediction variance # PCor <- nsCorr(Pdist1_sq, Pdist2_sq, Pdist12, PSigma11_j, PSigma22_j, PSigma12_j, nu) PsigmaMat <- exp(Plog_sigma_vec_j[m]) # PCov <- PsigmaMat %*% Cor %*% PsigmaMat if(predict.process){ # Do not include the nugget variance PC <- exp(Plog_sigma_vec_j[m])^2 } else{ # Include nugget variance PC <- exp(Plog_tau_vec_j[m])^2 + exp(Plog_sigma_vec_j[m])^2 } # Cross-covariance XCor <- nsCrosscorr(matrix(Pdist1_3d[m,k+1,1:k], nrow = 1, ncol = k), matrix(Pdist2_3d[m,k+1,1:k], nrow = 1, ncol = k), matrix(Pdist12_3d[m,k+1,1:k], nrow = 1, ncol = k), Sigma11_j[P_nID[m,]], Sigma22_j[P_nID[m,]], Sigma12_j[P_nID[m,]], PSigma11_j[m], PSigma22_j[m], PSigma12_j[m], nu, d) XCov <- PsigmaMat %*% XCor %*% sigmaMat # Conditional mean/covariance crscov_covinv <- t(backsolve(C_chol, backsolve(C_chol, t(XCov), transpose = TRUE))) condMean <- Pmu[m] + crscov_covinv %*% (z[P_nID[m,]] - mu[P_nID[m,]]) condCov <- PC - crscov_covinv %*% t(XCov) condCov_chol <- chol(condCov) pred_sample <- condMean + t(condCov_chol) %*% rnorm(1) # Store postPredDraws[j,m] <- pred_sample } } else if( modelsList$likelihood == "SGV" ){ # Predictions for the SGV likelihood
# SGV prediction ==============================================
U_j <- calculateU_ns(
dist1_3d = Alldist1_3d, dist2_3d = Alldist2_3d, dist12_3d = Alldist12_3d,
Sigma11 = c(Sigma11_j, PSigma11_j),
Sigma22 = c(Sigma22_j, PSigma22_j),
Sigma12 = c(Sigma12_j, PSigma12_j),
log_sigma_vec = c(log_sigma_vec_j, Plog_sigma_vec_j),
log_tau_vec = c(log_tau_vec_j, Plog_tau_vec_j), nu = nu,
nID = prednID_SGV, cond_on_y = predSGV_setup$condition_on_y_ord, N = N, k = k, M = M, d = d ) Usm <- sparseMatrix(i = U_j[,1], j = U_j[,2], x = U_j[,3]) Asm <- Usm[c(seq(from = 1, to = 2*N, by = 2), 2*N + 1:M),] Bsm <- Usm[seq(from = 2, to = 2*N, by = 2),] # Calculate V = rchol(W) Aoo_sm <- Asm[1:N,1:(2*N)] Woo_sm <- tcrossprod(Aoo_sm)[N:1,N:1] Voo_sm <- t(chol(Woo_sm))[N:1,N:1] Vsm <- Asm[,-(1:N)] Vsm[1:N,1:N] <- Voo_sm # Kriging predictor (mean zero) ABtz <- crossprod(t(Asm), crossprod(Bsm, z - mu)) krigPredictor <- -as.numeric(solve(tcrossprod(Vsm), ABtz)) # Draw mean zero pred_sample <- solve(t(Vsm), rnorm(N+M), system = "L") # Combine (with mean) # postpred_draw <- Pmu + pred_sample[-(1:N)] + krigPredictor[-(1:N)] postpred_draw <- c(mu,Pmu) + as.numeric(pred_sample + krigPredictor) # Store postPredDraws[j,] <- postpred_draw PPD_obs_ord <- postPredDraws[,1:N] PPD_pred_ord <- postPredDraws[,-(1:N)] PPD_obs_orig <- PPD_obs_ord[,order(model_constants$ord)]
PPD_pred_orig <- PPD_pred_ord[,order(predSGV_setup$ord_pred)] } else stop('') # Progress if( j %% ceiling(J/56) == 0 ){nimCat("-")} } nimCat("|\n") if( modelsList$likelihood == "fullGP" ){ # Predictions for the full GP likelihood
output <- list(obs = NULL, pred = postPredDraws, pred.mmd.seed = NULL)
} else if( modelsList$likelihood == "NNGP" ){ # Predictions for the NNGP likelihood output <- list(obs = NULL, pred = postPredDraws, pred.mmd.seed = NULL) } else if( modelsList$likelihood == "SGV" ){ # Predictions for the SGV likelihood
output <- list(obs = PPD_obs_orig, pred = PPD_pred_orig, pred.mmd.seed = pred.mmd.seed)
} else stop('')

return(output)

}


## Try the BayesNSGP package in your browser

Any scripts or data that you put into this service are public.

BayesNSGP documentation built on Jan. 9, 2022, 9:07 a.m.