R/adaptive_is.R

In thaos/AISamplR:

compute_nextmu_apis <- function(xmat, weights){
  #apply(sweep(xmat, 2, weights, "*"), 1, sum)
  if(any(is.na(weights))) stop("some weights are NA")
  apply(xmat, 1, weighted.mean, w = weights)
}

compute_nextmu_pmc <- function(xmat, weights){
  if(any(is.na(weights))) stop("some weights are NA")
  xmat[, as.logical(rmultinom(1, 1, weights))]
}

# better to create a function compute_ais_weights(obj, compute_denom, reuse_weight)
compute_ais_weights <- function(pchain, sigma, compute_denom, reuse_weights){
  if(reuse_weights & !is.null(pchain$denom)) denom_arr <- pchain$denom
  else denom_arr <- compute_denom(pchain$x, pchain$mu, sigma)
  weights <- compute_weights(pchain$pi, denom_arr)
  list(x = pchain$x, mu = pchain$mu, w = weights, pi = pchain$pi, denom = denom_arr, marglik = mean(weights))
}

create_adaptive_is <- function(parallel_chain){
  adaptive_is <- function(logposterior, d, N = 10, T = 100, M = 2, mu, sigma, compute_denom = compute_denomtable_byrow, reuse_weights = FALSE){
    pchain <- parallel_chain(d = d, N = N, T = T, M = M, mu = mu, sigma = sigma, logposterior = logposterior)
    compute_ais_weights(pchain, sigma, compute_denom, reuse_weights)
  }
}

#' @export
#' Layered Adaptive Importance Sampling (LAIS).
#'
#' \code{lais} returns samples and the associated weights generated by LAIS.
#'
#' This function performs adaptive important sampling using the method Layered Adaptive Importance Sampling (LAIS). New location parameters for the proposal distributions are provided by a Metripolis-Hasting (MH) step. The proposal distribution are gaussians the means of which are given by the MH step used for adaptation.
#'
#' @param logposterior A function taking as input a vector of length d. It corresponds to the log-posterior likehood of the distrubution one wants to simulate from.
#' @param d  An integer indicating the number of dimension of the logposterior.
#' @param N An integer indicating the number of chains used for the adaptation.
#' @param T An integer indicating the length of the chains used for the adaptation.
#' @param M An integer indicating the numbers of draws to be made for each proposals.
#' @param mu A numerical matrix of dimension d x N used to initialized the location parameters of the N proposal chains.
#' @param sigma A numerical vector of length d indicating the standard deviation to be used for each dimension of the gaussian proposal (each dimension is drawn independently)
#' @param sigma_mh A numerical vector of length d indicating the standard deviation to be used for each dimension of the proposal used in the MH step (only for the lais function). By default, it is equal to sigma.
#' @param compute_denom A function that indicates how to compute denominators used to compute the weights. By default, compute_denomtable_byrow is used where the denominator for one drawn sample x_t,m,n is computed as mixture of the N proposals at time t.
#' @param reuse_weights a boolean indicating whether to reuse the weight computed in the adaptive steps for the importance sampling step. For lais, the weight in the a
#' @return a list with the following elements. 
##' \itemize{
##'  \item{"x"}{an array of dimension T x d x M x N containing the samples drawn from LAIS.}
##'  \item{"mu"}{an array of dimension T x d x N containing the location parameters of the poposal distribution used in the importance sampling step.}
##'  \item{"weights"}{an array of dimension T x M x N containing the unormalized weights associated with the samples in x.}
##'  \item{"pi"}{an array of dimension T x M x N the loglikelihood evaluated for each sample drawn in x.} 
##'  \item{"denom"}{an array of dimension T x M x N the denominators of the weight for each sample drawn in x.} 
##'  \item{"marglik"}{an estimation of the marginal likelihood.}
##' }
#' @examples
#' # draw samples from the distribution defined by the loglikelihood lposterior_3 which is a mixture of two gaussians of dimension 2.
#' lais_chain <- lais(logposterior = lposterior_3  d = 2, N = 20, T = 100, M = 5, mu = rnorm(40, sd = 5), sigma = rep(sqrt(13), 2), sigma_mh = rep(5, 2), compute_denom = compute_denomtable_byrow)
#' # estimate the expected value of the distribution given by lposterior_3 (the real expected value is [2.5, 8]
#' compute_expectation(lais_chain$x, lais_chain$w)
#' \dontrun{
#' # Similar example with a mixture of 5 gaussians. the real expected value is [1.6, 1.4]
#' lais_chain <- lais(logposterior = lposterior_1,  d = 2, N = 100, T = 100, M = 10, mu = rnorm(100, sd = 5), sigma = rep(sqrt(13), 2), sigma_mh = rep(5, 2), compute_denom = compute_denomtable_byrow)
#' compute_expectation(lais_chain$x, lais_chain$w)
#' }
lais <- function(logposterior, d, N = 10, T = 100, M = 2, mu, sigma, sigma_mh = sigma, compute_denom = compute_denomtable_byrow, reuse_weights = FALSE){
    pchain <- indep_chains_mcmc(logposterior = logposterior, d = d, N = N, T = T, M = M, mu = mu, sigma = sigma, sigma_mh = sigma_mh)
    compute_ais_weights(pchain, sigma, compute_denom, reuse_weights = FALSE)
}
indep_apis <- create_adaptive_is(indep_chains_apis)
indep_pmc <- create_adaptive_is(indep_chains_pmc)
apis <- create_adaptive_is(fchain_apis )
pmc <- create_adaptive_is(fchain_pmc)