scripts/deep_mra_block_mh.R

In jtipton25/sgMRA: Multi-Resolution Gaussian Process Approximations with Stochastic Gradient Descent

deep_mra_mh <- function(y, locs, n_mcmc = 1000,
                        alpha_x1 = NULL,
                        sample_alpha_x1 = TRUE,
                        alpha_y1 = NULL,
                        sample_alpha_y1 = TRUE,
                        alpha_x2 = NULL,
                        sample_alpha_x2 = TRUE,
                        alpha_y2 = NULL,
                        sample_alpha_y2 = TRUE,
                        alpha = NULL,
                        sample_alpha = TRUE,
                        M = 1, n_coarse_grid = 20,
                         n_message = 50) {

    library(spam)
    library(Matrix)
    library(igraph)
    library(tidyverse)
    library(BayesMRA)

    source(here::here("R", "update-tuning.R"))
    N <- length(y)

    grid <- make_grid(locs, M = M, n_coarse_grid = n_coarse_grid)

    # initialize the algorithm ----
    # construct the first layer
    # W1 <- mra_wendland_2d(as.matrix(locs), M=M, n_coarse_grid = n_coarse_grid)$W
    W1 <- eval_basis(as.matrix(locs), grid)$W
    Q1 <- make_Q_alpha_2d(sqrt(MRA1$n_dims), phi=0.9)
    class(Q1) <- "spam"
    if (is.null(alpha_x1)) {
        alpha_x1 <- drop(rmvnorm.prec(1, mu = rep(0, nrow(Q1)), Q = Q1))
    }
    if (is.null(alpha_y1)) {
        alpha_y1 <- drop(rmvnorm.prec(1, mu = rep(0, nrow(Q1)), Q = Q1))
    }

    # # construct the second layer
    # W2 <- mra_wendland_2d(cbind(W1 %*% alpha_x1, W1 %*% alpha_y1), M=M, n_coarse_grid = n_coarse_grid)$W
    W2 <- eval_basis(cbind(W1 %*% alpha_x1, W1 %*% alpha_y1), grid)$W
    Q2 <- make_Q_alpha_2d(sqrt(MRA2$n_dims), phi=0.9)
    class(Q2) <- "spam"
    if (is.null(alpha_x2)) {
        alpha_x2 <- drop(rmvnorm.prec(1, mu = rep(0, nrow(Q2)), Q = Q2))
    }
    if (is.null(alpha_y2)) {
        alpha_y2 <- drop(rmvnorm.prec(1, mu = rep(0, nrow(Q2)), Q = Q2))
    }

    # construct the final layer
    # W <- mra_wendland_2d(cbind(W2 %*% alpha_x2, W2 %*% alpha_y2), M=M, n_coarse_grid = n_coarse_grid)$W
    W <- eval_basis(cbind(W2 %*% alpha_x2, W2 %*% alpha_y2), grid)$W
    Q <- make_Q_alpha_2d(sqrt(MRA$n_dims), phi=0.9)
    class(Q) <- "spam"
    if (is.null(alpha)) {
        alpha <- drop(rmvnorm.prec(1, mu = rep(0, nrow(Q)), Q = Q))
    }

    sigma <- max(min(rgamma(1, 1, 1), 5), 0.05)

    alpha_x1_save <- matrix(0, n_mcmc, nrow(Q1))
    alpha_y1_save <- matrix(0, n_mcmc, nrow(Q1))
    # W1_save       <- vector('list', n_mcmc)
    alpha_x2_save <- matrix(0, n_mcmc, nrow(Q2))
    alpha_y2_save <- matrix(0, n_mcmc, nrow(Q2))
    # W2_save       <- vector('list', n_mcmc)
    alpha_save    <- matrix(0, n_mcmc, nrow(Q))
    # W_save        <- vector('list', n_mcmc)
    # Walpha1_save  <- matrix(0, n_mcmc, N)
    # Walpha2_save  <- matrix(0, n_mcmc, N)
    Walpha_save   <- matrix(0, n_mcmc, N)
    sigma_save    <- rep(0, n_mcmc)

    alpha_x1_tune <- 0.01
    alpha_x1_accept <- 0
    alpha_x1_accept_batch <- 0

    alpha_y1_tune <- 0.01
    alpha_y1_accept <- 0
    alpha_y1_accept_batch <- 0

    alpha_x2_tune <- 0.01
    alpha_x2_accept <- 0
    alpha_x2_accept_batch <- 0

    alpha_y2_tune <- 0.01
    alpha_y2_accept <- 0
    alpha_y2_accept_batch <- 0

    alpha_tune    <- 0.01
    alpha_accept <- 0
    alpha_accept_batch <- 0

    # setup the ess pars ----
    # pars <- list(
    #     y             = y,
    #     alpha_x1      = alpha_x1,
    #     alpha_y1      = alpha_y1,
    #     W1            = W1,
    #     alpha_x2      = alpha_x2,
    #     alpha_y2      = alpha_y2,
    #     W2            = W2,
    #     alpha         = alpha,
    #     W             = W,
    #     sigma         = sigma,
    #     M             = M,
    #     n_coarse_grid = n_coarse_grid,
    #     num_calls_ess = 0,
    #     verbose = TRUE)


    message("Starting MCMC, will run for ", n_mcmc, " iterations")
    for (k in 1:n_mcmc) {
        if (k %% n_message == 0) {
            message("On iteration ", k, " out of ", n_mcmc)
        }

        # update alpha_x1 ----
        if (sample_alpha_x1) {
            message("sampling alpha_x1")
            alpha_x1_prop <- rnorm(length(alpha_x1), alpha_x1, alpha_x1_tune)
            # W2_prop <- mra_wendland_2d(cbind(W1 %*% alpha_x1_prop, W1 %*% alpha_y1), M=M, n_coarse_grid = n_coarse_grid)$W
            W2_prop <- eval_basis(cbind(W1 %*% alpha_x1_prop, W1 %*% alpha_y1), grid)$W
            # W_prop  <- mra_wendland_2d(cbind(W2_prop %*% alpha_x2, W2_prop %*% alpha_y2), M=M, n_coarse_grid = n_coarse_grid)$W
            W_prop  <- eval_basis(cbind(W2_prop %*% alpha_x2, W2_prop %*% alpha_y2), grid)$W
            mh1 <- sum(dnorm(y, W_prop %*% alpha, sigma, log=TRUE)) -
                0.5 * sum(alpha_x1_prop * Q1 %*% alpha_x1_prop)
            mh2 <- sum(dnorm(y, W %*% alpha, sigma, log=TRUE)) -
                0.5 * sum(alpha_x1 * Q1 %*% alpha_x1)
            mh <- exp(mh1 - mh2)
            if (mh > runif(1)) {
                alpha_x1 <- alpha_x1_prop
                W2 <- W2_prop
                W  <- W_prop
                alpha_x1_accept_batch <- alpha_x1_accept_batch + 1 / 50
                alpha_x1_accept <- alpha_x1_accept + 1 / n_mcmc
            }

            if (k %% 50 == 0) {
                out_tune <- update_tuning(k, alpha_x1_accept_batch, alpha_x1_tune)
                alpha_x1_accept_batch <- out_tune$accept
                alpha_x1_tune <- out_tune$tune
            }
        }

        # update alpha_y1 ----
      if (sample_alpha_y1) {

          message("sampling alpha_y1")
          alpha_y1_prop <- rnorm(length(alpha_y1), alpha_y1, alpha_y1_tune)
          # W2_prop <- mra_wendland_2d(cbind(W1 %*% alpha_x1, W1 %*% alpha_y1_prop), M=M, n_coarse_grid = n_coarse_grid)$W
          W2_prop <- eval_basis(cbind(W1 %*% alpha_x1, W1 %*% alpha_y1_prop), grid)$W
          # W_prop  <- mra_wendland_2d(cbind(W2_prop %*% alpha_x2, W2_prop %*% alpha_y2), M=M, n_coarse_grid = n_coarse_grid)$W
          W_prop  <- eval_basis(cbind(W2_prop %*% alpha_x2, W2_prop %*% alpha_y2), grid)$W
          mh1 <- sum(dnorm(y, W_prop %*% alpha, sigma, log=TRUE)) -
              0.5 * sum(alpha_y1_prop * Q1 %*% alpha_y1_prop)
          mh2 <- sum(dnorm(y, W %*% alpha, sigma, log=TRUE)) -
              0.5 * sum(alpha_y1 * Q1 %*% alpha_y1)
          mh <- exp(mh1 - mh2)
          if (mh > runif(1)) {
              alpha_y1 <- alpha_y1_prop
              W2 <- W2_prop
              W  <- W_prop
              alpha_y1_accept_batch <- alpha_y1_accept_batch + 1 / 50
              alpha_y1_accept <- alpha_y1_accept + 1 / n_mcmc
          }

          if (k %% 50 == 0) {
              out_tune <- update_tuning(k, alpha_y1_accept_batch, alpha_y1_tune)
              alpha_y1_accept_batch <- out_tune$accept
              alpha_y1_tune <- out_tune$tune
          }
      }

        # update alpha_x2 ----
        if (sample_alpha_x2) {
            message("sampling alpha_x2")
            alpha_x2_prop <- rnorm(length(alpha_x2), alpha_x2, alpha_x2_tune)
            # W_prop  <- mra_wendland_2d(cbind(W2 %*% alpha_x2_prop, W2 %*% alpha_y2), M=M, n_coarse_grid = n_coarse_grid)$W
            W_prop  <- eval_basis(cbind(W2 %*% alpha_x2_prop, W2 %*% alpha_y2), grid)$W
            mh1 <- sum(dnorm(y, W_prop %*% alpha, sigma, log=TRUE)) -
                0.5 * sum(alpha_x2_prop * Q2 %*% alpha_x2_prop)
            mh2 <- sum(dnorm(y, W %*% alpha, sigma, log=TRUE)) -
                0.5 * sum(alpha_x2 * Q2 %*% alpha_x2)
            mh <- exp(mh1 - mh2)
            if (mh > runif(1)) {
                alpha_x2 <- alpha_x2_prop
                W  <- W_prop
                alpha_x2_accept_batch <- alpha_x2_accept_batch + 1 / 50
                alpha_x2_accept <- alpha_x2_accept + 1 / n_mcmc
            }

            if (k %% 50 == 0) {
                out_tune <- update_tuning(k, alpha_x2_accept_batch, alpha_x2_tune)
                alpha_x2_accept_batch <- out_tune$accept
                alpha_x2_tune <- out_tune$tune
            }
        }

        # update alpha_y2 ----
        if (sample_alpha_y2) {
            message("sampling alpha_y2")
            alpha_y2_prop <- rnorm(length(alpha_y2), alpha_y2, alpha_y2_tune)
            # W_prop  <- mra_wendland_2d(cbind(W2 %*% alpha_x2, W2 %*% alpha_y2_prop), M=M, n_coarse_grid = n_coarse_grid)$W
            W_prop  <- eval_basis(cbind(W2 %*% alpha_x2, W2 %*% alpha_y2_prop), grid)$W
            mh1 <- sum(dnorm(y, W_prop %*% alpha, sigma, log=TRUE)) -
                0.5 * sum(alpha_y2_prop * Q2 %*% alpha_y2_prop)
            mh2 <- sum(dnorm(y, W %*% alpha, sigma, log=TRUE)) -
                0.5 * sum(alpha_y2 * Q2 %*% alpha_y2)
            mh <- exp(mh1 - mh2)
            if (mh > runif(1)) {
                alpha_y2 <- alpha_y2_prop
                W  <- W_prop
                alpha_y2_accept_batch <- alpha_y2_accept_batch + 1 / 50
                alpha_y2_accept <- alpha_y2_accept + 1 / n_mcmc
            }

            if (k %% 50 == 0) {
                out_tune <- update_tuning(k, alpha_y2_accept_batch, alpha_y2_tune)
                alpha_y2_accept_batch <- out_tune$accept
                alpha_y2_tune <- out_tune$tune
            }
        }

        # update alpha ----
        if (sample_alpha) {
            message("sampling alpha")
            # sample using MH

            # alpha_prop <- rnorm(length(alpha), alpha, alpha_tune)
            # mh1 <- sum(dnorm(y, W %*% alpha_prop, sigma, log=TRUE)) -
            #     0.5 * sum(alpha_prop * Q %*% alpha_prop)
            # mh2 <- sum(dnorm(y, W %*% alpha, sigma, log=TRUE)) -
            #     0.5 * sum(alpha * Q %*% alpha)
            # mh <- exp(mh1 - mh2)
            # if (mh > runif(1)) {
            #     alpha <- alpha_prop
            #     alpha_accept_batch <- alpha_accept_batch + 1 / 50
            #     alpha_accept <- alpha_accept + 1 / n_mcmc
            # }
            #
            # if (k %% 50 == 0) {
            #     out_tune <- update_tuning(k, alpha_accept_batch, alpha_tune)
            #     alpha_accept_batch <- out_tune$accept
            #     alpha_tune <- out_tune$tune
            # }

            # sample using Gibbs update
            A_alpha <- 1 / sigma^2 * t(W) %*% W + Q
            b_alpha <- 1 / sigma^2 * t(W) %*% y
            alpha   <- as.vector(rmvnorm.canonical.const(1, b_alpha, A_alpha))

        }

        # update sigma ----
        message("sampling sigma")
        sigma <- sqrt(1 / rgamma(1, 1 + 0.5 * N, 1 + 0.5 * sum((y - W %*% alpha)^2)))

        message("sigma = ", sigma)

        # save the parameters
        alpha_x1_save[k, ] <- alpha_x1
        alpha_y1_save[k, ] <- alpha_y1
        # W1_save[[k]]       <- W1
        alpha_x2_save[k, ] <- alpha_x2
        alpha_y2_save[k, ] <- alpha_y2
        # W2_save[[k]]       <- W2
        alpha_save[k, ]    <- alpha
        # W_save[[k]]        <- W
        Walpha_save[k, ]   <- W %*% alpha
        sigma_save[k]      <- sigma

    }

    message("Acceptance rate for alpha_x1 = ", alpha_x1_accept)
    message("Acceptance rate for alpha_y1 = ", alpha_y1_accept)
    message("Acceptance rate for alpha_x2 = ", alpha_x2_accept)
    message("Acceptance rate for alpha_y2 = ", alpha_y2_accept)
    message("Acceptance rate for alpha = ", alpha_accept)

    return(list(alpha_x1 = alpha_x1_save,
                alpha_y1 = alpha_y1_save,
                # W1       = W1_save,
                alpha_x2 = alpha_x2_save,
                alpha_y2 = alpha_y2_save,
                # W2       = W2_save,
                alpha    = alpha_save,
                # W        = W_save,
                Walpha   = Walpha_save,
                sigma    = sigma_save))
}