R/MYULA.R

In qhengncsu/BNRPxMCMC: Bayesian l1 Trend Filtering via Proximal MCMC

#' MYULA transition step
#' @param theta current theta iterate (vector)
#' @param sigma2 current sigma2 iterate
#' @param wimax the maximum number of observations at a given location
#' @param lambda lambda parameter of Moreau envelope
#' @param gradf function handle that computes the gradient of f
#' @param prox function handle of the proximal operator of g
#' @param delta percentage of max step size to use
#' @param t number of theta sampling steps to carry out sigma2
#' @param MH_correct whether to use MH correction or not
#' @param logpi the target conditional density of theta
#' @return the new theta iterate (vector); if using MH correction, return (thetanew,accept)
MYULA <- function(theta,sigma2,wimax,lambda,gradf,prox,delta=1,t=1,MH_correct=F,logpi=NULL){
  if(t>1 & MH_correct){
    stop('If using MH correction, t should be 1!')
  }
  n <- length(theta)

  #calculate the Lipschitz constant L
  L <- wimax/sigma2+1/lambda

  #calculate stepsize gamma
  gamma <- delta/L

  for(i in 1:t){
    #gradient of f at current iterate theta
    gradf_theta <- gradf(theta,sigma2)

    #proximal projection of current iterate theta
    prox_proj_theta <- prox(theta)


    #the constant part of the transition kernel
    normal_mean_forward <- (1-gamma/lambda)*theta-gamma*gradf_theta+gamma/lambda*prox_proj_theta

    #add Gaussian noise to get the new MCMC iterate
    thetanew <- normal_mean_forward + sqrt(2*gamma)*rnorm(n)

    accept <- T

    if(MH_correct){
      #gradient of f at new iterate thetanew
      gradf_thetanew <- gradf(thetanew,sigma2)

      #proximal projection of new iterate thetanew
      prox_proj_thetanew <- prox(thetanew)

      #the constant part of the inverse transition kernel
      normal_mean_backward <- (1-gamma/lambda)*thetanew-gamma*gradf_thetanew+gamma/lambda*prox_proj_thetanew

      #calculate the log difference of q
      log_q_diff <- (norm(thetanew-normal_mean_forward,'2')^2-norm(theta-normal_mean_backward,'2')^2)/(4*gamma)

      if(is.null(logpi)){
        stop('logpi must be specified for approximate correction!')
      }

      #calculate the log difference of pilambda
      log_pi_diff <- logpi(thetanew,prox_proj_thetanew,sigma2)-logpi(theta,prox_proj_theta,sigma2)

      #add them together to get log p
      log_p <- min(0,log_pi_diff+log_q_diff)

      #accept thetanew with probability r
      if(log(runif(1))>log_p){
        thetanew <- theta
        accept <- FALSE
      }
    }

    theta <- thetanew
  }

  if(!MH_correct){
    return(thetanew)
  }else{
    return(list(thetanew=thetanew,accept=accept))
  }
}