Better documentation

docs/Project.toml

@@ -1,3 +1,9 @@
 [deps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
+
+[compat]
+AdvancedHMC = "0.7"
+Documenter = "1"
+DocumenterCitations = "1"

docs/make.jl

@@ -1,8 +1,28 @@
 using Pkg
 
-using Documenter
+using Documenter, DocumenterCitations
 using AdvancedHMC
 
-# cp(joinpath(@__DIR__, "../README.md"), joinpath(@__DIR__, "src/index.md"))
+bib = CitationBibliography(joinpath(@__DIR__, "src", "refs.bib"))
 
-makedocs(; sitename="AdvancedHMC", format=Documenter.HTML(), warnonly=[:cross_references])
+makedocs(;
+    sitename="AdvancedHMC",
+    format=Documenter.HTML(;
+        assets=["assets/favicon.ico"],
+        canonical="https://turinglang.org/AdvancedHMC.jl/stable/",
+    ),
+    warnonly=[:cross_references],
+    plugins=[bib],
+    pages=[
+        "AdvancedHMC.jl" => "index.md",
+        "Get Started" => "get_started.md",
+        "Automatic Differentiation Backends" => "autodiff.md",
+        "Detailed API" => "api.md",
+        "Interfaces" => "interfaces.md",
+        "News" => "news.md",
+        "Change Log" => "changelog.md",
+        "References" => "references.md",
+    ],
+)
+
+deploydocs(; repo="github.com/TuringLang/AdvancedHMC.jl.git", push_preview=true)

docs/src/api.md

@@ -1,27 +1,79 @@
-# AdvancedHMC.jl
+# Detailed API for AdvancedHMC.jl
 
-Documentation for AdvancedHMC.jl
+An important design goal of AdvancedHMC.jl is modularity; we would like to support algorithmic research on HMC.
+This modularity means that different HMC variants can be easily constructed by composing various components, such as preconditioning metric (i.e., mass matrix), leapfrog integrators, trajectories (static or dynamic), adaption schemes, etc. In this section, we will explain the detailed usage of different modules in AdancedHMC.jl to provide a comprehensive udnerstanding of how AdvancedHMC.jl can achieve both modularity and efficiency. The section highlights the key components of AdvancedHMC.jl, with a complete documentation provided at the end.
 
-```@contents
-```
+### [Hamiltonian mass matrix (`metric`)](@id hamiltonian_mm)
 
-## Types
+  - Unit metric: `UnitEuclideanMetric(dim)`
+  - Diagonal metric: `DiagEuclideanMetric(dim)`
+  - Dense metric: `DenseEuclideanMetric(dim)`
 
-```@docs
-ClassicNoUTurn
-HMCSampler
-HMC
-NUTS
-HMCDA
-```
+where `dim` is the dimensionality of the sampling space.
+
+### [Integrator (`integrator`)](@id integrator)
+
+  - Ordinary leapfrog integrator: `Leapfrog(ϵ)`
+  - Jittered leapfrog integrator with jitter rate `n`: `JitteredLeapfrog(ϵ, n)`
+  - Tempered leapfrog integrator with tempering rate `a`: `TemperedLeapfrog(ϵ, a)`
+
+where `ϵ` is the step size of leapfrog integration.
+
+### [Kernel (`kernel`)](@id kernel)
+
+  - Static HMC with a fixed number of steps (`n_steps`) from [neal2011mcmc](@Citet): `HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(integrator)))`
+  - HMC with a fixed total trajectory length (`trajectory_length`) from [neal2011mcmc](@Citet): `HMCKernel(Trajectory{EndPointTS}(integrator, FixedIntegrationTime(trajectory_length)))`
+  - Original NUTS with slice sampling from [hoffman2014no](@Citet): `HMCKernel(Trajectory{SliceTS}(integrator, ClassicNoUTurn()))`
+  - Generalised NUTS with slice sampling from [betancourt2017conceptual](@Citet): `HMCKernel(Trajectory{SliceTS}(integrator, GeneralisedNoUTurn()))`
+  - Original NUTS with multinomial sampling from [betancourt2017conceptual](@Citet): `HMCKernel(Trajectory{MultinomialTS}(integrator, ClassicNoUTurn()))`
+  - Generalised NUTS with multinomial sampling from [betancourt2017conceptual](@Citet): `HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))`
 
-## Functions
+### Adaptor (`adaptor`)
 
-```@docs
-sample
+  - Adapt the mass matrix `metric` of the Hamiltonian dynamics: `mma = MassMatrixAdaptor(metric)`
+
+      + This is lowered to `UnitMassMatrix`, `WelfordVar` or `WelfordCov` based on the type of the mass matrix `metric`
+
+  - Adapt the step size of the leapfrog integrator `integrator`: `ssa = StepSizeAdaptor(δ, integrator)`
+
+      + It uses Nesterov's dual averaging with `δ` as the target acceptance rate.
+  - Combine the two above *naively*: `NaiveHMCAdaptor(mma, ssa)`
+  - Combine the first two using Stan's windowed adaptation: `StanHMCAdaptor(mma, ssa)`
+
+## The `sample` functions
+
+```julia
+sample(
+    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
+    h::Hamiltonian,
+    κ::HMCKernel,
+    θ::AbstractVector{<:AbstractFloat},
+    n_samples::Int;
+    adaptor::AbstractAdaptor=NoAdaptation(),
+    n_adapts::Int=min(div(n_samples, 10), 1_000),
+    drop_warmup=false,
+    verbose::Bool=true,
+    progress::Bool=false,
+)
 ```
 
-## More types
+Draw `n_samples` samples using the kernel `κ` under the Hamiltonian system `h`
+
+  - The randomness is controlled by `rng`.
+
+      + If `rng` is not provided, the default random number generator (`Random.default_rng()`) will be used.
+
+  - The initial point is given by `θ`.
+  - The adaptor is set by `adaptor`, for which the default is no adaptation.
+
+      + It will perform `n_adapts` steps of adaptation, for which the default is `1_000` or 10% of `n_samples`, whichever is lower.
+  - `drop_warmup` specifies whether to drop samples.
+  - `verbose` controls the verbosity.
+  - `progress` controls whether to show the progress meter or not.
+
+Note that the function signature of the `sample` function exported by `AdvancedHMC.jl` differs from the [`sample`](https://turinglang.org/dev/docs/using-turing/guide#modelling-syntax-explained) function used by `Turing.jl`. We refer to the documentation of `Turing.jl` for more details on the latter.
+
+## Full documentation of APIs in AdvancedHMC.jl
 
 ```@autodocs; canonical=false
 Modules = [AdvancedHMC, AdvancedHMC.Adaptation]

docs/src/autodiff.md

@@ -0,0 +1,5 @@
+# Gradient in AdvancedHMC.jl
+
+AdvancedHMC.jl supports automatic differentiation using [`LogDensityProblemsAD`](https://github.com/tpapp/LogDensityProblemsAD.jl) across various AD backends and allows user-specified gradients. While the default AD backend for AdvancedHMC.jl is ForwardDiff.jl, we can seamlessly change to other backend like Mooncake.jl using various syntax like `Hamiltonian(metric, ℓπ, AutoZygote())`. Different AD backend can also be pluged in using `Hamiltonian(metric, ℓπ, Zygote)`, `Hamiltonian(metric, ℓπ, Val(:Zygote))` but we recommend using ADTypes since that would allow you to have more freedom for specifying the AD backend.
+
+In order to use user-specified gradients, please replace ForwardDiff.jl with `ℓπ_grad` in the `Hamiltonian` constructor as `Hamiltonian(metric, ℓπ, ℓπ_grad)`, where the gradient function `ℓπ_grad` should return a tuple containing both the log-posterior and its gradient, for example `ℓπ_grad(x) = (log_posterior, grad)`.

docs/src/changelog.md

@@ -0,0 +1,19 @@
+**CHANGELOG**
+
+  - [v0.5.0] **Breaking!** Convenience constructors for common samplers changed to:
+
+      + `HMC(leapfrog_stepsize::Real, n_leapfrog::Int)`
+      + `NUTS(target_acceptance::Real)`
+      + `HMCDA(target_acceptance::Real, integration_time::Real)`
+
+  - [v0.2.22] Three functions are renamed.
+
+      + `Preconditioner(metric::AbstractMetric)` -> `MassMatrixAdaptor(metric)` and
+      + `NesterovDualAveraging(δ, integrator::AbstractIntegrator)` -> `StepSizeAdaptor(δ, integrator)`
+      + `find_good_eps` -> `find_good_stepsize`
+  - [v0.2.15] `n_adapts` is no longer needed to construct `StanHMCAdaptor`; the old constructor is deprecated.
+  - [v0.2.8] Two Hamiltonian trajectory sampling methods are renamed to avoid a name clash with Distributions.
+
+      + `Multinomial` -> `MultinomialTS`
+      + `Slice` -> `SliceTS`
+  - [v0.2.0] The gradient function passed to `Hamiltonian` is supposed to return a value-gradient tuple now.

docs/src/get_started.md

@@ -0,0 +1,203 @@
+# Sampling from a multivariate Gaussian using NUTS
+
+In this section, we demonstrate a minimal example of sampling from a multivariate Gaussian (10-dimensional) using the No U-Turn Sampler (NUTS). Below we describe the major components of the Hamiltonian system which are essential to sample using this approach:
+
+  - **Metric**: In many sampling problems the sample space is associated with a metric that allows us to measure the distance between any two points, and other similar quantities. In the example in this section, we use a special metric called the **Euclidean Metric**, represented with a `D × D` matrix from which we can compute distances.[^1]
+
+  - **Leapfrog integration**: Leapfrog integration is a second-order numerical method for integrating differential equations (In this case they are equations of motion for the relative position of one particle with respect to the other). The order of this integration signifies its rate of convergence. Any algorithm with a finite time step size will have numerical errors, and the order is related to this error. For a second-order algorithm, this error scales as the second power of the time step, hence the name. High-order integrators are usually complex to code and have a limited region of convergence; thus they do not allow arbitrarily large time steps. A second-order integrator is suitable for our purpose. Hence we opt for the leapfrog integrator. It is called `leapfrog` due to the ways this algorithm is written, where the positions and velocities of particles "leap over" each other.[^2]
+  - **Kernel for trajectories (static or dynamic)**: Different kernels, which may be static or dynamic, can be used. At each iteration of any variant of the HMC algorithm, there are two main steps - the first step changes the momentum and the second step may change both the position and the momentum of a particle.[^3]
+
+```julia
+using AdvancedHMC, ForwardDiff
+using LogDensityProblems
+using LinearAlgebra
+
+# Define the target distribution using the `LogDensityProblem` interface
+struct LogTargetDensity
+    dim::Int
+end
+LogDensityProblems.logdensity(p::LogTargetDensity, θ) = -sum(abs2, θ) / 2  # standard multivariate normal
+LogDensityProblems.dimension(p::LogTargetDensity) = p.dim
+function LogDensityProblems.capabilities(::Type{LogTargetDensity})
+    return LogDensityProblems.LogDensityOrder{0}()
+end
+
+# Choose parameter dimensionality and initial parameter value
+D = 10;
+initial_θ = rand(D);
+ℓπ = LogTargetDensity(D)
+
+# Set the number of samples to draw and warmup iterations
+n_samples, n_adapts = 2_000, 1_000
+
+# Define a Hamiltonian system
+metric = DiagEuclideanMetric(D)
+hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)
+
+# Define a leapfrog solver, with the initial step size chosen heuristically
+initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
+integrator = Leapfrog(initial_ϵ)
+
+# Define an HMC sampler with the following components
+#   - multinomial sampling scheme,
+#   - generalised No-U-Turn criteria, and
+#   - windowed adaption for step-size and diagonal mass matrix
+kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
+
+# Run the sampler to draw samples from the specified Gaussian, where
+#   - `samples` will store the samples
+#   - `stats` will store diagnostic statistics for each sample
+samples, stats = sample(
+    hamiltonian, kernel, initial_θ, n_samples, adaptor, n_adapts; progress=true
+)
+```
+
+## Parallel Sampling
+
+AdvancedHMC enables parallel sampling (either distributed or multi-thread) via Julia's [parallel computing functions](https://docs.julialang.org/en/v1/manual/parallel-computing/).
+It also supports vectorized sampling for static HMC.
+
+The below example utilizes the `@threads` macro to sample 4 chains across 4 threads.
+
+```julia
+# Ensure that Julia was launched with an appropriate number of threads
+println(Threads.nthreads())
+
+# Number of chains to sample
+nchains = 4
+
+# Cache to store the chains
+chains = Vector{Any}(undef, nchains)
+
+# The `samples` from each parallel chain is stored in the `chains` vector 
+# Adjust the `verbose` flag as per need
+Threads.@threads for i in 1:nchains
+    samples, stats = sample(
+        hamiltonian, kernel, initial_θ, n_samples, adaptor, n_adapts; verbose=false
+    )
+    chains[i] = samples
+end
+```
+
+## Using the `AbstractMCMC` Interface
+
+Users can also use the `AbstractMCMC` interface to sample, which is also used in Turing.jl.
+In order to show how this is done let us start from our previous example where we defined a `LogTargetDensity`, `ℓπ`.
+
+```julia
+using AbstractMCMC, LogDensityProblemsAD
+
+# Wrap the previous LogTargetDensity as LogDensityModel 
+# where ℓπ::LogTargetDensity
+model = AdvancedHMC.LogDensityModel(LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓπ))
+
+# Wrap the previous sampler as a HMCSampler <: AbstractMCMC.AbstractSampler
+D = 10;
+initial_θ = rand(D);
+n_samples, n_adapts, δ = 1_000, 2_000, 0.8
+sampler = HMCSampler(kernel, metric, adaptor)
+
+# Now sample
+samples = AbstractMCMC.sample(
+    model, sampler, n_adapts + n_samples; n_adapts=n_adapts, initial_params=initial_θ
+)
+```
+
+## Convenience Constructors
+
+In the previous examples, we built the sampler by manually specifying the integrator, metric, kernel, and adaptor to build our own sampler. However, in many cases, users might want to initialize a standard NUTS sampler. In such cases having to define each of these aspects manually is tedious and error-prone. For these reasons `AdvancedHMC` also provides users with a series of convenience constructors for standard samplers. We will now show how to use them.
+
+  - HMC:
+
+    ```julia
+    # HMC Sampler
+    # step size, number of leapfrog steps 
+    n_leapfrog, ϵ = 25, 0.1
+    hmc = HMC(ϵ, n_leapfrog)
+    ```
+
+    is equivalent to:
+
+    ```julia
+    metric = DiagEuclideanMetric(D)
+    hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)
+    integrator = Leapfrog(0.1)
+    kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(n_leapfrog)))
+    adaptor = NoAdaptation()
+    hmc = HMCSampler(kernel, metric, adaptor)
+    ```
+
+  - NUTS:
+
+    ```julia
+    # NUTS Sampler
+    # adaptation steps, target acceptance probability,
+    δ = 0.8
+    nuts = NUTS(δ)
+    ```
+
+    is equivalent to:
+
+    ```julia
+    metric = DiagEuclideanMetric(D)
+    hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)
+    initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
+    integrator = Leapfrog(initial_ϵ)
+    kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+    adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(δ, integrator))
+    nuts = HMCSampler(kernel, metric, adaptor)
+    ```
+  - HMCDA:
+
+    ```julia
+    #HMCDA (dual averaging)
+    # adaptation steps, target acceptance probability, target trajectory length 
+    δ, λ = 0.8, 1.0
+    hmcda = HMCDA(δ, λ)
+    ```
+
+    is equivalent to:
+
+    ```julia
+    metric = DiagEuclideanMetric(D)
+    hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)
+    initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
+    integrator = Leapfrog(initial_ϵ)
+    kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedIntegrationTime(λ)))
+    adaptor = StepSizeAdaptor(δ, initial_ϵ)
+    hmcda = HMCSampler(kernel, metric, adaptor)
+    ```
+
+Moreover, there's some flexibility in how these samplers can be initialized.
+For example, a user can initialize a NUTS (HMC and HMCDA) sampler with their own metrics and integrators.
+This can be done as follows:
+
+```julia
+nuts = NUTS(δ; metric=:diagonal) #metric = DiagEuclideanMetric(D) (Default!)
+nuts = NUTS(δ; metric=:unit)     #metric = UnitEuclideanMetric(D)
+nuts = NUTS(δ; metric=:dense)    #metric = DenseEuclideanMetric(D)
+# Provide your own AbstractMetric
+metric = DiagEuclideanMetric(10)
+nuts = NUTS(δ; metric=metric)
+
+nuts = NUTS(δ; integrator=:leapfrog)         #integrator = Leapfrog(ϵ) (Default!)
+nuts = NUTS(δ; integrator=:jitteredleapfrog) #integrator = JitteredLeapfrog(ϵ, 0.1ϵ)
+nuts = NUTS(δ; integrator=:temperedleapfrog) #integrator = TemperedLeapfrog(ϵ, 1.0)
+
+# Provide your own AbstractIntegrator
+integrator = JitteredLeapfrog(0.1, 0.2)
+nuts = NUTS(δ; integrator=integrator)
+```
+
+## GPU Sampling with CUDA
+
+There is experimental support for running static HMC on the GPU using CUDA.
+To do so, the user needs to have [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) installed, ensure the logdensity of the `Hamiltonian` can be executed on the GPU and that the initial points are a `CuArray`.
+A small working example can be found at `test/cuda.jl`.
+
+## Footnotes
+
+[^1]: The Euclidean metric is also known as the mass matrix in the physical perspective. See [Hamiltonian mass matrix](@ref hamiltonian-mm) for available metrics.
+[^2]: About the leapfrog integration scheme: Suppose ${\bf x}$ and ${\bf v}$ are the position and velocity of an individual particle respectively; $i$ and $i+1$ are the indices for time values $t_i$ and $t_{i+1}$ respectively; $dt = t_{i+1} - t_i$ is the time step size (constant and regularly spaced intervals), and ${\bf a}$ is the acceleration induced on a particle by the forces of all other particles. Furthermore, suppose positions are defined at times $t_i, t_{i+1}, t_{i+2}, \dots $, spaced at constant intervals $dt$, the velocities are defined at halfway times in between, denoted by $t_{i-1/2}, t_{i+1/2}, t_{i+3/2}, \dots $, where $t_{i+1} - t_{i + 1/2} = t_{i + 1/2} - t_i = dt / 2$, and the accelerations ${\bf a}$ are defined only on integer times, just like the positions. Then the leapfrog integration scheme is given as: $x_{i} = x_{i-1} + v_{i-1/2} dt; \quad v_{i+1/2} = v_{i-1/2} + a_i dt$. For available integrators refer to [Integrator](@ref integrator).
+[^3]: On kernels: In the classical HMC approach, during the first step, new values for the momentum variables are randomly drawn from their Gaussian distribution, independently of the current values of the position variables. A Metropolis update is performed during the second step, using Hamiltonian dynamics to provide a new state. For available kernels refer to [Kernel](@ref kernel).