R/vecchia_prediction.R

In GPvecchia: Scalable Gaussian-Process Approximations

#' Vecchia prediction
#'
#' @param z observed data
#' @param vecchia.approx a vecchia object as generated by vecchia_specify()
#' @param covparms covariance parameters as a vector
#' @param nuggets nugget
#' @param var.exact should prediction variances be computed exactly, or is a (faster) approximation acceptable
#' @param covmodel covariance model, 'matern' by default.
#' @param return.values either 'mean' only, 'meanvar', 'meanmat', or 'all'
#'
#' @return posterior mean and variances at observed and unobserved locations; V matrix
#' @examples
#' z=rnorm(5); locs=matrix(1:5,ncol=1); vecchia.approx=vecchia_specify(locs,m=3,locs.pred=locs+.5)
#' vecchia_prediction(z,vecchia.approx,covparms=c(1,2,.5),nuggets=.2)
#' @export

vecchia_prediction=function(z,vecchia.approx,covparms,nuggets,var.exact,
                            covmodel='matern',return.values='all') {

  # remove NAs in data and U
  removeNAs()

  # create the U matrix
  U.obj=createU(vecchia.approx,covparms,nuggets,covmodel)

  # compute cholesky V for posterior inference
  V.ord=U2V(U.obj)

  if(length(U.obj$zero.nugg)>0)
    warning('Rows/cols of V have been removed for data with zero noise')

  # compute the posterior mean
  vecchia.mean=vecchia_mean(z,U.obj,V.ord)

  # return what is requested
  return.list=list(mu.pred=vecchia.mean$mu.pred,mu.obs=vecchia.mean$mu.obs,
                   var.pred=NULL,var.obs=NULL,V.ord=NULL,U.obj=NULL)

  if(return.values=='meanmat' | return.values=='all'){
    return.list$V.ord=V.ord; return.list$U.obj=U.obj
  }

  if(return.values=='meanvar' | return.values=='all'){

    # compute posterior variances
    if(missing(var.exact)) var.exact = (sum(!vecchia.approx$obs)<4*1e4)
    vars.vecchia=vecchia_var(U.obj,V.ord,exact=var.exact)

    return.list$var.pred=vars.vecchia$vars.pred
    return.list$var.obs=vars.vecchia$vars.obs

  }

  return(return.list)

}



######  compute V for posterior inference   #######

U2V=function(U.obj){
### when changing this function make sure it returns a dtCMatrix!
### Otherwise solve in the parent function will be very slow
  
  U.y=U.obj$U[U.obj$latent,]

  if(U.obj$cond.yz=='zy') {

    V.ord=revMat(U.y[,U.obj$latent,drop=FALSE])

  } else if(U.obj$ord.pred!='obspred'){

    W=Matrix::tcrossprod(U.y)
    W.rev=revMat(W)

    if(U.obj$ic0){
        V.ord=Matrix::t(ichol(W.rev))
    } else {
        V.ord=Matrix::t(Matrix::chol(W.rev))
    }

    V.ord = methods::as(V.ord, 'dtCMatrix')
      
    
  } else {  # for obspred ordering

    last.obs=max(which(!U.obj$latent))
    latents.before=sum(U.obj$latent[1:last.obs])
    latents.after=sum(U.obj$latent[-(1:last.obs)])

    # pred columns are unchanged
    V.pr=revMat(U.y[,(last.obs+1):ncol(U.y),drop=FALSE])

    # have to compute cholesky for obs block
    U.oo=U.y[1:latents.before,1:last.obs]
    A=Matrix::tcrossprod(U.oo)
    A.rev=revMat(A)
    if(U.obj$ic0){ V.oor=Matrix::t(ichol(A.rev))
      }     else   V.oor=Matrix::t(Matrix::chol(A.rev))
    
    # combine the blocks into one matrix
    zeromat.sparse=Matrix::sparseMatrix(c(),c(),dims=c(latents.after,latents.before))
    V.or=rbind(zeromat.sparse,V.oor)

    V.ord=methods::as(cbind(V.pr,V.or),'dtCMatrix')

  }

  return(V.ord)
}




######  posterior mean (predictions)   #######

vecchia_mean=function(z,U.obj,V.ord){
  U=U.obj$U
  # compute entire posterior mean vector
  z.ord=z[U.obj$ord.z]
  z1=Matrix::crossprod(U[!U.obj$latent,],z.ord)
  z2=as.numeric(U[U.obj$latent,]%*%z1)
  temp=Matrix::solve(V.ord,rev(z2))
  mu.rev=-Matrix::solve(Matrix::t(V.ord),temp)
  mu.ord=rev(mu.rev)

  # for zero nugget, observations are posterior means
  if(length(U.obj$zero.nugg)>0){
    obs.zero=z.ord[U.obj$zero.nugg$inds.z]
    mu.ord=c(mu.ord,obs.zero)
  }

  # extract obs and pred parts; return to original ordering
  orig.order=order(U.obj$ord)
  mu=mu.ord[orig.order]
  obs.orig=U.obj$obs[orig.order]
  mu.obs=mu[obs.orig]
  mu.pred=mu[!obs.orig]

  return(list(mu.obs=mu.obs,mu.pred=mu.pred))
}



######  linear combination   #######

#' linear combination of predictions
#' compute the distribution of a linear combination Hy
#' @param H sparse matrix with n.all columns specifying the linear combination
#' @param U.obj U matrix is the full joint approximated cholesky matrix
#' @param V.ord ordered V matrix from vecchia_prediction() or U2V()
#' @param cov.mat logical TRUE or FALSE -- should the entire covariance matrix be returned (only do if H has a small number of rows)
#'
#' @return Variance of linear combination of predictions.
#' @examples
#' n=5; z=rnorm(n); locs=matrix(1:n,ncol=1); n.p=5
#' vecchia.approx = vecchia_specify(locs,m=3,locs.pred=locs+.5)
#' preds=vecchia_prediction(z,vecchia.approx,covparms=c(1,2,.5),nuggets=.2)
#' H=Matrix::sparseMatrix(i=rep(1,n.p),j=n+(1:n.p),x=1/n.p)
#' vecchia_lincomb(H,vecchia.approx,preds$V.ord,cov.mat=TRUE)
#' @export

vecchia_lincomb=function(H,U.obj,V.ord,cov.mat=FALSE) {
  ord=U.obj$ord
  if(length(U.obj$zero.nugg)>0){
    ord=order(order(ord[1:(length(ord)-length(U.obj$zero.nugg$inds.U))]))
  }
  H.tt=Matrix::t(H[,rev(ord),drop=FALSE])
  temp=Matrix::solve(V.ord,H.tt)
  if(cov.mat){
    lincomb.cov=as.matrix(Matrix::t(temp)%*%temp)
    return(lincomb.cov)
  } else {
    lincomb.vars=as.numeric(Matrix::t(temp*temp)%*%rep(1,ncol(H)))
    return(lincomb.vars)
  }
}



######  selected inverse of a sparse matrix   #######

#' selected inverse of a sparse matrix
#'
#' @param cholmat cholesky factor L of a positive definite sparseMatrix A
#'
#' @return sparse inverse of A, with same sparsity pattern as L
#' @examples
#' A=Matrix::sparseMatrix(1:9,1:9,x=4); L=Matrix::chol(A)
#' SelInv(L)
#' @export
SelInv=function(cholmat){
  n.all=nrow(cholmat)
  sparseinv::Takahashi_Davis(Q=Matrix::sparseMatrix(c(),c(),dims=c(n.all,1)),
                  cholQp=cholmat,P=Matrix::sparseMatrix(i=1:n.all,j=1:n.all,x=1))
}



######  posterior variances   #######

vecchia_var=function(U.obj,V.ord,exact=FALSE){

  # compute selected inverse and extract variances
  inv.sparse=SelInv(V.ord)
  vars.ord=rev(Matrix::diag(inv.sparse))

  # for zero nugget, add zero variances
  if(length(U.obj$zero.nugg)>0){
    vars.ord=c(vars.ord,rep(0,length(U.obj$zero.nugg$inds.z)))
  }

  # return to original ordering
  orig.order=order(U.obj$ord)
  vars=vars.ord[orig.order]

  # extract obs and pred variances
  obs.orig=U.obj$obs[orig.order]
  vars.obs=vars[obs.orig]
  vars.pred=vars[!obs.orig]

  # if exact variances are desired:
  if(exact & U.obj$ord.pred=='obspred'){

    n=sum(U.obj$obs)
    n.p=sum(!U.obj$obs)
    H=Matrix::sparseMatrix(i=1:(n+n.p),j=1:(n+n.p),x=1)

    if(U.obj$cond.yz!='zy'){

      if(sum(!obs.orig)>0) {
        vars.pred=vecchia_lincomb(H[(n+1):(n+n.p),],U.obj,V.ord)
      } else vars.pred=c()

    } else {

      vars.exact=vecchia_lincomb(H,U.obj,V.ord)
      vars.obs=vars.exact[1:n]
      vars.pred= (if(n.p>0) vars.exact[n+(1:n.p)] else c())

    }

  }

  return(list(vars.obs=vars.obs,vars.pred=vars.pred))
}



######  compute covariance matrix from V.ord   #######
# do not do this for large n or n.p!!!

#' compute covariance matrix from V.ord
#' Do not run this function for large n or n.p!!!
#' @param preds Object returned by vecchia_prediction()
#'
#' @return Covariance matrix at all locations in original order
#' @examples
#' z=rnorm(5)
#' locs=matrix(1:5,ncol=1)
#' vecchia_specify=function(z,locs,m=5,locs.pred=(1:5)+.5)
#' V2covmat(vecchia_prediction(vecchia.approx,covparms=c(1,2,.5),nuggets=.2))
#' @export
V2covmat=function(preds){

  # compute joint covariance matrix
  Sigma.ord=Matrix::solve(as.matrix(revMat(preds$V.ord%*%Matrix::t(preds$V.ord))))

  # for zero nugget, add zero rows/columns
  if(length(preds$U.obj$zero.nugg)>0){
    k=nrow(Sigma.ord)
    l=length(preds$U.obj$zero.nugg$inds.z)
    Sigma.ord=rbind(cbind(Sigma.ord,matrix(0,nrow=k,ncol=l)),
                    matrix(0,nrow=l,ncol=k+l))
  }

  # return to original ordering
  orig.order=order(preds$U.obj$ord)
  Sigma=Sigma.ord[orig.order,orig.order]

  # extract parts corresponding to obs and pred locs
  obs.orig=preds$U.obj$obs[orig.order]
  Sigma.obs=Sigma[obs.orig,obs.orig]
  Sigma.pred=Sigma[!obs.orig,!obs.orig]

  return(list(Sigma.obs=Sigma.obs,Sigma.pred=Sigma.pred))
}