R/pmvt_MLMC.R

In VeccTMVN: Multivariate Normal Probabilities using Vecchia Approximation

#' Applying the multi-level Monte Carlo (MLMC) technique to the pmvt function
#' The function uses \code{NLevel1 = 1} for \code{m = m2} and the same
#' exponential tilting parameter as \code{m = m1} to compute one MC estimate.
#' This MC estimate is used to correct the bias from the Vecchia approximation
#'
#' @importFrom truncnorm etruncnorm
#' @importFrom nleqslv nleqslv
#' @importFrom stats runif
#' @importFrom TruncatedNormal cholperm
#' @importFrom utils getFromNamespace
#'
#' @param lower lower bound vector for TMVT
#' @param upper upper bound vector for TMVT
#' @param delta MVT shifting parameter
#' @param df degrees of freedom
#' @param locs location (feature) matrix n X d
#' @param covName covariance function name from the `GpGp` package
#' @param covParms parameters for `covName`
#' @param m1 the smaller Vecchia conditioning set size for Level 1 MC
#' @param m2 the bigger Vecchia conditioning set size for Level 2 MC
#' @param sigma dense covariance matrix, not needed when `locs` is not null
#' @param reorder whether to reorder integration variables. `0` for no,
#' `1` for FIC-based univariate ordering, `2` for Vecchia-based univariate
#' ordering, and `3` for univariate reordering, which appeared faster than `2`
#' @param NLevel1 first level Monte Carlo sample size
#' @param NLevel2 second level Monte Carlo sample size
#' @param verbose verbose or not
#' @param retlog TRUE or FALSE for whether to return loglk or not
#' @param ... could be
#' m_ord for conditioning set size for reordering
#' @return estimated MVT probability and estimation error
#'
#' @export
pmvt_MLMC <- function(lower, upper, delta, df, locs = NULL, covName = "matern15_isotropic",
                      covParms = c(1.0, 0.1, 0.0), m1 = 30, m2 = 100, sigma = NULL, reorder = 0,
                      NLevel1 = 12, NLevel2 = 1e4, verbose = FALSE, retlog = FALSE, ...) {
  # standardize the input MVN prob -----------------------------
  lower <- lower - delta
  upper <- upper - delta
  if (is.null(sigma)) {
    n <- nrow(locs)
    use_sigma <- FALSE
    margin_sd <- sqrt(covParms[1])
    upper <- upper / margin_sd
    lower <- lower / margin_sd
    covParms[1] <- 1
  } else {
    n <- nrow(sigma)
    use_sigma <- TRUE
    margin_sd <- sqrt(diag(sigma))
    upper <- upper / margin_sd
    lower <- lower / margin_sd
    sigma <- t(t(sigma / margin_sd) / margin_sd)
  }
  if (any(upper < lower)) {
    stop("Invalid MVN probability. Truncated marginal
         probabilities have negative value(s)\n")
  }
  if (m1 > m2) {
    stop("m1 should be greater than m2 when using
         Multi-level Monte Carlo(MLMC)\n")
  }
  lower_upper <- matrix(0, n, 2)
  lower_upper[, 1] <- lower
  lower_upper[, 2] <- upper
  lower <- lower_upper[, 1]
  upper <- lower_upper[, 2]
  # reorder --------------------------------
  if (is.null(list(...)[["m_ord"]])) {
    m_ord <- m1
  } else {
    m_ord <- list(...)[["m_ord"]]
  }
  if (reorder == 1) {
    if (use_sigma) {
      ord <- FIC_reorder_univar(lower, upper, m_ord, covMat = sigma)
    } else {
      ord <- FIC_reorder_univar(
        lower, upper, m_ord, locs, covName,
        covParms
      )
    }
    lower <- lower[ord]
    upper <- upper[ord]
    if (use_sigma) {
      sigma <- sigma[ord, ord, drop = FALSE]
    } else {
      locs <- locs[ord, , drop = FALSE]
    }
  } else if (reorder == 2) {
    if (use_sigma) {
      ord <- Vecc_reorder(lower, upper, m_ord, covMat = sigma)$order
    } else {
      ord <- Vecc_reorder(
        lower, upper, m_ord, locs, covName, covParms
      )$order
    }
    lower <- lower[ord]
    upper <- upper[ord]
    if (use_sigma) {
      sigma <- sigma[ord, ord, drop = FALSE]
    } else {
      locs <- locs[ord, , drop = FALSE]
    }
  } else if (reorder == 3) {
    if (!use_sigma) {
      cov_func_GpGp <- utils::getFromNamespace(covName, "GpGp")
      sigma <- cov_func_GpGp(covParms, locs)
    }
    ord <- univar_order(lower, upper, sigma)
    lower <- lower[ord]
    upper <- upper[ord]
    if (use_sigma) {
      sigma <- sigma[ord, ord, drop = FALSE]
    } else {
      locs <- locs[ord, , drop = FALSE]
    }
  }
  # find nearest neighbors for Vecchia --------------------------------
  if (use_sigma) {
    NN_m2 <- find_nn_corr(sigma, m2)
    NN_m1 <- NN_m2[, 1:(m1 + 1)]
  } else {
    NN_m2 <- GpGp::find_ordered_nn(locs, m2)
    NN_m1 <- NN_m2[, 1:(m1 + 1)]
  }
  # find Vecchia approx object -----------------------------------
  if (use_sigma) {
    vecc_obj_m1 <- vecc_cond_mean_var_sp(NN_m1, covMat = sigma)
    vecc_obj_m2 <- vecc_cond_mean_var_sp(NN_m2, covMat = sigma)
  } else {
    vecc_obj_m1 <- vecc_cond_mean_var_sp(NN_m1,
      locs = locs, covName = covName,
      covParms = covParms
    )
    vecc_obj_m2 <- vecc_cond_mean_var_sp(NN_m2,
      locs = locs, covName = covName,
      covParms = covParms
    )
  }
  # find tilting parameter beta -----------------------------------
  trunc_expect <- truncnorm::etruncnorm(lower, upper)
  x0 <- c(trunc_expect, rep(0, n))
  x0[2 * n] <- sqrt(df)
  x0[n] <- log(x0[2 * n])
  solv_idea_5_sp <- nleqslv::nleqslv(
    x = x0, fn = gradpsiT_idea5, veccCondMeanVarObj = vecc_obj_m1,
    a = lower, b = upper, nu = df, global = "pwldog", method = "Broyden"
  )
  soln <- solv_idea_5_sp$x
  exitflag <- solv_idea_5_sp$termcd
  if (!(exitflag %in% c(1, 2)) || !all.equal(solv_idea_5_sp$fvec, rep(0, length(x0)))) {
    warning("Did not find a solution to the nonlinear system in `pmvt`!")
  }
  soln[n] <- exp(soln[n])
  x <- soln[1:n]
  beta <- soln[(n + 1):(2 * n)]
  eta <- beta[n]
  # compute MVT probs and est error ---------------------------------
  const <- log(2 * pi) / 2 - lgamma(df / 2) - (df / 2 - 1) * log(2) +
    TruncatedNormal::lnNpr(-eta, Inf) + 0.5 * eta^2
  seed <- round(runif(1) * 1e6)
  set.seed(seed)
  exp_psi_m1 <- sample_psiT_idea5_cpp(vecc_obj_m1, lower, upper, df,
    beta = beta, N_level1 = NLevel1,
    N_level2 = NLevel2
  )
  set.seed(seed)
  exp_psi_m2 <- sample_psiT_idea5_cpp(vecc_obj_m2, lower, upper, df,
    beta = beta, N_level1 = 1,
    N_level2 = NLevel2
  )
  if (retlog) {
    exponent <- min(exp_psi[[2]])
    # correction using exp_psi_m2
    log_est_prob <- exponent +
      log(mean(exp_psi_m1[[1]] * exp(exp_psi_m1[[2]] - exponent)) +
        exp_psi_m2[[1]][1] * exp(exp_psi_m2[[2]][1] - exponent) -
        exp_psi_m1[[1]][1] * exp(exp_psi_m1[[2]][1] - exponent)) + const
    return(log_est_prob)
  } else {
    exp_psi <- exp_psi_m1[[1]] * exp(exp_psi_m1[[2]])
    exp_psi_m2 <- exp_psi_m2[[1]] * exp(exp_psi_m2[[2]])
    # correction using exp_psi_m2
    est_prob <- (mean(exp_psi) + (exp_psi_m2[1] - exp_psi[1])) * exp(const)
    est_prob_err <- stats::sd(exp_psi) / sqrt(NLevel1) * exp(const)
    attr(est_prob, "error") <- est_prob_err
    return(est_prob)
  }
}


# TEST -------------------------------------------------------

# library(VeccTMVN)
# library(TruncatedNormal)
# library(GpGp)
#
# ## example MVN probabilities --------------------------------
# set.seed(123)
# n1 <- 10
# n2 <- 10
# n <- n1 * n2
# locs <- as.matrix(expand.grid((1:n1) / n1, (1:n2) / n2))
# covparms <- c(2, 0.3, 0)
# nu <- 10
# cov_mat <- matern15_isotropic(covparms, locs)
# a_list <- list(rep(-Inf, n), rep(-1, n), -runif(n) * 2 - 4)
# b_list <- list(rep(-2, n), rep(1, n), -runif(n) * 2)
#
# ## Compute MVN probs --------------------------------
# N_level1 <- 12 # Level 1 MC size
# N_level2 <- 1e4 # Level 2 MC size
# m <- 30 # num of nearest neighbors
# for (i in 1:length(a_list)) {
#   a <- a_list[[i]]
#   b <- b_list[[i]]
#   ### Compute MVN prob with idea V -----------------------
#   est_Vecc <- VeccTMVN::pmvt(a, b, 0, nu, locs, covName = "matern15_isotropic",
#                              covParms = covparms, m = m, verbose = FALSE)
#   est_Vecc_MLMC <- VeccTMVN::pmvt_MLMC(a, b, 0, nu, locs, covName = "matern15_isotropic",
#                                        covParms = covparms, m1 = m, m2 = 2 * m, verbose = FALSE)
#   est_TN <- TruncatedNormal::pmvt(rep(0, n), cov_mat, nu, lb = a, ub = b)
#   cat(
#     "est_Vecc", est_Vecc, "err_Vecc", attributes(est_Vecc)$error, "\n"
#   )
#   cat(
#     "est_Vecc_MLMC", est_Vecc_MLMC, "err_Vecc", attributes(est_Vecc_MLMC)$error,
#     "\n"
#   )
#   cat(
#     "est_TN", est_TN, "err_TN", attributes(est_TN)$relerr * est_TN, "\n"
#   )
# }