R/locus_z_core.R

In hruffieux/locus: Large-Scale Variational Inference for Combined Selection of Covariate and Response Variables in Regression Models

# This file is part of the `locus` R package:
#     https://github.com/hruffieux/locus
#
# Internal core function to call the variational algorithm for identity link,
# fixed covariates and no external annotation variables.
# See help of `locus` function for details.
#
locus_z_core_ <- function(Y, X, Z, list_hyper, gam_vb, alpha_vb, mu_beta_vb,
                          sig2_alpha_vb, sig2_beta_vb, tau_vb, tol, maxit, anneal,
                          verbose, batch = "y", full_output = FALSE, debug = TRUE) {

  # Y must have been centered, and X and Z, standardized (except the intercept in Z).

  d <- ncol(Y)
  n <- nrow(Y)
  p <- ncol(X)
  q <- ncol(Z)
  
  # Preparing annealing if any
  #
  if (is.null(anneal)) {
    annealing <- FALSE
    c <- 1
  } else {
    annealing <- TRUE
    ladder <- get_annealing_ladder_(anneal, verbose)
    c <- ladder[1]
  }
  
  eps <- .Machine$double.eps^0.5

  with(list_hyper, {  # list_init not used with the with() function to avoid
                      # copy-on-write for large objects

    m2_alpha <- update_m2_alpha_(alpha_vb, sig2_alpha_vb)

    beta_vb <- update_beta_vb_(gam_vb, mu_beta_vb)
    m2_beta <- update_m2_beta_(gam_vb, mu_beta_vb, sig2_beta_vb, sweep = TRUE)

    mat_x_m1 <- update_mat_x_m1_(X, beta_vb)
    mat_z_mu <- update_mat_z_mu_(Z, alpha_vb)

    rs_gam <- rowSums(gam_vb)
    sum_gam <- sum(rs_gam)

    converged <- FALSE
    lb_new <- -Inf
    it <- 0

    while ((!converged) & (it < maxit)) {

      lb_old <- lb_new
      it <- it + 1

      if (verbose & (it == 1 | it %% 5 == 0))
        cat(paste0("Iteration ", format(it), "... \n"))

      digam_sum <- digamma(c * (a + b + d) - 2 * c + 2) 
      
      # % #
      phi_vb <- update_phi_z_vb_(phi, d, c = c)
      xi_vb <- update_xi_z_vb_(xi, tau_vb, m2_alpha, c = c) ###

      zeta2_inv_vb <- phi_vb / xi_vb
      # % #

      # % #
      lambda_vb <- update_lambda_vb_(lambda, sum_gam, c = c)
      nu_vb <- update_nu_vb_(nu, m2_beta, tau_vb, c = c)

      sig2_inv_vb <- lambda_vb / nu_vb
      # % #

      # % #
      eta_vb <- update_eta_z_vb_(n, q, eta, gam_vb, c = c)
      kappa_vb <- update_kappa_z_vb_(Y, Z, kappa, alpha_vb, beta_vb,
                                     m2_alpha, m2_beta, mat_x_m1, mat_z_mu,
                                     sig2_inv_vb, zeta2_inv_vb, c = c)

      tau_vb <- eta_vb / kappa_vb
      # % #

      sig2_alpha_vb <- update_sig2_alpha_vb_(n, zeta2_inv_vb, tau_vb, c = c)
      sig2_beta_vb <- update_sig2_beta_vb_(n, sig2_inv_vb, tau_vb, c = c)

      log_tau_vb <- update_log_tau_vb_(eta_vb, kappa_vb)
      log_sig2_inv_vb <- update_log_sig2_inv_vb_(lambda_vb, nu_vb)


      # different possible batch-coordinate ascent schemes:

      if (batch == "y") { # optimal scheme

        log_om_vb <- update_log_om_vb(a, digam_sum, rs_gam, c = c)
        log_1_min_om_vb <- update_log_1_min_om_vb(b, d, digam_sum, rs_gam, c = c)

        for (i in sample(1:q)) {
          mat_z_mu <- mat_z_mu - tcrossprod(Z[, i], alpha_vb[i, ])

          alpha_vb[i, ] <- c * sig2_alpha_vb[i, ] * (tau_vb *
                                                      crossprod(Y  - mat_z_mu - mat_x_m1, Z[, i]))

          mat_z_mu <- mat_z_mu + tcrossprod(Z[, i], alpha_vb[i, ])
        }

        # C++ Eigen call for expensive updates
        shuffled_ind <- as.numeric(sample(0:(p-1))) # Zero-based index in C++

        coreZLoop(X, Y, gam_vb, log_om_vb, log_1_min_om_vb, log_sig2_inv_vb,
                  log_tau_vb, beta_vb, mat_x_m1, mat_z_mu, mu_beta_vb,
                  sig2_beta_vb, tau_vb, shuffled_ind, c = c)

        rs_gam <- rowSums(gam_vb)

      } else if (batch == "x") { # used only internally, convergence not ensured

        log_om_vb <- update_log_om_vb(a, digam_sum, rs_gam, c = c)
        log_1_min_om_vb <- update_log_1_min_om_vb(b, d, digam_sum, rs_gam, c = c)

        for (k in sample(1:d)) {

          alpha_vb[, k] <- c * sig2_alpha_vb[, k] * tau_vb[k] *
            (crossprod(Y[, k]  - mat_z_mu[, k] - mat_x_m1[, k], Z) +  (n - 1) * alpha_vb[, k])

          mat_z_mu[, k] <- Z %*% alpha_vb[, k]

          mu_beta_vb[, k] <- c * sig2_beta_vb[k] * tau_vb[k] *
            (crossprod(Y[, k] -  mat_z_mu[, k] - mat_x_m1[, k], X) + (n - 1) * beta_vb[, k])


          gam_vb[, k] <- exp(-log_one_plus_exp_(c * (log_1_min_om_vb - log_om_vb -
                                                  log_tau_vb[k] / 2 - log_sig2_inv_vb / 2 -
                                                  mu_beta_vb[, k] ^ 2 / (2 * sig2_beta_vb[k]) -
                                                  log(sig2_beta_vb[k]) / 2)))

          beta_vb[, k] <- mu_beta_vb[, k] * gam_vb[, k]

          mat_x_m1[, k] <- X %*% beta_vb[, k]

        }

        rs_gam <- rowSums(gam_vb)

      } else if (batch == "x-y") { # used only internally, convergence not ensured

        if (annealing)
          stop("Annealing not implemented for this scheme. Exit.")
        
        log_om_vb <- update_log_om_vb(a, digam_sum, rs_gam, c = c)
        log_1_min_om_vb <- update_log_1_min_om_vb(b, d, digam_sum, rs_gam, c = c)

        alpha_vb <- sweep(sig2_alpha_vb * (crossprod(Z, Y - mat_z_mu - mat_x_m1) + (n - 1) * alpha_vb), 2, tau_vb, `*`)

        mat_z_mu <- Z %*% alpha_vb

        # C++ Eigen call for expensive updates
        coreZBatch(X, Y, gam_vb, log_om_vb, log_1_min_om_vb, log_sig2_inv_vb,
                   log_tau_vb, beta_vb, mat_x_m1, mat_z_mu, mu_beta_vb, sig2_beta_vb, tau_vb)

        rs_gam <- rowSums(gam_vb)

      } else if (batch == "0"){ # no batch, used only internally

        for (k in sample(1:d)) {

          log_om_vb <- update_log_om_vb(a, digam_sum, rs_gam, c = c)
          log_1_min_om_vb <- update_log_1_min_om_vb(b, d, digam_sum, rs_gam, c = c)

          for (i in sample(1:q)) {

            mat_z_mu[, k] <- mat_z_mu[, k] - Z[, i] * alpha_vb[i, k]

            alpha_vb[i, k] <- c * sig2_alpha_vb[i, k] * tau_vb[k] *
              crossprod(Z[, i], Y[,k]  - mat_z_mu[, k] - mat_x_m1[, k])

            mat_z_mu[, k] <- mat_z_mu[, k] + Z[, i] * alpha_vb[i, k]
          }

          for (j in sample(1:p)) {

            mat_x_m1[, k] <- mat_x_m1[, k] - X[, j] * beta_vb[j, k]

            mu_beta_vb[j, k] <- c * sig2_beta_vb[k] * tau_vb[k] *
              crossprod(Y[,k] - mat_x_m1[, k] - mat_z_mu[, k], X[, j])

            gam_vb[j, k] <- exp(-log_one_plus_exp_(c * (log_1_min_om_vb[j] - log_om_vb[j] -
                                                     log_tau_vb[k] / 2 - log_sig2_inv_vb / 2 -
                                                     mu_beta_vb[j, k] ^ 2 / (2 * sig2_beta_vb[k]) -
                                                     log(sig2_beta_vb[k]) / 2)))

            beta_vb[j, k] <- mu_beta_vb[j, k] * gam_vb[j, k]

            mat_x_m1[, k] <- mat_x_m1[, k] + X[, j] * beta_vb[j, k]

          }

          rs_gam <- rowSums(gam_vb)

        }

      } else {

        stop ("Batch scheme not defined. Exit.")

      }

      m2_alpha <- update_m2_alpha_(alpha_vb, sig2_alpha_vb)
      m2_beta <- update_m2_beta_(gam_vb, mu_beta_vb, sig2_beta_vb, sweep = TRUE)

      a_vb <- update_a_vb(a, rs_gam, c = c)
      b_vb <- update_b_vb(b, d, rs_gam, c = c)
      om_vb <- a_vb / (a_vb + b_vb)

      sum_gam <- sum(rs_gam)

      
      if (annealing) {
        
        if (verbose & (it == 1 | it %% 5 == 0))
          cat(paste0("Temperature = ", format(1 / c, digits = 4), "\n\n"))
        
        c <- ifelse(it < length(ladder), ladder[it + 1], 1)
        
        if (isTRUE(all.equal(c, 1))) {
          
          annealing <- FALSE
          
          if (verbose)
            cat("** Exiting annealing mode. **\n\n")
        }
        
      } else {
        
        lb_new <- elbo_z_(Y, Z, a, a_vb, b, b_vb, beta_vb, eta, gam_vb, kappa, lambda,
                          alpha_vb, nu, phi, phi_vb, sig2_alpha_vb,
                          sig2_beta_vb, sig2_inv_vb, tau_vb, xi, zeta2_inv_vb,
                          m2_alpha, m2_beta, mat_x_m1, mat_z_mu, sum_gam)
  
        if (verbose & (it == 1 | it %% 5 == 0))
          cat(paste0("ELBO = ", format(lb_new), "\n\n"))
        
        if (debug && lb_new + eps < lb_old)
          stop("ELBO not increasing monotonically. Exit. ")
        
        converged <- (abs(lb_new - lb_old) < tol)

      }
      
    }

    if (verbose) {
      if (converged) {
        cat(paste0("Convergence obtained after ", format(it), " iterations. \n",
                  "Optimal marginal log-likelihood variational lower bound ",
                  "(ELBO) = ", format(lb_new), ". \n\n"))
      } else {
        warning("Maximal number of iterations reached before convergence. Exit.")
      }
    }

    lb_opt <- lb_new

    names_x <- colnames(X)
    names_y <- colnames(Y)
    names_z <- colnames(Z)
    
    rownames(gam_vb) <- rownames(beta_vb) <- names_x
    colnames(gam_vb) <- colnames(beta_vb) <- names_y
    
    names(om_vb) <- names_x
    rownames(alpha_vb) <- names_z
    colnames(alpha_vb) <- names_y
    
    diff_lb <- abs(lb_opt - lb_old)
    
    if (full_output) { # for internal use only
      
      create_named_list_(a, a_vb, b, b_vb, beta_vb, eta, gam_vb, kappa, lambda,
                         alpha_vb, mu_beta_vb, nu, om_vb, phi, phi_vb, 
                         sig2_alpha_vb, sig2_beta_vb, sig2_inv_vb, tau_vb, xi, 
                         zeta2_inv_vb, m2_alpha, m2_beta, sum_gam, converged, 
                         it, lb_opt, diff_lb, annealing)
    } else {

      create_named_list_(beta_vb, gam_vb, om_vb, alpha_vb, converged, it, lb_opt, diff_lb, annealing)
    }
  })
}



# Internal function which implements the marginal log-likelihood variational
# lower bound (ELBO) corresponding to the `locus_z_core` algorithm.
#
elbo_z_ <- function(Y, Z, a, a_vb, b, b_vb, beta_vb, eta, gam_vb, kappa, lambda,
                    alpha_vb, nu, phi, phi_vb, sig2_alpha_vb, sig2_beta_vb,
                    sig2_inv_vb, tau_vb, xi, zeta2_inv_vb, m2_alpha, m2_beta, 
                    mat_x_m1, mat_z_mu, sum_gam) {


  n <- nrow(Y)
  q <- ncol(Z)

  xi_vb <- update_xi_z_vb_(xi, tau_vb, m2_alpha)

  eta_vb <- update_eta_z_vb_(n, q, eta, gam_vb)

  kappa_vb <- update_kappa_z_vb_(Y, Z, kappa, alpha_vb, beta_vb, m2_alpha,
                                 m2_beta, mat_x_m1, mat_z_mu, sig2_inv_vb,
                                 zeta2_inv_vb)

  lambda_vb <- update_lambda_vb_(lambda, sum_gam)
  nu_vb <- update_nu_vb_(nu, m2_beta, tau_vb)

  log_tau_vb <- digamma(eta_vb) - log(kappa_vb)
  log_zeta2_inv_vb <- digamma(phi_vb) - log(xi_vb)
  log_sig2_inv_vb <- digamma(lambda_vb) - log(nu_vb)
  log_om_vb <- digamma(a_vb) - digamma(a_vb + b_vb)
  log_1_min_om_vb <- digamma(b_vb) - digamma(a_vb + b_vb)


  elbo_A <- e_y_(n, kappa, kappa_vb, log_tau_vb, m2_beta, sig2_inv_vb, tau_vb,
                m2_alpha, zeta2_inv_vb)

  elbo_B <- e_beta_gamma_(gam_vb, log_om_vb, log_1_min_om_vb, log_sig2_inv_vb,
                         log_tau_vb, m2_beta, sig2_beta_vb, sig2_inv_vb, tau_vb)

  elbo_C <- e_tau_(eta, eta_vb, kappa, kappa_vb, log_tau_vb, tau_vb)

  elbo_D <- e_sig2_inv_(lambda, lambda_vb, log_sig2_inv_vb, nu, nu_vb, sig2_inv_vb)

  elbo_E <- e_omega_(a, a_vb, b, b_vb, log_om_vb, log_1_min_om_vb)

  elbo_F <- e_alpha_(m2_alpha, log_tau_vb, log_zeta2_inv_vb, sig2_alpha_vb,
                       tau_vb, zeta2_inv_vb)

  elbo_G <- e_zeta2_inv_(log_zeta2_inv_vb, phi, phi_vb, xi, xi_vb, zeta2_inv_vb)


  elbo_A + elbo_B + elbo_C + elbo_D + elbo_E + elbo_F + elbo_G
}