R/expect_invariant.R

In mcunit: Unit Tests for MC Methods

#' @importFrom Rdpack reprompt

## Determine a lag such that the ACF is low enough
##
## Internal function.
getthinning <- function(object,control,acfbound=0.5){
    thinning <- 1
    theta <- object$genprior()
    dat <- object$gendata(theta)
    yakt <- object$test(theta)
    res <- matrix(nrow=length(yakt),
                  ncol=500)
    res[,1] <- yakt
    for (i in 2:500){
        theta <- object$stepMCMC(theta,dat,1)
        res[,i] <- object$test(theta)
    }
    thinning <- sapply(1:length(yakt),function(i){
        cor <- stats::acf(res[i,],plot=FALSE)
        if (all(cor$acf>acfbound))
            testthat::fail(paste("Could not automatically compute thinning.",
                                 "Autocorrelation does not decrease below ",
                                 acfbound,"in", length(cor)-1, "steps.\n"))

        min(which(c(cor$acf)<=acfbound))-1})
    
    if (any(is.na(thinning))){
        testthat::fail("Could not automatically compute step size.\n")
        return(-1)
    }
    else{
        if (!is.null(control$debug)&&control$debug>0){
            message("Step size thinning=",thinning," chosen.")
        }
        return(max(thinning))
    }
}

check_object <- function(object){
  if (!is.list(object)) testthat::fail("object must be a list")
  if (!is.function(object$genprior)) testthat::fail("object$genprior must be a function")
  if (!is.function(object$gendata)) testthat::fail("object$gendata must be a function")
  if (!is.function(object$stepMCMC)) testthat::fail("object$stepMCMC must be a function")
}

update_control <- function(control){
  if (is.null(control)) control <- list()
  if (is.null(control$debug)) control$debug <- 0
  control
}

update_object <- function(object){
  if (is.null(object$test)){
      object$test <- function(x, y) x
  }
  else{
    test_nargs <- length(methods::formalArgs(object$test))
    if (test_nargs == 1){
      func <- object$test
      object$test <- function(x,y) func(x)
    }
  }
  object
}


#' Test of MCMC chain assuming reversibility of the chain
#'
#' Test of MCMC steps having the correct stationary distribution
#' assuming reversibility of the chain. Uses ideas from
#' \insertCite{besag_clifford_1989;textual}{mcunit} as extended to
#' possible ties in \insertCite{gandy_scott_2019;textual}{mcunit}.
#'
#' @param object A list describing the MCMC sampler with the following elements:
#' * genprior: A function with no arguments that generates a sample of the prior distribution. No default value. 
#' * gendata: A function that takes as input the parameter value (of the type generated by genprior) and returns the observed data as an arbitrary R object. No default value.
#' * stepMCMC: A function that takes three arguments:
#'     -  theta: the current position of the chain (of the same type as produced by the prior),
#'     - dat: the observed data (of the same type as produced by gendat)
#'     - thinning: the number of steps the chain should take. 1 corresponds to one step. 
#' * test: Function that takes either one or two arguments, and returns a vector with components which will be used for checking the MCMC sampler. The first argument is interpreted as a parameter value, and if a second argument exists, it is interpreted as a data value. An example is the identity function: function(x) x. Alternatively, if you have access to the model's likelihood function, you could use: function(x,y) likelihood(x,y).
#' @inheritParams expect_mc_test
#' @param thinning Amount of thinning for the MCMC chain. 1 corresponds to no thinning. Default: automatically computed to ensure an autocorrelation of at most 0.5 at lag 1.
#' @param  nsteps Number of steps of the chain to use. Default: 10.
#' @param p The probability with which the MCMC updates the parameter as opposed to the data in a given step. If less than 1.0, then the MCMC is a random scan on both data and parameters. Default: 1.0.
#' @param tolerance Absolute error where value of the samplers are treated as equal. Default: 1e-8.
#' @references
#' \insertAllCited{}
#' @examples
#' object <- list(genprior=function() rnorm(1),
#'                gendata=function(theta) rnorm(5,theta),
#'                stepMCMC=function(theta,data,thinning){
#'                   f <- function(x) prod(dnorm(data,x))*dnorm(x)  
#'                   for (i in 1:thinning){
#'                      thetanew = rnorm(1,mean=theta,sd=1)
#'                      if (runif(1)<f(thetanew)/f(theta))
#'                      theta <- thetanew
#'                   }
#'                   theta
#'                },
#'                test=function(x) x
#'                )
#' expect_mcmc_reversible(object)
#' @seealso [expect_mcmc]
#' @return The first argument, invisibly, to allow chaining of expectations.
#' @export
expect_mcmc_reversible <- function(object, control=NULL, thinning=NULL, nsteps=10, p = 1.0, tolerance=1e-8) {
  act <- testthat::quasi_label(rlang::enquo(object), arg = "object")

  check_object(object)
  object  <- update_object(object)
  control <- update_control(control)

  if (is.null(thinning)){ ##autosetting of step size
      thinning <- getthinning(object,control)
  }
  
  if (nsteps<2){
      testthat::fail("nsteps must be at least 2 for this method to work.")
  }
  ##compute dimension of output
  res <-  object$test(object$genprior(), object$gendata(object$genprior()))
  if (!is.vector(res))
      testthat::fail("object$test must return a vector.")
  dimres <- length(res)
  if (dimres==0)
      testthat::fail("object$test returned a vector of length 0.")
  
  pvalsampler <- function(n){
      ranks <- replicate(n,{
          theta <- object$genprior()
          dat <- object$gendata(theta)
          yakt <- object$test(theta, dat)
          k <- sample.int(nsteps,size=1)
          r <- rep(1,length(yakt))
          rties <- rep(0,length(yakt))
          if (k>1){
              x <- theta
              y <- dat
              for (i in (k-1):1){
                  if(stats::runif(1) < p) x <- object$stepMCMC(x,y,thinning)
                  else y <- object$gendata(x)
                  xakt <- object$test(x,y)
                  ##adjust ranks
                  r <- r+(xakt<yakt-tolerance)
                  rties <- rties+(abs(xakt-yakt)<=tolerance) ##ties
              }
          }
          if (k<nsteps){
              x <- theta
              y <- dat
              for (i in (k+1):nsteps){
                  if(stats::runif(1) < p) x <- object$stepMCMC(x,y,thinning)
                  else y <- object$gendata(x)
                  xakt <- object$test(x,y)
                  ##adjust ranks
                  r <- r+(xakt<yakt)
                  rties <- rties+(abs(xakt-yakt)<=tolerance) ##ties
              }
          }
          r + sapply(rties + 1, sample.int, 1) - 1      
      })
      res <- list()
      if (is.vector(ranks))
          res[[1]] <- stats::chisq.test(tabulate(ranks,nbins=nsteps))
      else{
          for (i in 1:(dim(ranks)[1]))
              res[[i]] <-stats::chisq.test(tabulate(ranks[i,],nbins=nsteps))
      }
      sapply(res,function(x) x$p.value)
  }
  expect_mc_test(pvalsampler,control,npval=dimres)

  invisible(act$val)
}


#' Test of MCMC chain
#'
#' Test of MCMC steps having the correct stationary distribution
#' without assuming reversibility of the chain. Details of this are in
#' \insertCite{gandy_scott_2019;textual}{mcunit}; it uses ideas
#' described in the appendix of
#' \insertCite{gandy_veraart_2017;textual}{mcunit}.

#'
#' @inheritParams expect_mcmc_reversible
#' @inheritParams expect_mc_test
#' @param joint If TRUE, then the MCMC uses systematic scan of both data and parameters, rather than just updating parameters with the sampler to be tested. Default: FALSE.
#' @examples
#' object <- list(genprior=function() rnorm(1),
#'                gendata=function(theta) rnorm(5,theta),
#'                stepMCMC=function(theta,data,thinning){
#'                   f <- function(x) prod(dnorm(data,x))*dnorm(x)  
#'                   for (i in 1:thinning){
#'                      thetanew = rnorm(1,mean=theta,sd=1)
#'                      if (runif(1)<f(thetanew)/f(theta))
#'                      theta <- thetanew
#'                   }
#'                   theta
#'                }
#'                )
#' expect_mcmc(object)
#' @references
#' \insertAllCited{}
#' @seealso [expect_mcmc_reversible]
#' @return The first argument, invisibly, to allow chaining of expectations.
#' @export
expect_mcmc <- function(object, control=NULL, thinning=NULL, joint = FALSE) {
  act <- testthat::quasi_label(rlang::enquo(object), arg = "object")

  check_object(object)
  object  <- update_object(object)
  control <- update_control(control)

  if (is.null(thinning)){ ##autosetting of step size
      thinning <- getthinning(object,control)
  }

  ##compute dimension of output
  res <-  object$test(object$genprior(), object$gendata(object$genprior()))
  if (!is.vector(res))
      testthat::fail("object$test must return a vector.")
  dimres <- length(res)
  if (dimres==0)
      testthat::fail("object$test returned a vector of length 0.")
  
  
  pvalsampler <- function(n){
      xall <- replicate(n,{
          x <- object$genprior()
          y <- object$gendata(x)
          if(joint) {
            for(i in 1:thinning)
            {
                x <- object$stepMCMC(x,y,1)
                y <- object$gendata(x)
            }
          }
          else x <- object$stepMCMC(x,y,thinning)
          object$test(x,y)
      })
      priorall <- replicate(n,{
          x <- object$genprior()
          y <- object$gendata(x)
          object$test(x,y)
      })
      res <- list()
      test <- function(x,y) {
          withCallingHandlers({stats::ks.test(x,y)},
                              warning=function(w){invokeRestart("muffleWarning")})
      }
      if (is.vector(xall))
          res[[1]] <- test(xall,priorall)
      else{
          for (i in 1:(dim(xall)[1]))
              res[[i]] <-test(xall[i,],priorall[i,])
      }
      sapply(res,function(x) x$p.value)
  }
  expect_mc_test(pvalsampler,control,npval=dimres)

  invisible(act$val)
}