add documentation

docs/Project.toml

@@ -1,3 +1,4 @@
 [deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 MCMCTesting = "9963b6a1-5d46-439c-8efc-3a487843c7fa"

docs/make.jl

@@ -5,7 +5,7 @@ DocMeta.setdocmeta!(MCMCTesting, :DocTestSetup, :(using MCMCTesting); recursive=
 
 makedocs(;
     modules=[MCMCTesting],
-    authors="Ray Kim <[email protected]> and contributors",
+    authors="Kyurae Kim <[email protected]> and contributors",
     repo="https://github.com/Red-Portal/MCMCTesting.jl/blob/{commit}{path}#{line}",
     sitename="MCMCTesting.jl",
     format=Documenter.HTML(;
@@ -15,11 +15,16 @@ makedocs(;
         assets=String[],
     ),
     pages=[
-        "Home" => "index.md",
+        "MCMCTesting"       => "introduction.md",
+        "Getting Started"   => "example.md",
+        "Two-Sample Tests"  => "twosampletest.md",
+        "Exact Rank Tests"  => "exactranktest.md",
+        "API"               => "index.md",
     ],
 )
 
 deploydocs(;
     repo="github.com/Red-Portal/MCMCTesting.jl",
+    push_preview=true,
     devbranch="main",
 )

docs/src/exactranktest.md

@@ -0,0 +1,36 @@
+
+# [Exact Rank Hypothesis Tests](@id exactrank)
+
+## Introduction
+
+The exact rank hypothesis testing strategy is based on Algorithm 2 by Gandy & Scott (2021)[^gandyandscott2021].
+
+## `ExactRankTest`
+The ranks are computed by simulating a single Markov chain backwards and forward.
+First, a random midpoint is simulated as
+```math
+\begin{aligned}
+   M \sim \mathrm{Uniform}(1,\;L),
+\end{aligned}
+```
+where $L = \texttt{n\_mcmc\_steps}$.
+Then, we simulate the Markov chain forward and backward as
+```math
+\begin{alignat*}{3}
+  \theta_{M-l} &\sim K\left(\theta_{L-l+1}, \cdot \right) \qquad &&\text{for}\; l = 1, \ldots, M-1 \\
+  \theta_{l}   &\sim K\left(\theta_{l-1}, \cdot \right) \qquad &&\text{for}\; l = M+1, \ldots, L
+\end{alignat*}
+```
+forming the chain
+```math
+\theta_1,\; \ldots, \; \theta_{M},\; \ldots, \; \theta_{L}.
+```
+The *rank* is the ranking of the statistics of $\theta_{M}$.
+If the sampler and the model are correct, the rank has an uniform distribution as long as the midpoint is independently sampled.
+
+```@docs
+ExactRankTest
+```
+
+# References
+[^gandyandscott2021]: Gandy, A., & Scott, J. (2020). Unit testing for MCMC and other Monte Carlo methods. arXiv preprint arXiv:2001.06465.

docs/src/example.md

@@ -0,0 +1,133 @@
+
+# Getting Started
+
+## Problem Setup
+Let's consider a simple Normal-Normal model with a shared scale:
+```math
+\begin{aligned}
+\theta &\sim \mathrm{normal}\left(0,   \sigma^2\right) \\
+y      &\sim \mathrm{normal}\left(\mu, \sigma^2\right)
+\end{aligned}
+nothing
+```
+
+The joint log-likelihood can be implemented as follows:
+```@example started
+using Random
+using Distributions
+
+struct Model
+    σ::Float64
+    y::Float64
+end
+
+function logdensity(model::Model, θ)
+    σ, y = model.σ, model.y
+    logpdf(Normal(0, σ), only(θ)) + logpdf(Normal(0, σ), y)
+end
+nothing
+```
+For sampling from the posterior, a simple Gibbs sampling strategy is possible:
+```@example started
+struct Gibbs end
+
+function step(rng::Random.AbstractRNG, model::Model, ::Gibbs, θ)
+    y = model.y
+	σ = model.σ
+    rand(rng, MvNormal([y]/2, σ))
+end
+nothing
+```
+We could also use the classic symmetric random walk Metropolis-Hastings sampler:
+```@example started
+struct RWMH
+    σ::Float64
+end
+
+function step(rng::Random.AbstractRNG, model::Model, kernel::RWMH, θ)
+	σ = kernel.σ
+    θ′ = rand(rng, MvNormal(θ, σ))
+    ℓπ = logdensity(model, θ)
+    ℓπ′ = logdensity(model, θ′)
+    ℓα = ℓπ′ - ℓπ
+    if log(rand(rng)) ≤ ℓα
+        θ′
+    else
+        θ
+    end
+end
+nothing
+```
+
+## Testing the Gibbs Sampler
+All of the functionalities of `MCMCTesting` assume that we can sample from the joint distribution $p(\theta, y)$.
+This is done by as follows:
+```@example started
+using MCMCTesting
+
+function MCMCTesting.sample_joint(rng::Random.AbstractRNG, model::Model)
+	σ = model.σ
+    θ = rand(rng, Normal(0, σ))
+    y = rand(rng, Normal(θ, σ))
+    [θ], [y]
+end
+nothing
+```
+
+The Gibbs sampler can be connected to `MCMCTesting` by implementing the following:
+
+```@example started
+using Accessors
+
+function MCMCTesting.markovchain_transition(
+    rng::Random.AbstractRNG, model::Model, kernel, θ, y
+)
+	model′ = @set model.y = only(y)
+    step(rng, model′, kernel, θ)
+end
+nothing
+```
+Let's check that the implementation is correct by 
+
+```@example started
+model = Model(1., 1.)
+kernel = Gibbs()
+test = TwoSampleTest(100, 100)
+subject = TestSubject(model, kernel)
+seqmcmctest(test, subject, 0.0001, 100; show_progress=false)
+```
+`true` means that the tests have passed.
+Now, let's consider two erroneous implementations:
+
+```@example started
+struct GibbsWrongMean end
+
+function step(rng::Random.AbstractRNG, model::Model, ::GibbsWrongMean, θ)
+    y = model.y
+	σ = model.σ
+    rand(rng, MvNormal([y], σ/2))
+end
+
+struct GibbsWrongVar end
+
+function step(rng::Random.AbstractRNG, model::Model, ::GibbsWrongVar, θ)
+    y = model.y
+	σ = model.σ
+    rand(rng, MvNormal([y/2], 2*σ))
+end
+nothing
+```
+The kernel with a wrong mean fails:
+```@example started
+kernel = GibbsWrongMean()
+subject = TestSubject(model, kernel)
+seqmcmctest(test, subject, 0.0001, 100; show_progress=false)
+```
+and so does the one with the wrong variance:
+```@example started
+kernel = GibbsWrongVar()
+subject = TestSubject(model, kernel)
+seqmcmctest(test, subject, 0.0001, 100; show_progress=false)
+```
+
+## Visualizing the ranks

docs/src/introduction.md

@@ -0,0 +1,16 @@
+
+# MCMCTesting
+
+This package provides the MCMC testing algorithms proposed by Gandy & Scott (2021)[^gandyandscott2021].
+These tests can be seen as an improvement of the hypothesis testing approach proposed by Geweke [^geweke2004].
+Unlike simulation-based calibration (SBC; [^talts2018][^yaoanddomke2023][^modrak2022]), these tests are more appropriate for testing the exactness of the MCMC algorithm rather than the identifiability of the models.
+This is because the tests focus on maximizing the power for verify the validity of *individual Markov transitions* instead of a *set of samples*.
+Furthermore, unlike SBC, the approach of Gandy & Scott[^gandyandscott2021] is able to exactly satisfy the assumptions required for the theoretical guarantees.
+
+# References
+[^gandyandscott2021]: Gandy, A., & Scott, J. (2020). Unit testing for MCMC and other Monte Carlo methods. arXiv preprint arXiv:2001.06465.
+[^geweke2004]: Geweke, J. (2004). Getting it right: Joint distribution tests of posterior simulators. Journal of the American Statistical Association, 99(467), 799-804.
+[^talts2018]: Talts, S., Betancourt, M., Simpson, D., Vehtari, A., & Gelman, A. (2018). Validating Bayesian inference algorithms with simulation-based calibration. arXiv preprint arXiv:1804.06788.
+[^yaoanddomke2023]: Yao, Y., & Domke, J. (2023, November). Discriminative Calibration: Check Bayesian Computation from Simulations and Flexible Classifier. In Thirty-seventh Conference on Neural Information Processing Systems.
+[^modrak2022]: Modrák, M., Moon, A. H., Kim, S., Bürkner, P., Huurre, N., Faltejsková, K., ... & Vehtari, A. (2022). Simulation-based calibration checking for Bayesian computation: The choice of test quantities shapes sensitivity. arXiv preprint arXiv:2211.02383.
+

docs/src/twosampletest.md

@@ -0,0 +1,53 @@
+
+# [Two Sample Hypothesis Tests](@id twosample)
+
+## Introduction
+
+The two-sample hypothesis testing strategies are based on Algorithm 1 by Gandy & Scott (2021)[^gandyandscott2021].
+
+## `TwoSampleTest`
+The first basic strategy generates the **treatment group**
+```math
+\begin{aligned}
+  \theta,\; y_{\text{trtm}} &\sim p\left(\theta, y\right) \\
+  \theta_{\text{trtm}} &\sim K\left(\theta, \cdot\right),
+\end{aligned}
+```
+and the **control group** as
+```math
+\begin{aligned}
+  \theta_{\text{ctrl}},\; y_{\text{ctrl}} &\sim p\left(\theta, y\right), \\
+\end{aligned}
+```
+where the test compares 
+```math
+(\theta_{\text{trtm}}, \, y_{\text{trtm}})
+\quad\text{versus}\quad 
+(\theta_{\text{ctrl}}, \, y_{\text{ctrl}}).
+```
+
+```@docs
+TwoSampleTest
+```
+
+## `TwoSampleGibbsTest`
+The second strategy performs applies an additional Gibbs sampling step when generating the **treatment group**  as
+```math
+\begin{aligned}
+  \theta,\; y &\sim p\left(\theta, y\right) \\
+  \theta_{\text{trtm}} &\sim K\left(\theta, \cdot\right), \\
+  y_{\text{trtm}} &\sim p\left(y \mid \theta_{\text{trtm}}\right)
+\end{aligned}
+```
+resulting in the treatment group 
+```math
+(\theta_{\text{trtm}}, \, y_{\text{trtm}}).
+```
+The control group is generated the same as `TwoSampleTest`.
+
+```@docs
+TwoSampleGibbsTest
+```
+
+# References
+[^gandyandscott2021]: Gandy, A., & Scott, J. (2020). Unit testing for MCMC and other Monte Carlo methods. arXiv preprint arXiv:2001.06465.

src/MCMCTesting.jl

@@ -19,11 +19,102 @@ using HypothesisTests
 using ProgressMeter
 using MultipleTesting
 
+"""
+    sample_joint(rng, model)
+
+Sample from the joint distribution of the prior and the predictive distribution of `model`.
+
+# Arguments
+- `rng::Random.AbstractRNG`: Random number generator.
+- `model`: Model subject to test.
+
+# Returns
+- `θ`: Model parameter sampled from the prior `p(θ)`.
+- `y`: Data generated from conditionally on `θ` from `p(y|θ)`
+"""
 function sample_joint            end
+
+"""
+    sample_predictive(rng, model, θ)
+
+Sample from the predictive distribution of `model` conditionally on `θ`
+
+# Arguments
+- `rng::Random.AbstractRNG`: Random number generator.
+- `model`: Model subject to test.
+- `θ`: Model parameters to condition on.
+
+# Returns
+- `y`: Data generated from conditionally on `θ` from `p(y|θ)`
+"""
 function sample_predictive       end
+
+"""
+    markovchain_transition(rng, model, kernel, θ, y)
+
+Perform a single Markov chain transition of `kernel` on the previous state `θ` targeting the posterior of `model` conditioned on `y`.
+
+# Arguments
+- `rng::Random.AbstractRNG`: Random number generator.
+- `model`: Model forming the posterior `p(θ|y)` conditioned on `y`.
+- `θ`: Previous state of the Markov chain.
+- `y`: Data to condition on.
+
+# Returns
+- `θ′`: Next state of the Markov chain.
+"""
 function markovchain_transition  end
-function mcmctest                end
 
+"""
+    mcmctest([rng,] test, subject; kwargs...)
+
+Sample a p-value according to `test` for `subject`
+
+# Arguments
+- `rng::Random.AbstractRNG`: Random number generator. (Default: `Random.default_rng()`.)
+- `test::AbstractMCMCTest`: Test strategy.
+- `subject::TestSubject`: MCMC algorithm and model subject to test.
+
+# Keyword Arguments
+- `show_progress::Bool`: Whether to show the progress bar. (Default: `true`.)
+- `statistics`: Function for computing test statistics from samples generated from the tests. (See section below for additional description.)
+- Check the documentation for the respective test strategy for additional keyword arugments.
+
+# Custom Test Statistics
+The statistics used for the hypothesis tests can be modified by passing a custom funciton to `statistics`.
+The default statistics are the first and second moments computed as below.
+```julia
+statistics = params -> vcat(params, params.^2)
+```
+The cross-interaction can also be tested by adding an additional entry as below.
+```julia
+statistics = params -> vcat(params, params.^2, reshape(params*params',:))
+```
+But naturally, adding more statistics increase the computational cost of computing the tests.
+
+Also, different tests may result in different statistics being computed through the same `statistics` function.
+For example, the two-sample test strategies generate both model parameters `θ` and data `y`.
+Therefore, `params = vcat(θ, y)`.
+On the other hand, the exac rank test only generates model parameters `θ`.
+Therefore, `params = θ`.
+Naturally, `statistics` can also be used to `select` a subset of parameters used for testing.
+For example, for the two-sample test strategies, if we only want to use `θ` for the tests, where `d = length(θ) > 0`, one can do the following:
+```julia
+statistics = params -> θ[1:d]
+```
+
+"""
+function mcmctest end
+
+"""
+    TestSubject(model, kernel)
+
+Model and MCMC kernel obejct subject to test.
+
+# Arguments
+- `model`: Model subject to test.
+- `kernel`: MCMC kernel subject to test.
+"""
 struct TestSubject{M, K}
     model ::M
     kernel::K