R/dynamicWhittle_MCMC_zigzag.R

In beyondWhittle: Bayesian Spectral Inference for Time Series

#' MH sampler for BDP-DW method
#'
#' Obtain samples of the posterior of the Whittle likelihood in conjunction with a Bernstein-Dirichlet prior on the spectral density.
#' @details Further detail can be found in the simulation study section in the references papers.
#' @param data numeric vector.
#' @param m window size needed to calculate moving periodogram.
#' @param likelihood_thinning the thinning factor of the dynamic Whittle likelihood.
#' @param mcmc_params a list generated by \link{bdp_dw_mcmc_params_gen}.
#' @param prior_params a list generated by \link{bdp_dw_prior_params_gen}.
#' @param monitor a Boolean value (default FALSE) indicating whether to display the real-time status
#' @param print_interval If monitor = TRUE, then this value indicates the number of iterations after which a status is printed to console;
#' If monitor = FALSE, it does not have any effect
#' @return list containing the following fields:
#'
#'    \item{k1,k2,tau,V,W1,W2}{posterior traces of PSD parameters}
#'    \item{tim}{total run time}
#'    \item{prior_params}{the specifications of the prior}
#' @references Y. Tang et al. (2023)
#' \emph{Bayesian nonparametric spectral analysis of locally stationary processes}
#' ArXiv preprint
#' <arXiv:2303.11561>
#'
#' @importFrom Rcpp evalCpp sourceCpp
#' @importFrom stats cov pbeta qlogis rbinom rnorm rpois runif
#' @useDynLib beyondWhittle, .registration = TRUE
#' @keywords internal
bdp_dw_mcmc <- function(data,
                    m,
                    likelihood_thinning = 1,
                    mcmc_params,
                    prior_params,
                    monitor = FALSE,
                    print_interval = 100)
{

  # MCMC parameters
  stopifnot(!is.null(mcmc_params$Ntotal)); stopifnot(mcmc_params$Ntotal>0)
  Ntotal <- mcmc_params$Ntotal
  stopifnot(!is.null(mcmc_params$burnin)); stopifnot(mcmc_params$burnin>=0 && mcmc_params$burnin<Ntotal)
  burnin <- mcmc_params$burnin
  stopifnot(!is.null(mcmc_params$thin)); stopifnot(mcmc_params$thin>=1)
  thin <- mcmc_params$thin

  #

  # Adaptive MCMC parameters
  stopifnot(!is.null(mcmc_params$adaption.batchSize)); stopifnot(mcmc_params$adaption.batchSize>0)
  adaption.batchSize <- mcmc_params$adaption.batchSize
  stopifnot(!is.null(mcmc_params$adaption.targetAcceptanceRate)); stopifnot(mcmc_params$adaption.targetAcceptanceRate>0 && mcmc_params$adaption.targetAcceptanceRate<1)
  adaption.targetAcceptanceRate <- mcmc_params$adaption.targetAcceptanceRate



  # PRIOR PAREMETERS
  stopifnot(!is.null(prior_params$M)); stopifnot(prior_params$M > 0)
  M <- prior_params$M
  stopifnot(!is.null(prior_params$g0.alpha) && !is.null(prior_params$g0.beta)); stopifnot(prior_params$g0.alpha > 0 && prior_params$g0.beta > 0)
  g0.alpha <- prior_params$g0.alpha
  g0.beta <- prior_params$g0.beta
  stopifnot(!is.null(prior_params$k1.theta)); stopifnot(prior_params$k1.theta > 0)
  k1.theta <- prior_params$k1.theta
  stopifnot(!is.null(prior_params$k2.theta)); stopifnot(prior_params$k2.theta > 0)
  k2.theta <- prior_params$k2.theta
  stopifnot(!is.null(prior_params$tau.alpha) && !is.null(prior_params$tau.beta)); stopifnot(prior_params$tau.alpha > 0 && prior_params$tau.beta > 0)
  tau.alpha <- prior_params$tau.alpha
  tau.beta <- prior_params$tau.beta
  stopifnot(!is.null(prior_params$k1max)); stopifnot(prior_params$k1max > 0)
  k1max <- prior_params$k1max
  stopifnot(!is.null(prior_params$k2max)); stopifnot(prior_params$k2max > 0)
  k2max <- prior_params$k2max
  stopifnot(!is.null(prior_params$L)); stopifnot(prior_params$L > 0)
  L <- prior_params$L
  stopifnot(!is.null(prior_params$bernstein1_l) && !is.null(prior_params$bernstein1_r)); stopifnot(prior_params$bernstein1_l >= 0 && prior_params$bernstein1_r <= 1)
  bernstein1_l <- prior_params$bernstein1_l
  bernstein1_r <- prior_params$bernstein1_r
  stopifnot(!is.null(prior_params$bernstein2_l) && !is.null(prior_params$bernstein2_r)); stopifnot(prior_params$bernstein2_l >= 0 && prior_params$bernstein2_r <= 1)
  bernstein2_l <- prior_params$bernstein2_l
  bernstein2_r <- prior_params$bernstein2_r

  #
  stopifnot(is.logical(monitor))

  stopifnot(!is.null(print_interval)); stopifnot(print_interval>0)

  stopifnot(!is.null(likelihood_thinning)); stopifnot(likelihood_thinning %in% 1:3)

  #
  n <- length(data) - 2*m

  # moving Fourier coefs and time_grid
  mf_list <- local_moving_FT_zigzag(data, m, likelihood_thinning)

  time_grid <- mf_list$time_grid # original scale

  # moving periodograms (with scale 1/(2*pi))
  FZ <- abs(mf_list$MF)^2/(2*pi)

  # grid points
  u1 <- time_grid/n # re-scaled

  u2 <- (2 * (1:m))/(2*m+1)

  # Bernstein basis
  db.list_1 <- dbList_dw_Bern(u1, k1max, bernstein1_l, bernstein1_r)
  db.list_2 <- dbList_dw_Bern_for_lambda(u2, k2max, bernstein2_l, bernstein2_r, m, time_grid)


  # block length
  blk_len_V <- L
  blk_len_W1 <- L+1
  blk_len_W2 <- L+1


  # Open objects for storage
  tau <- rep(NA, Ntotal)
  tilde.V <- matrix(NA, nrow = L, ncol = Ntotal) # reside in [0,1]
  tilde.W1 <- matrix(NA, nrow = L + 1, ncol = Ntotal)
  tilde.W2 <- matrix(NA, nrow = L + 1, ncol = Ntotal)
  k1 <- rep(NA, Ntotal)
  k2 <- rep(NA, Ntotal)
  lpost_storage <- rep(NA, Ntotal)

  # Starting values
  tau[1] <- tau.hat <- mean(FZ)
  k1[1] <- round(k1max / 2)
  k2[1] <- round(k2max / 2)
  beta_basis_1_k <- db.list_1[[k1[1]]]
  beta_basis_2_k <- db.list_2[[k2[1]]]

  tilde.V[, 1] <-  qlogis(rep(1/L, L)) # logistic transformation s.t. the domain becomes the whole space
  tilde.W1[, 1] <- qlogis(seq(from=1/(2*k1[1]), to=1-1/(2*k1[1]), length.out=L+1))
  tilde.W2[, 1] <- qlogis(seq(from=1/(2*k2[1]), to=1-1/(2*k2[1]), length.out=L+1))


  lp.current <- lprior_dw(tilde.V[,1], tilde.W1[, 1], tilde.W2[, 1], k1[1], k2[1], tau[1],
                              M, g0.alpha, g0.beta, k1.theta, k2.theta, tau.alpha, tau.beta)

  norm_psd.current <- qpsd_dw.tilde_zigzag_cpp_expedited(tilde.V[,1], tilde.W1[, 1], tilde.W2[, 1], k1[1], k2[1], beta_basis_1_k, beta_basis_2_k)

  ll.current <- llike_dw(FZ, norm_psd.current, tau[1])
  
  lpost_storage[1] <- lp.current + ll.current

  # First one only used for W0.

  # proposal variances
  eps_tau <- 0.1

  dwd_tau <- log(eps_tau) / 2


  ####

  # Metropolis-within-Gibbs sampler
  tim.start <- Sys.time()
  tim0 <- proc.time()
  for (i in 1:(Ntotal-1)) {

    if (monitor){

      if ((i+1)%%print_interval == 0) {
      tim <- proc.time()
      print_mcmc_state(i+1, Ntotal, tim, tim0)
      }

    }

    # adaptation

        if ((i > 1) && (i <= burnin) && (i %% adaption.batchSize == 1)) {
          batch <- (i - 1):(i - adaption.batchSize)
          adaption.delta <- min(0.05, 1/sqrt(i)) # cf. Section 3 of Roberts and Rosenthal (2009)

          ### tau
          batch.tau <- tau[batch]
          batch.tau.acceptanceRate <- acceptanceRate(batch.tau)
          dwd_tau <- dwd_tau + ((batch.tau.acceptanceRate > adaption.targetAcceptanceRate)*2-1) * adaption.delta
          eps_tau <- exp(2*dwd_tau)
        }


    # STEP 1.1: Sample Bernstein polynomial degree k1

    k1.star <- k1[i] + (rpois(1, lambda = 1) + 1) * (2*rbinom(1, 1, 0.5) - 1)
    k1.star <- (k1.star-1)%%as.integer(k1max) + 1 # circular symmetric proposal


    beta_basis_1_k_star <- db.list_1[[k1.star]]

    lp.k1.star <- lprior_dw(tilde.V[,i], tilde.W1[, i], tilde.W2[, i], k1.star, k2[i], tau[i],
                            M, g0.alpha, g0.beta, k1.theta, k2.theta, tau.alpha, tau.beta)

    norm_psd.k1.star <- qpsd_dw.tilde_zigzag_cpp_expedited(tilde.V[,i], tilde.W1[, i], tilde.W2[, i], k1.star, k2[i], beta_basis_1_k_star, beta_basis_2_k)

    ll.k1.star <- llike_dw(FZ, norm_psd.k1.star, tau[i])

    #####
    # Accept/reject

    alpha_k1 <- min(0, (lp.k1.star + ll.k1.star) - (lp.current + ll.current))
    if (isTRUE(log(runif(1, 0, 1)) < alpha_k1)) {
      k1[i + 1] <- k1.star  # Accept k1.star
      lp.current <- lp.k1.star
      norm_psd.current <- norm_psd.k1.star
      ll.current <- ll.k1.star

      beta_basis_1_k <- beta_basis_1_k_star
      lpost_storage[i + 1] <- lp.current + ll.current
    } else {
      k1[i + 1] <- k1[i]  # Reject and use previous
      lpost_storage[i + 1] <- lpost_storage[i]
    }


    # STEP 1.2: Sample Bernstein polynomial degree k2

    k2.star <- k2[i] + (rpois(1, lambda = 1) + 1) * (2*rbinom(1, 1, 0.5) - 1)
    k2.star <- (k2.star-1)%%as.integer(k2max) + 1 # circular symmetric proposal


    beta_basis_2_k_star <- db.list_2[[k2.star]]


    lp.k2.star <- lprior_dw(tilde.V[,i], tilde.W1[, i], tilde.W2[, i], k1[i+1], k2.star, tau[i],
                            M, g0.alpha, g0.beta, k1.theta, k2.theta, tau.alpha, tau.beta)

    norm_psd.k2.star <- qpsd_dw.tilde_zigzag_cpp_expedited(tilde.V[,i], tilde.W1[, i], tilde.W2[, i], k1[i+1], k2.star, beta_basis_1_k, beta_basis_2_k_star)

    ll.k2.star <- llike_dw(FZ, norm_psd.k2.star, tau[i])

    #####
    # Accept/reject

    alpha_k2 <- min(0, (lp.k2.star + ll.k2.star) - (lp.current + ll.current))
    if (isTRUE(log(runif(1, 0, 1)) < alpha_k2)) {
      k2[i + 1] <- k2.star  # Accept k.star
      lp.current <- lp.k2.star
      norm_psd.current <- norm_psd.k2.star
      ll.current <- ll.k2.star
      
      beta_basis_2_k <- beta_basis_2_k_star
      lpost_storage[i + 1] <- lp.current + ll.current
    } else {
      k2[i + 1] <- k2[i]  # Reject and use previous
      lpost_storage[i + 1] <- lpost_storage[i]
    }


    # block sampling V  (cf. Section 2 of Roberts and Rosenthal (2009))
    vec.curr_V <- tilde.V[, i]

    if (i <= 2 * blk_len_V){
      tilde.V.star <- vec.curr_V + 0.1/sqrt(blk_len_V) * rnorm(blk_len_V)
    } else{
      if (i == 2 * blk_len_V + 1){
        # a matrix of dimension blk_len_V by blk_len_V
        cov.curr_V <- cov(t(tilde.V[, (1:i)]))

        # a vector of length blk_len_V
        mean.curr_V <- colMeans(t(tilde.V[, (1:i)]))
      } else{
        mean.curr_V <- (1 - 1/i) * mean.curr_V + vec.curr_V/i
        cov.curr_V <- (1 - 1/i) * (cov.curr_V + tcrossprod(vec.curr_V - mean.curr_V, vec.curr_V - mean.curr_V)/i)
      }

      if (rbinom(1, size = 1, prob = 0.95)){
        eS <- eigen(cov.curr_V, symmetric = TRUE)
        tilde.V.star <- vec.curr_V + 2.38/sqrt(blk_len_V) * drop(eS$vectors %*% diag(sqrt(pmax(eS$values, 0))) %*% rnorm(blk_len_V))

      } else{
        tilde.V.star <- vec.curr_V + 0.1/sqrt(blk_len_V) * rnorm(blk_len_V)
      }
    }


    lp.V.star <- lprior_dw(tilde.V.star, tilde.W1[,i], tilde.W2[,i], k1[i+1], k2[i+1], tau[i],
                                 M, g0.alpha, g0.beta, k1.theta, k2.theta, tau.alpha, tau.beta)

    norm_psd.V.star <- qpsd_dw.tilde_zigzag_cpp_expedited(tilde.V.star, tilde.W1[,i], tilde.W2[,i], k1[i+1], k2[i+1], beta_basis_1_k, beta_basis_2_k)

    ll.V.star <- llike_dw(FZ, norm_psd.V.star, tau[i])

    #####
    # Accept/reject

    alpha_V <- min(0, (ll.V.star - ll.current) + (lp.V.star - lp.current))
    if (isTRUE(log(runif(1, 0, 1)) < alpha_V)) {# accept the proposed value
      tilde.V[, i+1] <- tilde.V.star

      lp.current <- lp.V.star
      norm_psd.current <- norm_psd.V.star
      ll.current <- ll.V.star
      
      lpost_storage[i + 1] <- lp.current + ll.current

    } else {# Reject and use previous
      tilde.V[, i+1] <- tilde.V[, i]
      
      lpost_storage[i + 1] <- lpost_storage[i]
    }


    # block sampling W1  (cf. Section 2 of Roberts and Rosenthal (2009))
    vec.curr_W1 <- tilde.W1[, i]

    if (i <= 2 * blk_len_W1){
      tilde.W1.star <- vec.curr_W1 + 0.1/sqrt(blk_len_W1) * rnorm(blk_len_W1)
    } else{
      if (i == 2 * blk_len_W1 + 1){
        # a matrix of dimension blk_len_W1 by blk_len_W1
        cov.curr_W1 <- cov(t(tilde.W1[, (1:i)]))

        # a vector of length blk_len_W1
        mean.curr_W1 <- colMeans(t(tilde.W1[, (1:i)]))
      } else{
        mean.curr_W1 <- (1 - 1/i) * mean.curr_W1 + vec.curr_W1/i
        cov.curr_W1 <- (1 - 1/i) * (cov.curr_W1 + tcrossprod(vec.curr_W1 - mean.curr_W1, vec.curr_W1 - mean.curr_W1)/i)
      }

      if (rbinom(1, size = 1, prob = 0.95)){
        eS <- eigen(cov.curr_W1, symmetric = TRUE)
        tilde.W1.star <- vec.curr_W1 + 2.38/sqrt(blk_len_W1) * drop(eS$vectors %*% diag(sqrt(pmax(eS$values, 0))) %*% rnorm(blk_len_W1))

      } else{
        tilde.W1.star <- vec.curr_W1 + 0.1/sqrt(blk_len_W1) * rnorm(blk_len_W1)
      }
    }


    lp.W1.star <- lprior_dw(tilde.V[,i+1], tilde.W1.star, tilde.W2[,i], k1[i+1], k2[i+1], tau[i],
                                  M, g0.alpha, g0.beta, k1.theta, k2.theta, tau.alpha, tau.beta)

    norm_psd.W1.star <- qpsd_dw.tilde_zigzag_cpp_expedited(tilde.V[,i+1], tilde.W1.star, tilde.W2[,i], k1[i+1], k2[i+1], beta_basis_1_k, beta_basis_2_k)

    ll.W1.star <- llike_dw(FZ, norm_psd.W1.star, tau[i])

    #####
    # Accept/reject

    alpha_W1 <- min(0, (ll.W1.star - ll.current) + (lp.W1.star - lp.current))
    if (isTRUE(log(runif(1, 0, 1)) < alpha_W1)) {# accept the proposed value
      tilde.W1[, i+1] <- tilde.W1.star

      lp.current <- lp.W1.star
      norm_psd.current <- norm_psd.W1.star
      ll.current <- ll.W1.star
      
      lpost_storage[i + 1] <- lp.current + ll.current

    } else {# Reject and use previous
      tilde.W1[, i+1] <- tilde.W1[, i]
      
      lpost_storage[i + 1] <- lpost_storage[i]
    }


    # block sampling W2  (cf. Section 2 of Roberts and Rosenthal (2009))
    vec.curr_W2 <- tilde.W2[, i]

    if (i <= 2 * blk_len_W2){
      tilde.W2.star <- vec.curr_W2 + 0.1/sqrt(blk_len_W2) * rnorm(blk_len_W2)
    } else{
      if (i == 2 * blk_len_W1 + 1){
        # a matrix of dimension blk_len_W2 by blk_len_W2
        cov.curr_W2 <- cov(t(tilde.W2[, (1:i)]))

        # a vector of length blk_len_W2
        mean.curr_W2 <- colMeans(t(tilde.W2[, (1:i)]))
      } else{
        mean.curr_W2 <- (1 - 1/i) * mean.curr_W2 + vec.curr_W2/i
        cov.curr_W2 <- (1 - 1/i) * (cov.curr_W2 + tcrossprod(vec.curr_W2 - mean.curr_W2, vec.curr_W2 - mean.curr_W2)/i)
      }

      if (rbinom(1, size = 1, prob = 0.95)){
        eS <- eigen(cov.curr_W2, symmetric = TRUE)
        tilde.W2.star <- vec.curr_W2 + 2.38/sqrt(blk_len_W2) * drop(eS$vectors %*% diag(sqrt(pmax(eS$values, 0))) %*% rnorm(blk_len_W2))

      } else{
        tilde.W2.star <- vec.curr_W2 + 0.1/sqrt(blk_len_W2) * rnorm(blk_len_W2)
      }
    }


    lp.W2.star <- lprior_dw(tilde.V[,i+1], tilde.W1[,i+1], tilde.W2.star, k1[i+1], k2[i+1], tau[i],
                                  M, g0.alpha, g0.beta, k1.theta, k2.theta, tau.alpha, tau.beta)

    norm_psd.W2.star <- qpsd_dw.tilde_zigzag_cpp_expedited(tilde.V[,i+1], tilde.W1[,i+1], tilde.W2.star, k1[i+1], k2[i+1], beta_basis_1_k, beta_basis_2_k)

    ll.W2.star <- llike_dw(FZ, norm_psd.W2.star, tau[i])

    #####
    # Accept/reject

    alpha_W2 <- min(0, (ll.W2.star - ll.current) + (lp.W2.star - lp.current))
    if (isTRUE(log(runif(1, 0, 1)) < alpha_W2)) {# accept the proposed value
      tilde.W2[, i+1] <- tilde.W2.star

      lp.current <- lp.W2.star
      norm_psd.current <- norm_psd.W2.star
      ll.current <- ll.W2.star
      
      lpost_storage[i + 1] <- lp.current + ll.current

    } else {# Reject and use previous
      tilde.W2[, i+1] <- tilde.W2[, i]
      
      lpost_storage[i + 1] <- lpost_storage[i]
    }

    # sampling tau

    tau.star <- exp(log(tau.hat) + rnorm(1, mean = 0, sd = eps_tau))

    lp.tau.star <- lprior_dw(tilde.V[, i+1], tilde.W1[, i+1], tilde.W2[, i+1], k1[i+1], k2[i+1], tau.star,
                             M, g0.alpha, g0.beta, k1.theta, k2.theta, tau.alpha, tau.beta)

    ll.tau.star <- llike_dw(FZ, norm_psd.current, tau.star)

    #####
    # Accept/reject

    alpha_tau <- min(0, (ll.tau.star  - ll.current) + (lp.tau.star  - lp.current))
    if (isTRUE(log(runif(1, 0, 1)) < alpha_tau)) {# Accept tau.star
      tau[i + 1] <- tau.star

      lp.current <- lp.tau.star
      ll.current <- ll.tau.star
      
      lpost_storage[i + 1] <- lp.current + ll.current
    } else {# Reject and use previous
      tau[i + 1] <- tau[i]
      
      lpost_storage[i + 1] <- lpost_storage[i]
    }

  }  # END: MCMC loop

  tim.end <- Sys.time()

  # Which iterations to keep
  keep <- seq(burnin + 1, Ntotal, by = thin)
  k1 <- k1[keep]
  k2 <- k2[keep]
  tau <- tau[keep]
  V <- plogis(tilde.V[, keep])
  W1 <- plogis(tilde.W1[, keep])
  W2 <- plogis(tilde.W2[, keep])
  lpost_storage <- lpost_storage[keep]

  return(list(k1 = k1, k2 = k2, tau = tau, V = V, W1 = W1, W2 = W2,
              tim = tim.end - tim.start, prior_params = prior_params, lpost = lpost_storage, data = data)
         )
}