diff --git a/index.html b/index.html index 3ac2596..6a5afc3 100644 --- a/index.html +++ b/index.html @@ -1,3 +1,2 @@ - diff --git a/previews/PR144/.documenter-siteinfo.json b/previews/PR144/.documenter-siteinfo.json index 0917d01..7a4b308 100644 --- a/previews/PR144/.documenter-siteinfo.json +++ b/previews/PR144/.documenter-siteinfo.json @@ -1 +1 @@ -{"documenter":{"julia_version":"1.10.5","generation_timestamp":"2024-09-22T20:02:53","documenter_version":"1.7.0"}} \ No newline at end of file +{"documenter":{"julia_version":"1.10.5","generation_timestamp":"2024-09-27T16:04:26","documenter_version":"1.7.0"}} \ No newline at end of file diff --git a/previews/PR144/api/index.html b/previews/PR144/api/index.html index 55512d2..89fc5f2 100644 --- a/previews/PR144/api/index.html +++ b/previews/PR144/api/index.html @@ -1,475 +1,17 @@ -API · AbstractMCMC - - - - - -
GitHub

API

AbstractMCMC defines an interface for sampling Markov chains.

Model

AbstractMCMC.LogDensityModel — Type
LogDensityModel <: AbstractMCMC.AbstractModel

Wrapper around something that implements the LogDensityProblem.jl interface.

Note that this does not implement the LogDensityProblems.jl interface itself, but it simply useful for indicating to the sample and other AbstractMCMC methods that the wrapped object implements the LogDensityProblems.jl interface.

Fields

  • logdensity: The object that implements the LogDensityProblems.jl interface.
source

Sampler

AbstractMCMC.AbstractSampler — Type
AbstractSampler

The AbstractSampler type is intended to be inherited from when implementing a custom sampler. Any persistent state information should be saved in a subtype of AbstractSampler.

When defining a new sampler, you should also overload the function transition_type, which tells the sample function what type of parameter it should expect to receive.

source

Sampling a single chain

StatsBase.sample — Method
sample(
+API · AbstractMCMC
GitHub
API
AbstractMCMC defines an interface for sampling Markov chains.
Model
AbstractMCMC.LogDensityModel — Type
LogDensityModel <: AbstractMCMC.AbstractModel
Wrapper around something that implements the LogDensityProblem.jl interface.
Note that this does not implement the LogDensityProblems.jl interface itself, but it simply useful for indicating to the sample and other AbstractMCMC methods that the wrapped object implements the LogDensityProblems.jl interface.
Fields
  • logdensity: The object that implements the LogDensityProblems.jl interface.
source
Sampler
AbstractMCMC.AbstractSampler — Type
AbstractSampler
The AbstractSampler type is intended to be inherited from when implementing a custom sampler. Any persistent state information should be saved in a subtype of AbstractSampler.
When defining a new sampler, you should also overload the function transition_type, which tells the sample function what type of parameter it should expect to receive.
source
Sampling a single chain
StatsBase.sample — Method
sample(
     rng::Random.AbatractRNG=Random.default_rng(),
     model::AbstractModel,
     sampler::AbstractSampler,
     N_or_isdone;
     kwargs...,
-)
Sample from the model with the Markov chain Monte Carlo sampler and return the samples.
If N_or_isdone is an Integer, exactly N_or_isdone samples are returned.
Otherwise, sampling is performed until a convergence criterion N_or_isdone returns true. The convergence criterion has to be a function with the signature
isdone(rng, model, sampler, samples, state, iteration; kwargs...)
where state and iteration are the current state and iteration of the sampler, respectively. It should return true when sampling should end, and false otherwise.
source
StatsBase.sample — Method
sample(
+)
Sample from the model with the Markov chain Monte Carlo sampler and return the samples.
If N_or_isdone is an Integer, exactly N_or_isdone samples are returned.
Otherwise, sampling is performed until a convergence criterion N_or_isdone returns true. The convergence criterion has to be a function with the signature
isdone(rng, model, sampler, samples, state, iteration; kwargs...)
where state and iteration are the current state and iteration of the sampler, respectively. It should return true when sampling should end, and false otherwise.
source
StatsBase.sample — Method
sample(
     rng::Random.AbstractRNG=Random.default_rng(),
     logdensity,
     sampler::AbstractSampler,
     N_or_isdone;
     kwargs...,
-)
Wrap the logdensity function in a LogDensityModel, and call sample with the resulting model instead of logdensity.
The logdensity function has to support the LogDensityProblems.jl interface.
source
Iterator
Iterator
AbstractMCMC.steps — Method
steps(
     rng::Random.AbstractRNG=Random.default_rng(),
     model::AbstractModel,
     sampler::AbstractSampler;
@@ -486,12 +28,12 @@
 julia> iterator = steps(MyModel(), MySampler());
 
 julia> collect(Iterators.take(iterator, 10)) == zeros(10)
-true
source
AbstractMCMC.steps — Method
steps(
     rng::Random.AbstractRNG=Random.default_rng(),
     logdensity,
     sampler::AbstractSampler;
     kwargs...,
-)
Wrap the logdensity function in a LogDensityModel, and call steps with the resulting model instead of logdensity.
The logdensity function has to support the LogDensityProblems.jl interface.
source
Transducer
Transducer
AbstractMCMC.Sample — Method
Sample(
     rng::Random.AbstractRNG=Random.default_rng(),
     model::AbstractModel,
     sampler::AbstractSampler;
@@ -508,12 +50,12 @@
 julia> transducer = Sample(MyModel(), MySampler());
 
 julia> collect(transducer(1:10)) == zeros(10)
-true
source
AbstractMCMC.Sample — Method
Sample(
     rng::Random.AbstractRNG=Random.default_rng(),
     logdensity,
     sampler::AbstractSampler;
     kwargs...,
-)
Wrap the logdensity function in a LogDensityModel, and call Sample with the resulting model instead of logdensity.
The logdensity function has to support the LogDensityProblems.jl interface.
source
Sampling multiple chains in parallel
Sampling multiple chains in parallel
StatsBase.sample — Method
sample(
     rng::Random.AbstractRNG=Random.default_rng(),
     model::AbstractModel,
     sampler::AbstractSampler,
@@ -521,7 +63,7 @@
     N::Integer,
     nchains::Integer;
     kwargs...,
-)
Sample nchains Monte Carlo Markov chains from the model with the sampler in parallel using the parallel algorithm, and combine them into a single chain.
source
StatsBase.sample — Method
sample(
+)
Sample nchains Monte Carlo Markov chains from the model with the sampler in parallel using the parallel algorithm, and combine them into a single chain.
source
StatsBase.sample — Method
sample(
     rng::Random.AbstractRNG=Random.default_rng(),
     logdensity,
     sampler::AbstractSampler,
@@ -529,5 +71,4 @@
     N::Integer,
     nchains::Integer;
     kwargs...,
-)
Wrap the logdensity function in a LogDensityModel, and call sample with the resulting model instead of logdensity.
The logdensity function has to support the LogDensityProblems.jl interface.
source
Two algorithms are provided for parallel sampling with multiple threads and multiple processes, and one allows for the user to sample multiple chains in serial (no parallelization):
Common keyword arguments
Common keyword arguments for regular and parallel sampling are:
  • progress (default: AbstractMCMC.PROGRESS[] which is true initially):  toggles progress logging
  • chain_type (default: Any): determines the type of the returned chain
  • callback (default: nothing): if callback !== nothing, then callback(rng, model, sampler, sample, state, iteration) is called after every sampling step, where sample is the most recent sample of the Markov chain and state and iteration are the current state and iteration of the sampler
  • discard_initial (default: 0): number of initial samples that are discarded
  • thinning (default: 1): factor by which to thin samples.
  • initial_state (default: nothing): if initial_state !== nothing, the first call to AbstractMCMC.step is passed initial_state as the state argument.
Info
The common keyword arguments progress, chain_type, and callback are not supported by the iterator AbstractMCMC.steps and the transducer AbstractMCMC.Sample.
There is no "official" way for providing initial parameter values yet. However, multiple packages such as EllipticalSliceSampling.jl and AdvancedMH.jl support an initial_params keyword argument for setting the initial values when sampling a single chain. To ensure that sampling multiple chains "just works" when sampling of a single chain is implemented, we decided to support initial_params in the default implementations of the ensemble methods:
  • initial_params (default: nothing): if initial_params isa AbstractArray, then the ith element of initial_params is used as initial parameters of the ith chain. If one wants to use the same initial parameters x for every chain, one can specify e.g. initial_params = FillArrays.Fill(x, N).
Progress logging can be enabled and disabled globally with AbstractMCMC.setprogress!(progress).
AbstractMCMC.setprogress! — Function
setprogress!(progress::Bool; silent::Bool=false)
Enable progress logging globally if progress is true, and disable it otherwise.  Optionally disable informational message if silent is true.
source
Chains
The chain_type keyword argument allows to set the type of the returned chain. A common choice is to return chains of type Chains from MCMCChains.jl.
AbstractMCMC defines the abstract type AbstractChains for Markov chains.
AbstractMCMC.AbstractChains — Type
AbstractChains
AbstractChains is an abstract type for an object that stores parameter samples generated through a MCMC process.
source
For chains of this type, AbstractMCMC defines the following two methods.
AbstractMCMC.chainscat — Function
chainscat(c::AbstractChains...)
Concatenate multiple chains.
By default, the chains are concatenated along the third dimension by calling cat(c...; dims=3).
source
AbstractMCMC.chainsstack — Function
chainsstack(c::AbstractVector)
Stack chains in c.
By default, the vector of chains is returned unmodified. If eltype(c) <: AbstractChains, then reduce(chainscat, c) is called.
source

-
+)

Wrap the logdensity function in a LogDensityModel, and call sample with the resulting model instead of logdensity.

The logdensity function has to support the LogDensityProblems.jl interface.

source

Two algorithms are provided for parallel sampling with multiple threads and multiple processes, and one allows for the user to sample multiple chains in serial (no parallelization):

Common keyword arguments

Common keyword arguments for regular and parallel sampling are:

  • progress (default: AbstractMCMC.PROGRESS[] which is true initially): toggles progress logging
  • chain_type (default: Any): determines the type of the returned chain
  • callback (default: nothing): if callback !== nothing, then callback(rng, model, sampler, sample, state, iteration) is called after every sampling step, where sample is the most recent sample of the Markov chain and state and iteration are the current state and iteration of the sampler
  • discard_initial (default: 0): number of initial samples that are discarded
  • thinning (default: 1): factor by which to thin samples.
  • initial_state (default: nothing): if initial_state !== nothing, the first call to AbstractMCMC.step is passed initial_state as the state argument.
Info

The common keyword arguments progress, chain_type, and callback are not supported by the iterator AbstractMCMC.steps and the transducer AbstractMCMC.Sample.

There is no "official" way for providing initial parameter values yet. However, multiple packages such as EllipticalSliceSampling.jl and AdvancedMH.jl support an initial_params keyword argument for setting the initial values when sampling a single chain. To ensure that sampling multiple chains "just works" when sampling of a single chain is implemented, we decided to support initial_params in the default implementations of the ensemble methods:

  • initial_params (default: nothing): if initial_params isa AbstractArray, then the ith element of initial_params is used as initial parameters of the ith chain. If one wants to use the same initial parameters x for every chain, one can specify e.g. initial_params = FillArrays.Fill(x, N).

Progress logging can be enabled and disabled globally with AbstractMCMC.setprogress!(progress).

AbstractMCMC.setprogress! — Function
setprogress!(progress::Bool; silent::Bool=false)

Enable progress logging globally if progress is true, and disable it otherwise. Optionally disable informational message if silent is true.

source

Chains

The chain_type keyword argument allows to set the type of the returned chain. A common choice is to return chains of type Chains from MCMCChains.jl.

AbstractMCMC defines the abstract type AbstractChains for Markov chains.

AbstractMCMC.AbstractChains — Type
AbstractChains

AbstractChains is an abstract type for an object that stores parameter samples generated through a MCMC process.

source

For chains of this type, AbstractMCMC defines the following two methods.

AbstractMCMC.chainscat — Function
chainscat(c::AbstractChains...)

Concatenate multiple chains.

By default, the chains are concatenated along the third dimension by calling cat(c...; dims=3).

source
AbstractMCMC.chainsstack — Function
chainsstack(c::AbstractVector)

Stack chains in c.

By default, the vector of chains is returned unmodified. If eltype(c) <: AbstractChains, then reduce(chainscat, c) is called.

source
diff --git a/previews/PR144/design/index.html b/previews/PR144/design/index.html index 2393ad4..e64efba 100644 --- a/previews/PR144/design/index.html +++ b/previews/PR144/design/index.html @@ -1,463 +1,5 @@ -Design · AbstractMCMC - - - - - -
GitHub

Design

This page explains the default implementations and design choices of AbstractMCMC. It is not intended for users but for developers that want to implement the AbstractMCMC interface for Markov chain Monte Carlo sampling. The user-facing API is explained in API.

Overview

AbstractMCMC provides a default implementation of the user-facing interface described in API. You can completely neglect these and define your own implementation of the interface. However, as described below, in most use cases the default implementation allows you to obtain support of parallel sampling, progress logging, callbacks, iterators, and transducers for free by just defining the sampling step of your inference algorithm, drastically reducing the amount of code you have to write. In general, the docstrings of the functions described below might be helpful if you intend to make use of the default implementations.

Basic structure

The simplified structure for regular sampling (the actual implementation contains some additional error checks and support for progress logging and callbacks) is

StatsBase.sample(
+Design · AbstractMCMC
GitHub
Design
This page explains the default implementations and design choices of AbstractMCMC. It is not intended for users but for developers that want to implement the AbstractMCMC interface for Markov chain Monte Carlo sampling. The user-facing API is explained in API.
Overview
AbstractMCMC provides a default implementation of the user-facing interface described in API. You can completely neglect these and define your own implementation of the interface. However, as described below, in most use cases the default implementation allows you to obtain support of parallel sampling, progress logging, callbacks, iterators, and transducers for free by just defining the sampling step of your inference algorithm, drastically reducing the amount of code you have to write. In general, the docstrings of the functions described below might be helpful if you intend to make use of the default implementations.
Basic structure
The simplified structure for regular sampling (the actual implementation contains some additional error checks and support for progress logging and callbacks) is
StatsBase.sample(
     rng::Random.AbstractRNG,
     model::AbstractMCMC.AbstractModel,
     sampler::AbstractMCMC.AbstractSampler,
@@ -482,5 +24,4 @@
     end
 
     return AbstractMCMC.bundle_samples(samples, model, sampler, state, chain_type; kwargs...)
-end
All other default implementations make use of the same structure and in particular call the same methods.
Sampling step
The only method for which no default implementation is provided (and hence which downstream packages have to implement) is AbstractMCMC.step. It defines the sampling step of the inference method.
AbstractMCMC.step — Function
step(rng, model, sampler[, state; kwargs...])
Return a 2-tuple of the next sample and the next state of the MCMC sampler for model.
Samples describe the results of a single step of the sampler. As an example, a sample might include a vector of parameters sampled from a prior distribution.
When sampling using sample, every step call after the first has access to the current state of the sampler.
source
Collecting samples
Note
This section does not apply to the iterator and transducer interface.
After the initial sample is obtained, the default implementations for regular and parallel sampling (not for the iterator and the transducer since it is not needed there) create a container for all samples (the initial one and all subsequent samples) using AbstractMCMC.samples.
AbstractMCMC.samples — Function
samples(sample, model, sampler[, N; kwargs...])
Generate a container for the samples of the MCMC sampler for the model, whose first sample is sample.
The method can be called with and without a predefined number N of samples.
source
In each step, the sample is saved in the container by AbstractMCMC.save!!. The notation !! follows the convention of the package BangBang.jl which is used in the default implementation of AbstractMCMC.save!!. It indicates that the sample is pushed to the container but a "widening" fallback is used if the container type does not allow to save the sample. Therefore AbstractMCMC.save!! always has to return the container.
AbstractMCMC.save!! — Function
save!!(samples, sample, iteration, model, sampler[, N; kwargs...])
Save the sample of the MCMC sampler at the current iteration in the container of samples.
The function can be called with and without a predefined number N of samples. By default, AbstractMCMC uses push!! from the Julia package BangBang to append to the container, and widen its type if needed.
source
For most use cases the default implementation of AbstractMCMC.samples and AbstractMCMC.save!! should work out of the box and hence need not to be overloaded in downstream code.
Creating chains
Note
This section does not apply to the iterator and transducer interface.
At the end of the sampling procedure for regular and paralle sampling we transform the collection of samples to the desired output type by calling AbstractMCMC.bundle_samples.
AbstractMCMC.bundle_samples — Function
bundle_samples(samples, model, sampler, state, chain_type[; kwargs...])
Bundle all samples that were sampled from the model with the given sampler in a chain.
The final state of the sampler can be included in the chain. The type of the chain can be specified with the chain_type argument.
By default, this method returns samples.
source
The default implementation should be fine in most use cases, but downstream packages could, e.g., save the final state of the sampler as well if they overload AbstractMCMC.bundle_samples.

-
+end

All other default implementations make use of the same structure and in particular call the same methods.

Sampling step

The only method for which no default implementation is provided (and hence which downstream packages have to implement) is AbstractMCMC.step. It defines the sampling step of the inference method.

AbstractMCMC.step — Function
step(rng, model, sampler[, state; kwargs...])

Return a 2-tuple of the next sample and the next state of the MCMC sampler for model.

Samples describe the results of a single step of the sampler. As an example, a sample might include a vector of parameters sampled from a prior distribution.

When sampling using sample, every step call after the first has access to the current state of the sampler.

source

Collecting samples

Note

This section does not apply to the iterator and transducer interface.

After the initial sample is obtained, the default implementations for regular and parallel sampling (not for the iterator and the transducer since it is not needed there) create a container for all samples (the initial one and all subsequent samples) using AbstractMCMC.samples.

AbstractMCMC.samples — Function
samples(sample, model, sampler[, N; kwargs...])

Generate a container for the samples of the MCMC sampler for the model, whose first sample is sample.

The method can be called with and without a predefined number N of samples.

source

In each step, the sample is saved in the container by AbstractMCMC.save!!. The notation !! follows the convention of the package BangBang.jl which is used in the default implementation of AbstractMCMC.save!!. It indicates that the sample is pushed to the container but a "widening" fallback is used if the container type does not allow to save the sample. Therefore AbstractMCMC.save!! always has to return the container.

AbstractMCMC.save!! — Function
save!!(samples, sample, iteration, model, sampler[, N; kwargs...])

Save the sample of the MCMC sampler at the current iteration in the container of samples.

The function can be called with and without a predefined number N of samples. By default, AbstractMCMC uses push!! from the Julia package BangBang to append to the container, and widen its type if needed.

source

For most use cases the default implementation of AbstractMCMC.samples and AbstractMCMC.save!! should work out of the box and hence need not to be overloaded in downstream code.

Creating chains

Note

This section does not apply to the iterator and transducer interface.

At the end of the sampling procedure for regular and paralle sampling we transform the collection of samples to the desired output type by calling AbstractMCMC.bundle_samples.

AbstractMCMC.bundle_samples — Function
bundle_samples(samples, model, sampler, state, chain_type[; kwargs...])

Bundle all samples that were sampled from the model with the given sampler in a chain.

The final state of the sampler can be included in the chain. The type of the chain can be specified with the chain_type argument.

By default, this method returns samples.

source

The default implementation should be fine in most use cases, but downstream packages could, e.g., save the final state of the sampler as well if they overload AbstractMCMC.bundle_samples.

diff --git a/previews/PR144/gibbs/index.html b/previews/PR144/gibbs/index.html index aa54f54..340b5e5 100644 --- a/previews/PR144/gibbs/index.html +++ b/previews/PR144/gibbs/index.html @@ -1,463 +1,5 @@ -state Interface Functions · AbstractMCMC - - - - - -
GitHub

state Interface Functions

We encourage sampler packages to implement the following interface functions for the state type(s) they maintain:

LogDensityProblems.logdensity(logdensity_model::AbstractMCMC.LogDensityModel, state::MHState; recompute_logp=true)

This function takes the logdensity model and the state, and returns the log probability of the state. If recompute_logp is true, it should recompute the log probability of the state. Otherwise, it should use the log probability stored in the state.

Base.vec(state)

This function takes the state and returns a vector of the parameter values stored in the state.

(state::StateType)(logp::Float64)

This function takes the state and a log probability value, and updates the state with the new log probability.

These function will provide a minimum interface to interact with the state datatype, which a sampler package doesn't have to expose.

Using the state Interface for block sampling within Gibbs

In this sections, we will demonstrate how a model package may use this state interface to support a Gibbs sampler that can support blocking sampling using different inference algorithms.

We consider a simple hierarchical model with a normal likelihood, with unknown mean and variance parameters.

\[\begin{align} +state Interface Functions · AbstractMCMC

GitHub

state Interface Functions

We encourage sampler packages to implement the following interface functions for the state type(s) they maintain:

LogDensityProblems.logdensity(logdensity_model::AbstractMCMC.LogDensityModel, state::MHState; recompute_logp=true)

This function takes the logdensity model and the state, and returns the log probability of the state. If recompute_logp is true, it should recompute the log probability of the state. Otherwise, it could use the log probability stored in the state.

Base.vec(state)

This function takes the state and returns a vector of the parameter values stored in the state.

(state::StateType)(logp::Float64)

This function takes the state and a log probability value, and returns a new state with the updated log probability.

These functions provide a minimal interface to interact with the state datatype, which a sampler package can optionally implement. The interface facilitates the implementation of "meta-algorithms" that combine different samplers. We will demonstrate how it can be used to implement Gibbs sampling in the following sections.

Using the state Interface for block sampling within Gibbs

In this sections, we will demonstrate how a model package may use this state interface to support a Gibbs sampler that can support blocking sampling using different inference algorithms.

We consider a simple hierarchical model with a normal likelihood, with unknown mean and variance parameters.

\[\begin{align} \mu &\sim \text{Normal}(0, 1) \\ \tau^2 &\sim \text{InverseGamma}(1, 1) \\ x_i &\sim \text{Normal}(\mu, \sqrt{\tau^2}) @@ -521,7 +63,7 @@ function LogDensityProblems.capabilities(::ConditionedHierNormal) return LogDensityProblems.LogDensityOrder{0}() -end

Sampler Packages

To illustrate the AbstractMCMC interface, we will first implement two very simple Metropolis-Hastings samplers: random walk and static proposal.

Although the interface doesn't force the sampler to implement Transition and State types, in practice, it has been the convention to do so.

Here we define some bare minimum types to represent the transitions and states.

struct MHTransition{T}
+end

Implementing A Sampler with AbstractMCMC Interface

To illustrate the AbstractMCMC interface, we will first implement two very simple Metropolis-Hastings samplers: random walk and static proposal.

Although the interface doesn't force the sampler to implement Transition and State types, in practice, it has been the convention to do so.

Here we define some bare minimum types to represent the transitions and states.

struct MHTransition{T}
     params::Vector{T}
 end
 
@@ -633,12 +175,8 @@
         logp_proposal - state.logp + logpdf(sampler.proposal_dist, state.params) -
         logpdf(sampler.proposal_dist, proposal),
     )
-end

At last, we can proceed to implement the Gibbs sampler.

"""
-    Gibbs(sampler_map::NamedTuple)
-
-A Gibbs sampler that allows for block sampling using different inference algorithms for each parameter.
-"""
-struct Gibbs{T<:NamedTuple} <: AbstractMCMC.AbstractSampler
+end

At last, we can proceed to implement a very simple Gibbs sampler.

struct Gibbs{T<:NamedTuple} <: AbstractMCMC.AbstractSampler
+    "Maps variables to their samplers."
     sampler_map::T
 end
 
@@ -663,16 +201,18 @@
 
 Update the trace with the values from the MCMC states of the sub-problems.
 """
-function update_trace(trace::NamedTuple, gibbs_state::GibbsState)
-    for parameter_variable in keys(gibbs_state.mcmc_states)
+function update_trace(
+    trace::NamedTuple{trace_names}, gibbs_state::GibbsState{TraceNT,StateNT,SizeNT}
+) where {trace_names,TraceNT,StateNT,SizeNT}
+    for parameter_variable in fieldnames(StateNT)
         sub_state = gibbs_state.mcmc_states[parameter_variable]
-        sub_state_params = Base.vec(sub_state)
-        unflattened_sub_state_params = unflatten(
-            sub_state_params,
-            NamedTuple{(parameter_variable,)}((
-                gibbs_state.variable_sizes[parameter_variable],
-            )),
+        sub_state_params_values = Base.vec(sub_state)
+        reshaped_sub_state_params_values = reshape(
+            sub_state_params_values, gibbs_state.variable_sizes[parameter_variable]
         )
+        unflattened_sub_state_params = NamedTuple{(parameter_variable,)}((
+            reshaped_sub_state_params_values,
+        ))
         trace = merge(trace, unflattened_sub_state_params)
     end
     return trace
@@ -693,8 +233,7 @@
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     logdensity_model::AbstractMCMC.LogDensityModel,
-    sampler::Gibbs{Tsamplingmap},
-    args...;
+    sampler::Gibbs{Tsamplingmap};
     initial_params::NamedTuple,
     kwargs...,
 ) where {Tsamplingmap}
@@ -710,30 +249,27 @@
         conditioning_variables_values = NamedTuple{Tuple(variables_to_be_conditioned_on)}(
             Tuple([initial_params[g] for g in variables_to_be_conditioned_on])
         )
-        sub_problem_parameters_values = NamedTuple{(parameter_variable,)}((
-            initial_params[parameter_variable],
-        ))
 
         # LogDensityProblems' `logdensity` function expects a single vector of real numbers
         # `Gibbs` stores the parameters as a named tuple, thus we need to flatten the sub_problem_parameters_values
         # and unflatten after the sampling step
-        flattened_sub_problem_parameters_values = flatten(sub_problem_parameters_values)
+        flattened_sub_problem_parameters_values = vec(initial_params[parameter_variable])
 
+        sub_logdensity_model = AbstractMCMC.LogDensityModel(
+            AbstractPPL.condition(
+                logdensity_model.logdensity, conditioning_variables_values
+            ),
+        )
         sub_state = last(
             AbstractMCMC.step(
                 rng,
-                AbstractMCMC.LogDensityModel(
-                    AbstractPPL.condition(
-                        logdensity_model.logdensity, conditioning_variables_values
-                    ),
-                ),
-                sub_sampler,
-                args...;
+                sub_logdensity_model,
+                sub_sampler;
                 initial_params=flattened_sub_problem_parameters_values,
                 kwargs...,
             ),
         )
-        (sub_state, Tuple(size(initial_params[parameter_variable])))
+        (sub_state, size(initial_params[parameter_variable]))
     end
 
     mcmc_states_tuple = first.(results)
@@ -754,11 +290,12 @@
     rng::Random.AbstractRNG,
     logdensity_model::AbstractMCMC.LogDensityModel,
     sampler::Gibbs{Tsamplingmap},
-    gibbs_state::GibbsState,
-    args...;
+    gibbs_state::GibbsState;
     kwargs...,
 ) where {Tsamplingmap}
-    (; trace, mcmc_states, variable_sizes) = gibbs_state
+    trace = gibbs_state.trace
+    mcmc_states = gibbs_state.mcmc_states
+    variable_sizes = gibbs_state.variable_sizes
 
     model_parameter_names = fieldnames(Tsamplingmap)
     mcmc_states = map(model_parameter_names) do parameter_variable
@@ -779,7 +316,7 @@
         sub_state = (sub_state)(logp)
         sub_state = last(
             AbstractMCMC.step(
-                rng, cond_logdensity_model, sub_sampler, sub_state, args...; kwargs...
+                rng, cond_logdensity_model, sub_sampler, sub_state; kwargs...
             ),
         )
         trace = update_trace(trace, gibbs_state)
@@ -788,34 +325,7 @@
     mcmc_states = NamedTuple{Tuple(model_parameter_names)}(mcmc_states)
 
     return GibbsTransition(trace), GibbsState(trace, mcmc_states, variable_sizes)
-end

where we use two utility functions flatten and unflatten to convert between the single vector of real numbers and the named tuple of parameters.

"""
-    flatten(trace::NamedTuple)
-
-Flatten all the values in the trace into a single vector. Variable names information is discarded.
-"""
-function flatten(trace::NamedTuple)
-    return reduce(vcat, vec.(values(trace)))
-end
-
-"""
-    unflatten(vec::AbstractVector, variable_names::Vector{Symbol}, variable_sizes::Vector{Tuple})
-
-Reverse operation of flatten. Reshape the vector into the original arrays using size information.
-"""
-function unflatten(
-    vec::AbstractVector, variable_names_and_sizes::NamedTuple{variable_names}
-) where {variable_names}
-    result = Dict{Symbol,Array}()
-    start_idx = 1
-    for name in variable_names
-        size = variable_names_and_sizes[name]
-        end_idx = start_idx + prod(size) - 1
-        result[name] = reshape(vec[start_idx:end_idx], size...)
-        start_idx = end_idx + 1
-    end
-
-    return NamedTuple{variable_names}(Tuple([result[name] for name in variable_names]))
-end

Some points worth noting:

  1. We are using NamedTuple to store the mapping between variables and samplers. The order will determine the order of the Gibbs sweeps. A limitation is that exactly one sampler for each variable is required, which means it is less flexible than Gibbs in Turing.jl.
  2. For each conditional probability problem, we need to store the sampler states for each variable group and also the values of all the variables from last iteration.
  3. The first step of the Gibbs sampler is to setup the states for each conditional probability problem.
  4. In the following steps of the Gibbs sampler, it will do a sweep over all the conditional probability problems, and update the sampler states for each problem. In each step of the sweep, it will do the following:
    • condition on the values of all variables that are not in the current group
    • recompute the log probability of the current state, because the values of the variables that are not in the current group may have changed
    • perform a step of the sampler for the conditional probability problem, and update the sampler state
    • update the vi with the new values from the sampler state

The state interface in AbstractMCMC allows the Gibbs sampler to be agnostic of the details of the sampler state, and acquire the values of the parameters from individual sampler states.

Now we can use the Gibbs sampler to sample from the hierarchical normal model.

First we generate some data,

N = 100  # Number of data points
+end

We are using NamedTuple to store the mapping between variables and samplers. The order will determine the order of the Gibbs sweeps. A limitation is that exactly one sampler for each variable is required, which means it is less flexible than Gibbs in Turing.jl.

We uses the AbstractPPL.condition to devide the full model into smaller conditional probability problems. And each conditional probability problem corresponds to a sampler and corresponding state.

The Gibbs sampler has the same interface as other samplers in AbstractMCMC (we don't implement the above state interface for GibbsState to keep it simple, but it can be implemented similarly).

The Gibbs sampler operates in two main phases:

  1. Initialization:

    • Set up initial states for each conditional probability problem.

  2. Iterative Sampling: For each iteration, the sampler performs a sweep over all conditional probability problems:

    a. Condition on other variables:

    • Fix the values of all variables except the current one.

    b. Update current variable:

    • Recompute the log probability of the current state, as other variables may have changed:
      • Use LogDensityProblems.logdensity(cond_logdensity_model, sub_state) to get the new log probability.
      • Update the state with sub_state = sub_state(logp) to incorporate the new log probability.
    • Perform a sampling step for the current conditional probability problem:
      • Use AbstractMCMC.step(rng, cond_logdensity_model, sub_sampler, sub_state; kwargs...) to generate a new state.
    • Update the global trace:
      • Extract parameter values from the new state using Base.vec(new_sub_state).
      • Incorporate these values into the overall Gibbs state trace.

This process allows the Gibbs sampler to iteratively update each variable while conditioning on the others, gradually exploring the joint distribution of all variables.

Now we can use the Gibbs sampler to sample from the hierarchical Normal model.

First we generate some data,

N = 100  # Number of data points
 mu_true = 0.5  # True mean
 tau2_true = 2.0  # True variance
 
@@ -832,5 +342,4 @@
 
 mean(mu_samples)
 mean(tau2_samples)
-(mu_mean, tau2_mean)

the result should looks like:

(4.995812149309413, 1.9372372289677886)

which is close to the true values (5, 2).

- +(mu_mean, tau2_mean)

the result should looks like:

(4.995812149309413, 1.9372372289677886)

which is close to the true values (5, 2).

diff --git a/previews/PR144/index.html b/previews/PR144/index.html index 26a6a2f..9d58424 100644 --- a/previews/PR144/index.html +++ b/previews/PR144/index.html @@ -1,461 +1,2 @@ -Home · AbstractMCMC - - - - - -
GitHub

AbstractMCMC.jl

Abstract types and interfaces for Markov chain Monte Carlo methods.

AbstractMCMC defines an interface for sampling and combining Markov chains. It comes with a default sampling algorithm that provides support of progress bars, parallel sampling (multithreaded and multicore), and user-provided callbacks out of the box. Typically developers only have to define the sampling step of their inference method in an iterator-like fashion to make use of this functionality. Additionally, the package defines an iterator and a transducer for sampling Markov chains based on the interface.

- +Home · AbstractMCMC
GitHub

AbstractMCMC.jl

Abstract types and interfaces for Markov chain Monte Carlo methods.

AbstractMCMC defines an interface for sampling and combining Markov chains. It comes with a default sampling algorithm that provides support of progress bars, parallel sampling (multithreaded and multicore), and user-provided callbacks out of the box. Typically developers only have to define the sampling step of their inference method in an iterator-like fashion to make use of this functionality. Additionally, the package defines an iterator and a transducer for sampling Markov chains based on the interface.

diff --git a/previews/PR144/objects.inv b/previews/PR144/objects.inv index e27e567..14ee8a0 100644 Binary files a/previews/PR144/objects.inv and b/previews/PR144/objects.inv differ diff --git a/previews/PR144/search_index.js b/previews/PR144/search_index.js index bc79e16..e90cf8f 100644 --- a/previews/PR144/search_index.js +++ b/previews/PR144/search_index.js @@ -1,3 +1,3 @@ var documenterSearchIndex = {"docs": -[{"location":"api/#API","page":"API","title":"API","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC defines an interface for sampling Markov chains.","category":"page"},{"location":"api/#Model","page":"API","title":"Model","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.AbstractModel\nAbstractMCMC.LogDensityModel","category":"page"},{"location":"api/#AbstractMCMC.AbstractModel","page":"API","title":"AbstractMCMC.AbstractModel","text":"AbstractModel\n\nAn AbstractModel represents a generic model type that can be used to perform inference.\n\n\n\n\n\n","category":"type"},{"location":"api/#AbstractMCMC.LogDensityModel","page":"API","title":"AbstractMCMC.LogDensityModel","text":"LogDensityModel <: AbstractMCMC.AbstractModel\n\nWrapper around something that implements the LogDensityProblem.jl interface.\n\nNote that this does not implement the LogDensityProblems.jl interface itself, but it simply useful for indicating to the sample and other AbstractMCMC methods that the wrapped object implements the LogDensityProblems.jl interface.\n\nFields\n\nlogdensity: The object that implements the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"type"},{"location":"api/#Sampler","page":"API","title":"Sampler","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.AbstractSampler","category":"page"},{"location":"api/#AbstractMCMC.AbstractSampler","page":"API","title":"AbstractMCMC.AbstractSampler","text":"AbstractSampler\n\nThe AbstractSampler type is intended to be inherited from when implementing a custom sampler. Any persistent state information should be saved in a subtype of AbstractSampler.\n\nWhen defining a new sampler, you should also overload the function transition_type, which tells the sample function what type of parameter it should expect to receive.\n\n\n\n\n\n","category":"type"},{"location":"api/#Sampling-a-single-chain","page":"API","title":"Sampling a single chain","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.sample(::AbstractRNG, ::AbstractMCMC.AbstractModel, ::AbstractMCMC.AbstractSampler, ::Any)\nAbstractMCMC.sample(::AbstractRNG, ::Any, ::AbstractMCMC.AbstractSampler, ::Any)\n","category":"page"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler, Any}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbatractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler,\n N_or_isdone;\n kwargs...,\n)\n\nSample from the model with the Markov chain Monte Carlo sampler and return the samples.\n\nIf N_or_isdone is an Integer, exactly N_or_isdone samples are returned.\n\nOtherwise, sampling is performed until a convergence criterion N_or_isdone returns true. The convergence criterion has to be a function with the signature\n\nisdone(rng, model, sampler, samples, state, iteration; kwargs...)\n\nwhere state and iteration are the current state and iteration of the sampler, respectively. It should return true when sampling should end, and false otherwise.\n\n\n\n\n\n","category":"method"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler, Any}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler,\n N_or_isdone;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call sample with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/#Iterator","page":"API","title":"Iterator","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.steps(::AbstractRNG, ::AbstractMCMC.AbstractModel, ::AbstractMCMC.AbstractSampler)\nAbstractMCMC.steps(::AbstractRNG, ::Any, ::AbstractMCMC.AbstractSampler)","category":"page"},{"location":"api/#AbstractMCMC.steps-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.steps","text":"steps(\n rng::Random.AbstractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nCreate an iterator that returns samples from the model with the Markov chain Monte Carlo sampler.\n\nExamples\n\njulia> struct MyModel <: AbstractMCMC.AbstractModel end\n\njulia> struct MySampler <: AbstractMCMC.AbstractSampler end\n\njulia> function AbstractMCMC.step(rng, ::MyModel, ::MySampler, state=nothing; kwargs...)\n # all samples are zero\n return 0.0, state\n end\n\njulia> iterator = steps(MyModel(), MySampler());\n\njulia> collect(Iterators.take(iterator, 10)) == zeros(10)\ntrue\n\n\n\n\n\n","category":"method"},{"location":"api/#AbstractMCMC.steps-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.steps","text":"steps(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call steps with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/#Transducer","page":"API","title":"Transducer","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.Sample(::AbstractRNG, ::AbstractMCMC.AbstractModel, ::AbstractMCMC.AbstractSampler)\nAbstractMCMC.Sample(::AbstractRNG, ::Any, ::AbstractMCMC.AbstractSampler)","category":"page"},{"location":"api/#AbstractMCMC.Sample-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.Sample","text":"Sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nCreate a transducer that returns samples from the model with the Markov chain Monte Carlo sampler.\n\nExamples\n\njulia> struct MyModel <: AbstractMCMC.AbstractModel end\n\njulia> struct MySampler <: AbstractMCMC.AbstractSampler end\n\njulia> function AbstractMCMC.step(rng, ::MyModel, ::MySampler, state=nothing; kwargs...)\n # all samples are zero\n return 0.0, state\n end\n\njulia> transducer = Sample(MyModel(), MySampler());\n\njulia> collect(transducer(1:10)) == zeros(10)\ntrue\n\n\n\n\n\n","category":"method"},{"location":"api/#AbstractMCMC.Sample-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.Sample","text":"Sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call Sample with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/#Sampling-multiple-chains-in-parallel","page":"API","title":"Sampling multiple chains in parallel","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.sample(\n ::AbstractRNG,\n ::AbstractMCMC.AbstractModel,\n ::AbstractMCMC.AbstractSampler,\n ::AbstractMCMC.AbstractMCMCEnsemble,\n ::Integer,\n ::Integer,\n)\nAbstractMCMC.sample(\n ::AbstractRNG,\n ::Any,\n ::AbstractMCMC.AbstractSampler,\n ::AbstractMCMC.AbstractMCMCEnsemble,\n ::Integer,\n ::Integer,\n)","category":"page"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler, AbstractMCMC.AbstractMCMCEnsemble, Integer, Integer}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler,\n parallel::AbstractMCMCEnsemble,\n N::Integer,\n nchains::Integer;\n kwargs...,\n)\n\nSample nchains Monte Carlo Markov chains from the model with the sampler in parallel using the parallel algorithm, and combine them into a single chain.\n\n\n\n\n\n","category":"method"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler, AbstractMCMC.AbstractMCMCEnsemble, Integer, Integer}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler,\n parallel::AbstractMCMCEnsemble,\n N::Integer,\n nchains::Integer;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call sample with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/","page":"API","title":"API","text":"Two algorithms are provided for parallel sampling with multiple threads and multiple processes, and one allows for the user to sample multiple chains in serial (no parallelization):","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.MCMCThreads\nAbstractMCMC.MCMCDistributed\nAbstractMCMC.MCMCSerial","category":"page"},{"location":"api/#AbstractMCMC.MCMCThreads","page":"API","title":"AbstractMCMC.MCMCThreads","text":"MCMCThreads\n\nThe MCMCThreads algorithm allows users to sample MCMC chains in parallel using multiple threads.\n\n\n\n\n\n","category":"type"},{"location":"api/#AbstractMCMC.MCMCDistributed","page":"API","title":"AbstractMCMC.MCMCDistributed","text":"MCMCDistributed\n\nThe MCMCDistributed algorithm allows users to sample MCMC chains in parallel using multiple processes.\n\n\n\n\n\n","category":"type"},{"location":"api/#AbstractMCMC.MCMCSerial","page":"API","title":"AbstractMCMC.MCMCSerial","text":"MCMCSerial\n\nThe MCMCSerial algorithm allows users to sample serially, with no thread or process parallelism.\n\n\n\n\n\n","category":"type"},{"location":"api/#Common-keyword-arguments","page":"API","title":"Common keyword arguments","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"Common keyword arguments for regular and parallel sampling are:","category":"page"},{"location":"api/","page":"API","title":"API","text":"progress (default: AbstractMCMC.PROGRESS[] which is true initially): toggles progress logging\nchain_type (default: Any): determines the type of the returned chain\ncallback (default: nothing): if callback !== nothing, then callback(rng, model, sampler, sample, state, iteration) is called after every sampling step, where sample is the most recent sample of the Markov chain and state and iteration are the current state and iteration of the sampler\ndiscard_initial (default: 0): number of initial samples that are discarded\nthinning (default: 1): factor by which to thin samples.\ninitial_state (default: nothing): if initial_state !== nothing, the first call to AbstractMCMC.step is passed initial_state as the state argument.","category":"page"},{"location":"api/","page":"API","title":"API","text":"info: Info\nThe common keyword arguments progress, chain_type, and callback are not supported by the iterator AbstractMCMC.steps and the transducer AbstractMCMC.Sample.","category":"page"},{"location":"api/","page":"API","title":"API","text":"There is no \"official\" way for providing initial parameter values yet. However, multiple packages such as EllipticalSliceSampling.jl and AdvancedMH.jl support an initial_params keyword argument for setting the initial values when sampling a single chain. To ensure that sampling multiple chains \"just works\" when sampling of a single chain is implemented, we decided to support initial_params in the default implementations of the ensemble methods:","category":"page"},{"location":"api/","page":"API","title":"API","text":"initial_params (default: nothing): if initial_params isa AbstractArray, then the ith element of initial_params is used as initial parameters of the ith chain. If one wants to use the same initial parameters x for every chain, one can specify e.g. initial_params = FillArrays.Fill(x, N).","category":"page"},{"location":"api/","page":"API","title":"API","text":"Progress logging can be enabled and disabled globally with AbstractMCMC.setprogress!(progress).","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.setprogress!","category":"page"},{"location":"api/#AbstractMCMC.setprogress!","page":"API","title":"AbstractMCMC.setprogress!","text":"setprogress!(progress::Bool; silent::Bool=false)\n\nEnable progress logging globally if progress is true, and disable it otherwise. Optionally disable informational message if silent is true.\n\n\n\n\n\n","category":"function"},{"location":"api/#Chains","page":"API","title":"Chains","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"The chain_type keyword argument allows to set the type of the returned chain. A common choice is to return chains of type Chains from MCMCChains.jl.","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC defines the abstract type AbstractChains for Markov chains.","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.AbstractChains","category":"page"},{"location":"api/#AbstractMCMC.AbstractChains","page":"API","title":"AbstractMCMC.AbstractChains","text":"AbstractChains\n\nAbstractChains is an abstract type for an object that stores parameter samples generated through a MCMC process.\n\n\n\n\n\n","category":"type"},{"location":"api/","page":"API","title":"API","text":"For chains of this type, AbstractMCMC defines the following two methods.","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.chainscat\nAbstractMCMC.chainsstack","category":"page"},{"location":"api/#AbstractMCMC.chainscat","page":"API","title":"AbstractMCMC.chainscat","text":"chainscat(c::AbstractChains...)\n\nConcatenate multiple chains.\n\nBy default, the chains are concatenated along the third dimension by calling cat(c...; dims=3).\n\n\n\n\n\n","category":"function"},{"location":"api/#AbstractMCMC.chainsstack","page":"API","title":"AbstractMCMC.chainsstack","text":"chainsstack(c::AbstractVector)\n\nStack chains in c.\n\nBy default, the vector of chains is returned unmodified. If eltype(c) <: AbstractChains, then reduce(chainscat, c) is called.\n\n\n\n\n\n","category":"function"},{"location":"design/#Design","page":"Design","title":"Design","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"This page explains the default implementations and design choices of AbstractMCMC. It is not intended for users but for developers that want to implement the AbstractMCMC interface for Markov chain Monte Carlo sampling. The user-facing API is explained in API.","category":"page"},{"location":"design/#Overview","page":"Design","title":"Overview","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC provides a default implementation of the user-facing interface described in API. You can completely neglect these and define your own implementation of the interface. However, as described below, in most use cases the default implementation allows you to obtain support of parallel sampling, progress logging, callbacks, iterators, and transducers for free by just defining the sampling step of your inference algorithm, drastically reducing the amount of code you have to write. In general, the docstrings of the functions described below might be helpful if you intend to make use of the default implementations.","category":"page"},{"location":"design/#Basic-structure","page":"Design","title":"Basic structure","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"The simplified structure for regular sampling (the actual implementation contains some additional error checks and support for progress logging and callbacks) is","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"StatsBase.sample(\n rng::Random.AbstractRNG,\n model::AbstractMCMC.AbstractModel,\n sampler::AbstractMCMC.AbstractSampler,\n nsamples::Integer;\n chain_type = ::Type{Any},\n kwargs...\n)\n # Obtain the initial sample and state.\n sample, state = AbstractMCMC.step(rng, model, sampler; kwargs...)\n\n # Save the sample.\n samples = AbstractMCMC.samples(sample, model, sampler, N; kwargs...)\n samples = AbstractMCMC.save!!(samples, sample, 1, model, sampler, N; kwargs...)\n\n # Step through the sampler.\n for i in 2:N\n # Obtain the next sample and state.\n sample, state = AbstractMCMC.step(rng, model, sampler, state; kwargs...)\n\n # Save the sample.\n samples = AbstractMCMC.save!!(samples, sample, i, model, sampler, N; kwargs...)\n end\n\n return AbstractMCMC.bundle_samples(samples, model, sampler, state, chain_type; kwargs...)\nend","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"All other default implementations make use of the same structure and in particular call the same methods.","category":"page"},{"location":"design/#Sampling-step","page":"Design","title":"Sampling step","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"The only method for which no default implementation is provided (and hence which downstream packages have to implement) is AbstractMCMC.step. It defines the sampling step of the inference method.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.step","category":"page"},{"location":"design/#AbstractMCMC.step","page":"Design","title":"AbstractMCMC.step","text":"step(rng, model, sampler[, state; kwargs...])\n\nReturn a 2-tuple of the next sample and the next state of the MCMC sampler for model.\n\nSamples describe the results of a single step of the sampler. As an example, a sample might include a vector of parameters sampled from a prior distribution.\n\nWhen sampling using sample, every step call after the first has access to the current state of the sampler.\n\n\n\n\n\n","category":"function"},{"location":"design/#Collecting-samples","page":"Design","title":"Collecting samples","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"note: Note\nThis section does not apply to the iterator and transducer interface.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"After the initial sample is obtained, the default implementations for regular and parallel sampling (not for the iterator and the transducer since it is not needed there) create a container for all samples (the initial one and all subsequent samples) using AbstractMCMC.samples.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.samples","category":"page"},{"location":"design/#AbstractMCMC.samples","page":"Design","title":"AbstractMCMC.samples","text":"samples(sample, model, sampler[, N; kwargs...])\n\nGenerate a container for the samples of the MCMC sampler for the model, whose first sample is sample.\n\nThe method can be called with and without a predefined number N of samples.\n\n\n\n\n\n","category":"function"},{"location":"design/","page":"Design","title":"Design","text":"In each step, the sample is saved in the container by AbstractMCMC.save!!. The notation !! follows the convention of the package BangBang.jl which is used in the default implementation of AbstractMCMC.save!!. It indicates that the sample is pushed to the container but a \"widening\" fallback is used if the container type does not allow to save the sample. Therefore AbstractMCMC.save!! always has to return the container.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.save!!","category":"page"},{"location":"design/#AbstractMCMC.save!!","page":"Design","title":"AbstractMCMC.save!!","text":"save!!(samples, sample, iteration, model, sampler[, N; kwargs...])\n\nSave the sample of the MCMC sampler at the current iteration in the container of samples.\n\nThe function can be called with and without a predefined number N of samples. By default, AbstractMCMC uses push!! from the Julia package BangBang to append to the container, and widen its type if needed.\n\n\n\n\n\n","category":"function"},{"location":"design/","page":"Design","title":"Design","text":"For most use cases the default implementation of AbstractMCMC.samples and AbstractMCMC.save!! should work out of the box and hence need not to be overloaded in downstream code.","category":"page"},{"location":"design/#Creating-chains","page":"Design","title":"Creating chains","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"note: Note\nThis section does not apply to the iterator and transducer interface.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"At the end of the sampling procedure for regular and paralle sampling we transform the collection of samples to the desired output type by calling AbstractMCMC.bundle_samples.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.bundle_samples","category":"page"},{"location":"design/#AbstractMCMC.bundle_samples","page":"Design","title":"AbstractMCMC.bundle_samples","text":"bundle_samples(samples, model, sampler, state, chain_type[; kwargs...])\n\nBundle all samples that were sampled from the model with the given sampler in a chain.\n\nThe final state of the sampler can be included in the chain. The type of the chain can be specified with the chain_type argument.\n\nBy default, this method returns samples.\n\n\n\n\n\n","category":"function"},{"location":"design/","page":"Design","title":"Design","text":"The default implementation should be fine in most use cases, but downstream packages could, e.g., save the final state of the sampler as well if they overload AbstractMCMC.bundle_samples.","category":"page"},{"location":"gibbs/#state-Interface-Functions","page":"state Interface Functions","title":"state Interface Functions","text":"","category":"section"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We encourage sampler packages to implement the following interface functions for the state type(s) they maintain:","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"LogDensityProblems.logdensity(logdensity_model::AbstractMCMC.LogDensityModel, state::MHState; recompute_logp=true)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"This function takes the logdensity model and the state, and returns the log probability of the state. If recompute_logp is true, it should recompute the log probability of the state. Otherwise, it should use the log probability stored in the state.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Base.vec(state)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"This function takes the state and returns a vector of the parameter values stored in the state.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"(state::StateType)(logp::Float64)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"This function takes the state and a log probability value, and updates the state with the new log probability.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"These function will provide a minimum interface to interact with the state datatype, which a sampler package doesn't have to expose.","category":"page"},{"location":"gibbs/#Using-the-state-Interface-for-block-sampling-within-Gibbs","page":"state Interface Functions","title":"Using the state Interface for block sampling within Gibbs","text":"","category":"section"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"In this sections, we will demonstrate how a model package may use this state interface to support a Gibbs sampler that can support blocking sampling using different inference algorithms.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We consider a simple hierarchical model with a normal likelihood, with unknown mean and variance parameters.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"beginalign\nmu sim textNormal(0 1) \ntau^2 sim textInverseGamma(1 1) \nx_i sim textNormal(mu sqrttau^2)\nendalign","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We can write the log joint probability function as follows, where for the sake of simplicity for the following steps, we will assume that the mu and tau2 parameters are one-element vectors. And x is the data.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"function log_joint(; mu::Vector{Float64}, tau2::Vector{Float64}, x::Vector{Float64})\n # mu is the mean\n # tau2 is the variance\n # x is data\n\n # μ ~ Normal(0, 1)\n # τ² ~ InverseGamma(1, 1)\n # xᵢ ~ Normal(μ, √τ²)\n\n logp = 0.0\n mu = only(mu)\n tau2 = only(tau2)\n\n mu_prior = Normal(0, 1)\n logp += logpdf(mu_prior, mu)\n\n tau2_prior = InverseGamma(1, 1)\n logp += logpdf(tau2_prior, tau2)\n\n obs_prior = Normal(mu, sqrt(tau2))\n logp += sum(logpdf(obs_prior, xi) for xi in x)\n\n return logp\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"To make using LogDensityProblems interface, we create a simple type for this model.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"abstract type AbstractHierNormal end\n\nstruct HierNormal{Tdata<:NamedTuple} <: AbstractHierNormal\n data::Tdata\nend\n\nstruct ConditionedHierNormal{Tdata<:NamedTuple,Tconditioned_vars<:NamedTuple} <:\n AbstractHierNormal\n data::Tdata\n\n \" The variable to be conditioned on and its value\"\n conditioned_values::Tconditioned_vars\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"where ConditionedHierNormal is a type that represents the model conditioned on some variables, and","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"function AbstractPPL.condition(hn::HierNormal, conditioned_values::NamedTuple)\n return ConditionedHierNormal(hn.data, conditioned_values)\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"then we can simply write down the LogDensityProblems interface for this model.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"function LogDensityProblems.logdensity(\n hier_normal_model::ConditionedHierNormal{Tdata,Tconditioned_vars},\n params::AbstractVector,\n) where {Tdata,Tconditioned_vars}\n variable_to_condition = only(fieldnames(Tconditioned_vars))\n data = hier_normal_model.data\n conditioned_values = hier_normal_model.conditioned_values\n\n if variable_to_condition == :mu\n return log_joint(; mu=conditioned_values.mu, tau2=params, x=data.x)\n elseif variable_to_condition == :tau2\n return log_joint(; mu=params, tau2=conditioned_values.tau2, x=data.x)\n else\n error(\"Unsupported conditioning variable: $variable_to_condition\")\n end\nend\n\nfunction LogDensityProblems.capabilities(::HierNormal)\n return LogDensityProblems.LogDensityOrder{0}()\nend\n\nfunction LogDensityProblems.capabilities(::ConditionedHierNormal)\n return LogDensityProblems.LogDensityOrder{0}()\nend","category":"page"},{"location":"gibbs/#Sampler-Packages","page":"state Interface Functions","title":"Sampler Packages","text":"","category":"section"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"To illustrate the AbstractMCMC interface, we will first implement two very simple Metropolis-Hastings samplers: random walk and static proposal.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Although the interface doesn't force the sampler to implement Transition and State types, in practice, it has been the convention to do so.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Here we define some bare minimum types to represent the transitions and states.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"struct MHTransition{T}\n params::Vector{T}\nend\n\nstruct MHState{T}\n params::Vector{T}\n logp::Float64\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Next we define the state interface functions mentioned at the beginning of this section.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"# Interface 1: LogDensityProblems.logdensity\n# This function takes the logdensity function and the state (state is defined by the sampler package)\n# and returns the logdensity. It allows for optional recomputation of the log probability.\n# If recomputation is not needed, it returns the stored log probability from the state.\nfunction LogDensityProblems.logdensity(\n logdensity_model::AbstractMCMC.LogDensityModel, state::MHState; recompute_logp=true\n)\n logdensity_function = logdensity_model.logdensity\n return if recompute_logp\n AbstractMCMC.LogDensityProblems.logdensity(logdensity_function, state.params)\n else\n state.logp\n end\nend\n\n# Interface 2: Base.vec\n# This function takes a state and returns a vector of the parameter values stored in the state.\n# It is part of the interface for interacting with the state object.\nBase.vec(state::MHState) = state.params\n\n# Interface 3: (state::MHState)(logp::Float64)\n# This function allows the state to be updated with a new log probability.\n# ! this makes state into a Julia functor\n(state::MHState)(logp::Float64) = MHState(state.params, logp)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"It is very simple to implement the samplers according to the AbstractMCMC interface, where we can use LogDensityProblems.logdensity to easily read the log probability of the current state.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"\"\"\"\n RandomWalkMH{T} <: AbstractMCMC.AbstractSampler\n\nA random walk Metropolis-Hastings sampler with a normal proposal distribution. The field σ\nis the standard deviation of the proposal distribution.\n\"\"\"\nstruct RandomWalkMH{T} <: AbstractMHSampler\n σ::T\nend\n\n\"\"\"\n IndependentMH{T} <: AbstractMCMC.AbstractSampler\n\nA Metropolis-Hastings sampler with an independent proposal distribution.\n\"\"\"\nstruct IndependentMH{T} <: AbstractMHSampler\n proposal_dist::T\nend\n\n# the first step of the sampler\nfunction AbstractMCMC.step(\n rng::AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::AbstractMHSampler,\n args...;\n initial_params,\n kwargs...,\n)\n logdensity_function = logdensity_model.logdensity\n transition = MHTransition(initial_params)\n state = MHState(\n initial_params,\n only(LogDensityProblems.logdensity(logdensity_function, initial_params)),\n )\n\n return transition, state\nend\n\n@inline get_proposal_dist(sampler::RandomWalkMH, current_params::Vector{Float64}) =\n MvNormal(current_params, sampler.σ)\n@inline get_proposal_dist(sampler::IndependentMH, current_params::Vector{T}) where {T} =\n sampler.proposal_dist\n\n# the subsequent steps of the sampler\nfunction AbstractMCMC.step(\n rng::AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::AbstractMHSampler,\n state::MHState,\n args...;\n kwargs...,\n)\n logdensity_function = logdensity_model.logdensity\n current_params = state.params\n proposal_dist = get_proposal_dist(sampler, current_params)\n proposed_params = rand(rng, proposal_dist)\n logp_proposal = only(\n LogDensityProblems.logdensity(logdensity_function, proposed_params)\n )\n\n if log(rand(rng)) <\n compute_log_acceptance_ratio(sampler, state, proposed_params, logp_proposal)\n return MHTransition(proposed_params), MHState(proposed_params, logp_proposal)\n else\n return MHTransition(current_params), MHState(current_params, state.logp)\n end\nend\n\nfunction compute_log_acceptance_ratio(\n ::RandomWalkMH, state::MHState, ::Vector{Float64}, logp_proposal::Float64\n)\n return min(0, logp_proposal - state.logp)\nend\n\nfunction compute_log_acceptance_ratio(\n sampler::IndependentMH, state::MHState, proposal::Vector{T}, logp_proposal::Float64\n) where {T}\n return min(\n 0,\n logp_proposal - state.logp + logpdf(sampler.proposal_dist, state.params) -\n logpdf(sampler.proposal_dist, proposal),\n )\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"At last, we can proceed to implement the Gibbs sampler.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"\"\"\"\n Gibbs(sampler_map::NamedTuple)\n\nA Gibbs sampler that allows for block sampling using different inference algorithms for each parameter.\n\"\"\"\nstruct Gibbs{T<:NamedTuple} <: AbstractMCMC.AbstractSampler\n sampler_map::T\nend\n\nstruct GibbsState{TraceNT<:NamedTuple,StateNT<:NamedTuple,SizeNT<:NamedTuple}\n \"Contains the values of all parameters up to the last iteration.\"\n trace::TraceNT\n\n \"Maps parameters to their sampler-specific MCMC states.\"\n mcmc_states::StateNT\n\n \"Maps parameters to their sizes.\"\n variable_sizes::SizeNT\nend\n\nstruct GibbsTransition{ValuesNT<:NamedTuple}\n \"Realizations of the parameters, this is considered a \\\"sample\\\" in the MCMC chain.\"\n values::ValuesNT\nend\n\n\"\"\"\n update_trace(trace::NamedTuple, gibbs_state::GibbsState)\n\nUpdate the trace with the values from the MCMC states of the sub-problems.\n\"\"\"\nfunction update_trace(trace::NamedTuple, gibbs_state::GibbsState)\n for parameter_variable in keys(gibbs_state.mcmc_states)\n sub_state = gibbs_state.mcmc_states[parameter_variable]\n sub_state_params = Base.vec(sub_state)\n unflattened_sub_state_params = unflatten(\n sub_state_params,\n NamedTuple{(parameter_variable,)}((\n gibbs_state.variable_sizes[parameter_variable],\n )),\n )\n trace = merge(trace, unflattened_sub_state_params)\n end\n return trace\nend\n\nfunction error_if_not_fully_initialized(\n initial_params::NamedTuple{ParamNames}, sampler::Gibbs{<:NamedTuple{SamplerNames}}\n) where {ParamNames,SamplerNames}\n if Set(ParamNames) != Set(SamplerNames)\n throw(\n ArgumentError(\n \"initial_params must contain all parameters in the model, expected $(SamplerNames), got $(ParamNames)\",\n ),\n )\n end\nend\n\nfunction AbstractMCMC.step(\n rng::Random.AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::Gibbs{Tsamplingmap},\n args...;\n initial_params::NamedTuple,\n kwargs...,\n) where {Tsamplingmap}\n error_if_not_fully_initialized(initial_params, sampler)\n\n model_parameter_names = fieldnames(Tsamplingmap)\n results = map(model_parameter_names) do parameter_variable\n sub_sampler = sampler.sampler_map[parameter_variable]\n\n variables_to_be_conditioned_on = setdiff(\n model_parameter_names, (parameter_variable,)\n )\n conditioning_variables_values = NamedTuple{Tuple(variables_to_be_conditioned_on)}(\n Tuple([initial_params[g] for g in variables_to_be_conditioned_on])\n )\n sub_problem_parameters_values = NamedTuple{(parameter_variable,)}((\n initial_params[parameter_variable],\n ))\n\n # LogDensityProblems' `logdensity` function expects a single vector of real numbers\n # `Gibbs` stores the parameters as a named tuple, thus we need to flatten the sub_problem_parameters_values\n # and unflatten after the sampling step\n flattened_sub_problem_parameters_values = flatten(sub_problem_parameters_values)\n\n sub_state = last(\n AbstractMCMC.step(\n rng,\n AbstractMCMC.LogDensityModel(\n AbstractPPL.condition(\n logdensity_model.logdensity, conditioning_variables_values\n ),\n ),\n sub_sampler,\n args...;\n initial_params=flattened_sub_problem_parameters_values,\n kwargs...,\n ),\n )\n (sub_state, Tuple(size(initial_params[parameter_variable])))\n end\n\n mcmc_states_tuple = first.(results)\n variable_sizes_tuple = last.(results)\n\n gibbs_state = GibbsState(\n initial_params,\n NamedTuple{Tuple(model_parameter_names)}(mcmc_states_tuple),\n NamedTuple{Tuple(model_parameter_names)}(variable_sizes_tuple),\n )\n\n trace = update_trace(NamedTuple(), gibbs_state)\n return GibbsTransition(trace), gibbs_state\nend\n\n# subsequent steps\nfunction AbstractMCMC.step(\n rng::Random.AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::Gibbs{Tsamplingmap},\n gibbs_state::GibbsState,\n args...;\n kwargs...,\n) where {Tsamplingmap}\n (; trace, mcmc_states, variable_sizes) = gibbs_state\n\n model_parameter_names = fieldnames(Tsamplingmap)\n mcmc_states = map(model_parameter_names) do parameter_variable\n sub_sampler = sampler.sampler_map[parameter_variable]\n sub_state = mcmc_states[parameter_variable]\n variables_to_be_conditioned_on = setdiff(\n model_parameter_names, (parameter_variable,)\n )\n conditioning_variables_values = NamedTuple{Tuple(variables_to_be_conditioned_on)}(\n Tuple([trace[g] for g in variables_to_be_conditioned_on])\n )\n cond_logdensity = AbstractPPL.condition(\n logdensity_model.logdensity, conditioning_variables_values\n )\n cond_logdensity_model = AbstractMCMC.LogDensityModel(cond_logdensity)\n\n logp = LogDensityProblems.logdensity(cond_logdensity_model, sub_state)\n sub_state = (sub_state)(logp)\n sub_state = last(\n AbstractMCMC.step(\n rng, cond_logdensity_model, sub_sampler, sub_state, args...; kwargs...\n ),\n )\n trace = update_trace(trace, gibbs_state)\n sub_state\n end\n mcmc_states = NamedTuple{Tuple(model_parameter_names)}(mcmc_states)\n\n return GibbsTransition(trace), GibbsState(trace, mcmc_states, variable_sizes)\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"where we use two utility functions flatten and unflatten to convert between the single vector of real numbers and the named tuple of parameters.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"\"\"\"\n flatten(trace::NamedTuple)\n\nFlatten all the values in the trace into a single vector. Variable names information is discarded.\n\"\"\"\nfunction flatten(trace::NamedTuple)\n return reduce(vcat, vec.(values(trace)))\nend\n\n\"\"\"\n unflatten(vec::AbstractVector, variable_names::Vector{Symbol}, variable_sizes::Vector{Tuple})\n\nReverse operation of flatten. Reshape the vector into the original arrays using size information.\n\"\"\"\nfunction unflatten(\n vec::AbstractVector, variable_names_and_sizes::NamedTuple{variable_names}\n) where {variable_names}\n result = Dict{Symbol,Array}()\n start_idx = 1\n for name in variable_names\n size = variable_names_and_sizes[name]\n end_idx = start_idx + prod(size) - 1\n result[name] = reshape(vec[start_idx:end_idx], size...)\n start_idx = end_idx + 1\n end\n\n return NamedTuple{variable_names}(Tuple([result[name] for name in variable_names]))\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Some points worth noting:","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We are using NamedTuple to store the mapping between variables and samplers. The order will determine the order of the Gibbs sweeps. A limitation is that exactly one sampler for each variable is required, which means it is less flexible than Gibbs in Turing.jl.\nFor each conditional probability problem, we need to store the sampler states for each variable group and also the values of all the variables from last iteration.\nThe first step of the Gibbs sampler is to setup the states for each conditional probability problem.\nIn the following steps of the Gibbs sampler, it will do a sweep over all the conditional probability problems, and update the sampler states for each problem. In each step of the sweep, it will do the following:\ncondition on the values of all variables that are not in the current group\nrecompute the log probability of the current state, because the values of the variables that are not in the current group may have changed\nperform a step of the sampler for the conditional probability problem, and update the sampler state\nupdate the vi with the new values from the sampler state","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"The state interface in AbstractMCMC allows the Gibbs sampler to be agnostic of the details of the sampler state, and acquire the values of the parameters from individual sampler states.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Now we can use the Gibbs sampler to sample from the hierarchical normal model.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"First we generate some data,","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"N = 100 # Number of data points\nmu_true = 0.5 # True mean\ntau2_true = 2.0 # True variance\n\nx_data = rand(Normal(mu_true, sqrt(tau2_true)), N)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Then we can create a HierNormal model, with the data we just generated.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"hn = HierNormal((x=x_data,))","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Using Gibbs sampling allows us to use random walk MH for mu and prior MH for tau2, because tau2 has support only on positive real numbers.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"samples = sample(\n hn,\n Gibbs((\n mu=RandomWalkMH(0.3),\n tau2=IndependentMH(product_distribution([InverseGamma(1, 1)])),\n )),\n 10000;\n initial_params=(mu=[0.0], tau2=[1.0]),\n)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Then we can extract the samples and compute the mean of the samples.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"mu_samples = [sample.values.mu for sample in samples][20001:end]\ntau2_samples = [sample.values.tau2 for sample in samples][20001:end]\n\nmean(mu_samples)\nmean(tau2_samples)\n(mu_mean, tau2_mean)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"the result should looks like:","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"(4.995812149309413, 1.9372372289677886)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"which is close to the true values (5, 2).","category":"page"},{"location":"#AbstractMCMC.jl","page":"Home","title":"AbstractMCMC.jl","text":"","category":"section"},{"location":"","page":"Home","title":"Home","text":"Abstract types and interfaces for Markov chain Monte Carlo methods.","category":"page"},{"location":"","page":"Home","title":"Home","text":"AbstractMCMC defines an interface for sampling and combining Markov chains. It comes with a default sampling algorithm that provides support of progress bars, parallel sampling (multithreaded and multicore), and user-provided callbacks out of the box. Typically developers only have to define the sampling step of their inference method in an iterator-like fashion to make use of this functionality. Additionally, the package defines an iterator and a transducer for sampling Markov chains based on the interface.","category":"page"}] +[{"location":"api/#API","page":"API","title":"API","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC defines an interface for sampling Markov chains.","category":"page"},{"location":"api/#Model","page":"API","title":"Model","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.AbstractModel\nAbstractMCMC.LogDensityModel","category":"page"},{"location":"api/#AbstractMCMC.AbstractModel","page":"API","title":"AbstractMCMC.AbstractModel","text":"AbstractModel\n\nAn AbstractModel represents a generic model type that can be used to perform inference.\n\n\n\n\n\n","category":"type"},{"location":"api/#AbstractMCMC.LogDensityModel","page":"API","title":"AbstractMCMC.LogDensityModel","text":"LogDensityModel <: AbstractMCMC.AbstractModel\n\nWrapper around something that implements the LogDensityProblem.jl interface.\n\nNote that this does not implement the LogDensityProblems.jl interface itself, but it simply useful for indicating to the sample and other AbstractMCMC methods that the wrapped object implements the LogDensityProblems.jl interface.\n\nFields\n\nlogdensity: The object that implements the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"type"},{"location":"api/#Sampler","page":"API","title":"Sampler","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.AbstractSampler","category":"page"},{"location":"api/#AbstractMCMC.AbstractSampler","page":"API","title":"AbstractMCMC.AbstractSampler","text":"AbstractSampler\n\nThe AbstractSampler type is intended to be inherited from when implementing a custom sampler. Any persistent state information should be saved in a subtype of AbstractSampler.\n\nWhen defining a new sampler, you should also overload the function transition_type, which tells the sample function what type of parameter it should expect to receive.\n\n\n\n\n\n","category":"type"},{"location":"api/#Sampling-a-single-chain","page":"API","title":"Sampling a single chain","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.sample(::AbstractRNG, ::AbstractMCMC.AbstractModel, ::AbstractMCMC.AbstractSampler, ::Any)\nAbstractMCMC.sample(::AbstractRNG, ::Any, ::AbstractMCMC.AbstractSampler, ::Any)\n","category":"page"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler, Any}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbatractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler,\n N_or_isdone;\n kwargs...,\n)\n\nSample from the model with the Markov chain Monte Carlo sampler and return the samples.\n\nIf N_or_isdone is an Integer, exactly N_or_isdone samples are returned.\n\nOtherwise, sampling is performed until a convergence criterion N_or_isdone returns true. The convergence criterion has to be a function with the signature\n\nisdone(rng, model, sampler, samples, state, iteration; kwargs...)\n\nwhere state and iteration are the current state and iteration of the sampler, respectively. It should return true when sampling should end, and false otherwise.\n\n\n\n\n\n","category":"method"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler, Any}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler,\n N_or_isdone;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call sample with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/#Iterator","page":"API","title":"Iterator","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.steps(::AbstractRNG, ::AbstractMCMC.AbstractModel, ::AbstractMCMC.AbstractSampler)\nAbstractMCMC.steps(::AbstractRNG, ::Any, ::AbstractMCMC.AbstractSampler)","category":"page"},{"location":"api/#AbstractMCMC.steps-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.steps","text":"steps(\n rng::Random.AbstractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nCreate an iterator that returns samples from the model with the Markov chain Monte Carlo sampler.\n\nExamples\n\njulia> struct MyModel <: AbstractMCMC.AbstractModel end\n\njulia> struct MySampler <: AbstractMCMC.AbstractSampler end\n\njulia> function AbstractMCMC.step(rng, ::MyModel, ::MySampler, state=nothing; kwargs...)\n # all samples are zero\n return 0.0, state\n end\n\njulia> iterator = steps(MyModel(), MySampler());\n\njulia> collect(Iterators.take(iterator, 10)) == zeros(10)\ntrue\n\n\n\n\n\n","category":"method"},{"location":"api/#AbstractMCMC.steps-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.steps","text":"steps(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call steps with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/#Transducer","page":"API","title":"Transducer","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.Sample(::AbstractRNG, ::AbstractMCMC.AbstractModel, ::AbstractMCMC.AbstractSampler)\nAbstractMCMC.Sample(::AbstractRNG, ::Any, ::AbstractMCMC.AbstractSampler)","category":"page"},{"location":"api/#AbstractMCMC.Sample-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.Sample","text":"Sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nCreate a transducer that returns samples from the model with the Markov chain Monte Carlo sampler.\n\nExamples\n\njulia> struct MyModel <: AbstractMCMC.AbstractModel end\n\njulia> struct MySampler <: AbstractMCMC.AbstractSampler end\n\njulia> function AbstractMCMC.step(rng, ::MyModel, ::MySampler, state=nothing; kwargs...)\n # all samples are zero\n return 0.0, state\n end\n\njulia> transducer = Sample(MyModel(), MySampler());\n\njulia> collect(transducer(1:10)) == zeros(10)\ntrue\n\n\n\n\n\n","category":"method"},{"location":"api/#AbstractMCMC.Sample-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler}","page":"API","title":"AbstractMCMC.Sample","text":"Sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call Sample with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/#Sampling-multiple-chains-in-parallel","page":"API","title":"Sampling multiple chains in parallel","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.sample(\n ::AbstractRNG,\n ::AbstractMCMC.AbstractModel,\n ::AbstractMCMC.AbstractSampler,\n ::AbstractMCMC.AbstractMCMCEnsemble,\n ::Integer,\n ::Integer,\n)\nAbstractMCMC.sample(\n ::AbstractRNG,\n ::Any,\n ::AbstractMCMC.AbstractSampler,\n ::AbstractMCMC.AbstractMCMCEnsemble,\n ::Integer,\n ::Integer,\n)","category":"page"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, AbstractMCMC.AbstractModel, AbstractMCMC.AbstractSampler, AbstractMCMC.AbstractMCMCEnsemble, Integer, Integer}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n model::AbstractModel,\n sampler::AbstractSampler,\n parallel::AbstractMCMCEnsemble,\n N::Integer,\n nchains::Integer;\n kwargs...,\n)\n\nSample nchains Monte Carlo Markov chains from the model with the sampler in parallel using the parallel algorithm, and combine them into a single chain.\n\n\n\n\n\n","category":"method"},{"location":"api/#StatsBase.sample-Tuple{AbstractRNG, Any, AbstractMCMC.AbstractSampler, AbstractMCMC.AbstractMCMCEnsemble, Integer, Integer}","page":"API","title":"StatsBase.sample","text":"sample(\n rng::Random.AbstractRNG=Random.default_rng(),\n logdensity,\n sampler::AbstractSampler,\n parallel::AbstractMCMCEnsemble,\n N::Integer,\n nchains::Integer;\n kwargs...,\n)\n\nWrap the logdensity function in a LogDensityModel, and call sample with the resulting model instead of logdensity.\n\nThe logdensity function has to support the LogDensityProblems.jl interface.\n\n\n\n\n\n","category":"method"},{"location":"api/","page":"API","title":"API","text":"Two algorithms are provided for parallel sampling with multiple threads and multiple processes, and one allows for the user to sample multiple chains in serial (no parallelization):","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.MCMCThreads\nAbstractMCMC.MCMCDistributed\nAbstractMCMC.MCMCSerial","category":"page"},{"location":"api/#AbstractMCMC.MCMCThreads","page":"API","title":"AbstractMCMC.MCMCThreads","text":"MCMCThreads\n\nThe MCMCThreads algorithm allows users to sample MCMC chains in parallel using multiple threads.\n\n\n\n\n\n","category":"type"},{"location":"api/#AbstractMCMC.MCMCDistributed","page":"API","title":"AbstractMCMC.MCMCDistributed","text":"MCMCDistributed\n\nThe MCMCDistributed algorithm allows users to sample MCMC chains in parallel using multiple processes.\n\n\n\n\n\n","category":"type"},{"location":"api/#AbstractMCMC.MCMCSerial","page":"API","title":"AbstractMCMC.MCMCSerial","text":"MCMCSerial\n\nThe MCMCSerial algorithm allows users to sample serially, with no thread or process parallelism.\n\n\n\n\n\n","category":"type"},{"location":"api/#Common-keyword-arguments","page":"API","title":"Common keyword arguments","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"Common keyword arguments for regular and parallel sampling are:","category":"page"},{"location":"api/","page":"API","title":"API","text":"progress (default: AbstractMCMC.PROGRESS[] which is true initially): toggles progress logging\nchain_type (default: Any): determines the type of the returned chain\ncallback (default: nothing): if callback !== nothing, then callback(rng, model, sampler, sample, state, iteration) is called after every sampling step, where sample is the most recent sample of the Markov chain and state and iteration are the current state and iteration of the sampler\ndiscard_initial (default: 0): number of initial samples that are discarded\nthinning (default: 1): factor by which to thin samples.\ninitial_state (default: nothing): if initial_state !== nothing, the first call to AbstractMCMC.step is passed initial_state as the state argument.","category":"page"},{"location":"api/","page":"API","title":"API","text":"info: Info\nThe common keyword arguments progress, chain_type, and callback are not supported by the iterator AbstractMCMC.steps and the transducer AbstractMCMC.Sample.","category":"page"},{"location":"api/","page":"API","title":"API","text":"There is no \"official\" way for providing initial parameter values yet. However, multiple packages such as EllipticalSliceSampling.jl and AdvancedMH.jl support an initial_params keyword argument for setting the initial values when sampling a single chain. To ensure that sampling multiple chains \"just works\" when sampling of a single chain is implemented, we decided to support initial_params in the default implementations of the ensemble methods:","category":"page"},{"location":"api/","page":"API","title":"API","text":"initial_params (default: nothing): if initial_params isa AbstractArray, then the ith element of initial_params is used as initial parameters of the ith chain. If one wants to use the same initial parameters x for every chain, one can specify e.g. initial_params = FillArrays.Fill(x, N).","category":"page"},{"location":"api/","page":"API","title":"API","text":"Progress logging can be enabled and disabled globally with AbstractMCMC.setprogress!(progress).","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.setprogress!","category":"page"},{"location":"api/#AbstractMCMC.setprogress!","page":"API","title":"AbstractMCMC.setprogress!","text":"setprogress!(progress::Bool; silent::Bool=false)\n\nEnable progress logging globally if progress is true, and disable it otherwise. Optionally disable informational message if silent is true.\n\n\n\n\n\n","category":"function"},{"location":"api/#Chains","page":"API","title":"Chains","text":"","category":"section"},{"location":"api/","page":"API","title":"API","text":"The chain_type keyword argument allows to set the type of the returned chain. A common choice is to return chains of type Chains from MCMCChains.jl.","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC defines the abstract type AbstractChains for Markov chains.","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.AbstractChains","category":"page"},{"location":"api/#AbstractMCMC.AbstractChains","page":"API","title":"AbstractMCMC.AbstractChains","text":"AbstractChains\n\nAbstractChains is an abstract type for an object that stores parameter samples generated through a MCMC process.\n\n\n\n\n\n","category":"type"},{"location":"api/","page":"API","title":"API","text":"For chains of this type, AbstractMCMC defines the following two methods.","category":"page"},{"location":"api/","page":"API","title":"API","text":"AbstractMCMC.chainscat\nAbstractMCMC.chainsstack","category":"page"},{"location":"api/#AbstractMCMC.chainscat","page":"API","title":"AbstractMCMC.chainscat","text":"chainscat(c::AbstractChains...)\n\nConcatenate multiple chains.\n\nBy default, the chains are concatenated along the third dimension by calling cat(c...; dims=3).\n\n\n\n\n\n","category":"function"},{"location":"api/#AbstractMCMC.chainsstack","page":"API","title":"AbstractMCMC.chainsstack","text":"chainsstack(c::AbstractVector)\n\nStack chains in c.\n\nBy default, the vector of chains is returned unmodified. If eltype(c) <: AbstractChains, then reduce(chainscat, c) is called.\n\n\n\n\n\n","category":"function"},{"location":"design/#Design","page":"Design","title":"Design","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"This page explains the default implementations and design choices of AbstractMCMC. It is not intended for users but for developers that want to implement the AbstractMCMC interface for Markov chain Monte Carlo sampling. The user-facing API is explained in API.","category":"page"},{"location":"design/#Overview","page":"Design","title":"Overview","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC provides a default implementation of the user-facing interface described in API. You can completely neglect these and define your own implementation of the interface. However, as described below, in most use cases the default implementation allows you to obtain support of parallel sampling, progress logging, callbacks, iterators, and transducers for free by just defining the sampling step of your inference algorithm, drastically reducing the amount of code you have to write. In general, the docstrings of the functions described below might be helpful if you intend to make use of the default implementations.","category":"page"},{"location":"design/#Basic-structure","page":"Design","title":"Basic structure","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"The simplified structure for regular sampling (the actual implementation contains some additional error checks and support for progress logging and callbacks) is","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"StatsBase.sample(\n rng::Random.AbstractRNG,\n model::AbstractMCMC.AbstractModel,\n sampler::AbstractMCMC.AbstractSampler,\n nsamples::Integer;\n chain_type = ::Type{Any},\n kwargs...\n)\n # Obtain the initial sample and state.\n sample, state = AbstractMCMC.step(rng, model, sampler; kwargs...)\n\n # Save the sample.\n samples = AbstractMCMC.samples(sample, model, sampler, N; kwargs...)\n samples = AbstractMCMC.save!!(samples, sample, 1, model, sampler, N; kwargs...)\n\n # Step through the sampler.\n for i in 2:N\n # Obtain the next sample and state.\n sample, state = AbstractMCMC.step(rng, model, sampler, state; kwargs...)\n\n # Save the sample.\n samples = AbstractMCMC.save!!(samples, sample, i, model, sampler, N; kwargs...)\n end\n\n return AbstractMCMC.bundle_samples(samples, model, sampler, state, chain_type; kwargs...)\nend","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"All other default implementations make use of the same structure and in particular call the same methods.","category":"page"},{"location":"design/#Sampling-step","page":"Design","title":"Sampling step","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"The only method for which no default implementation is provided (and hence which downstream packages have to implement) is AbstractMCMC.step. It defines the sampling step of the inference method.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.step","category":"page"},{"location":"design/#AbstractMCMC.step","page":"Design","title":"AbstractMCMC.step","text":"step(rng, model, sampler[, state; kwargs...])\n\nReturn a 2-tuple of the next sample and the next state of the MCMC sampler for model.\n\nSamples describe the results of a single step of the sampler. As an example, a sample might include a vector of parameters sampled from a prior distribution.\n\nWhen sampling using sample, every step call after the first has access to the current state of the sampler.\n\n\n\n\n\n","category":"function"},{"location":"design/#Collecting-samples","page":"Design","title":"Collecting samples","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"note: Note\nThis section does not apply to the iterator and transducer interface.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"After the initial sample is obtained, the default implementations for regular and parallel sampling (not for the iterator and the transducer since it is not needed there) create a container for all samples (the initial one and all subsequent samples) using AbstractMCMC.samples.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.samples","category":"page"},{"location":"design/#AbstractMCMC.samples","page":"Design","title":"AbstractMCMC.samples","text":"samples(sample, model, sampler[, N; kwargs...])\n\nGenerate a container for the samples of the MCMC sampler for the model, whose first sample is sample.\n\nThe method can be called with and without a predefined number N of samples.\n\n\n\n\n\n","category":"function"},{"location":"design/","page":"Design","title":"Design","text":"In each step, the sample is saved in the container by AbstractMCMC.save!!. The notation !! follows the convention of the package BangBang.jl which is used in the default implementation of AbstractMCMC.save!!. It indicates that the sample is pushed to the container but a \"widening\" fallback is used if the container type does not allow to save the sample. Therefore AbstractMCMC.save!! always has to return the container.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.save!!","category":"page"},{"location":"design/#AbstractMCMC.save!!","page":"Design","title":"AbstractMCMC.save!!","text":"save!!(samples, sample, iteration, model, sampler[, N; kwargs...])\n\nSave the sample of the MCMC sampler at the current iteration in the container of samples.\n\nThe function can be called with and without a predefined number N of samples. By default, AbstractMCMC uses push!! from the Julia package BangBang to append to the container, and widen its type if needed.\n\n\n\n\n\n","category":"function"},{"location":"design/","page":"Design","title":"Design","text":"For most use cases the default implementation of AbstractMCMC.samples and AbstractMCMC.save!! should work out of the box and hence need not to be overloaded in downstream code.","category":"page"},{"location":"design/#Creating-chains","page":"Design","title":"Creating chains","text":"","category":"section"},{"location":"design/","page":"Design","title":"Design","text":"note: Note\nThis section does not apply to the iterator and transducer interface.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"At the end of the sampling procedure for regular and paralle sampling we transform the collection of samples to the desired output type by calling AbstractMCMC.bundle_samples.","category":"page"},{"location":"design/","page":"Design","title":"Design","text":"AbstractMCMC.bundle_samples","category":"page"},{"location":"design/#AbstractMCMC.bundle_samples","page":"Design","title":"AbstractMCMC.bundle_samples","text":"bundle_samples(samples, model, sampler, state, chain_type[; kwargs...])\n\nBundle all samples that were sampled from the model with the given sampler in a chain.\n\nThe final state of the sampler can be included in the chain. The type of the chain can be specified with the chain_type argument.\n\nBy default, this method returns samples.\n\n\n\n\n\n","category":"function"},{"location":"design/","page":"Design","title":"Design","text":"The default implementation should be fine in most use cases, but downstream packages could, e.g., save the final state of the sampler as well if they overload AbstractMCMC.bundle_samples.","category":"page"},{"location":"gibbs/#state-Interface-Functions","page":"state Interface Functions","title":"state Interface Functions","text":"","category":"section"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We encourage sampler packages to implement the following interface functions for the state type(s) they maintain:","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"LogDensityProblems.logdensity(logdensity_model::AbstractMCMC.LogDensityModel, state::MHState; recompute_logp=true)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"This function takes the logdensity model and the state, and returns the log probability of the state. If recompute_logp is true, it should recompute the log probability of the state. Otherwise, it could use the log probability stored in the state.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Base.vec(state)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"This function takes the state and returns a vector of the parameter values stored in the state.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"(state::StateType)(logp::Float64)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"This function takes the state and a log probability value, and returns a new state with the updated log probability.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"These functions provide a minimal interface to interact with the state datatype, which a sampler package can optionally implement. The interface facilitates the implementation of \"meta-algorithms\" that combine different samplers. We will demonstrate how it can be used to implement Gibbs sampling in the following sections.","category":"page"},{"location":"gibbs/#Using-the-state-Interface-for-block-sampling-within-Gibbs","page":"state Interface Functions","title":"Using the state Interface for block sampling within Gibbs","text":"","category":"section"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"In this sections, we will demonstrate how a model package may use this state interface to support a Gibbs sampler that can support blocking sampling using different inference algorithms.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We consider a simple hierarchical model with a normal likelihood, with unknown mean and variance parameters.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"beginalign\nmu sim textNormal(0 1) \ntau^2 sim textInverseGamma(1 1) \nx_i sim textNormal(mu sqrttau^2)\nendalign","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We can write the log joint probability function as follows, where for the sake of simplicity for the following steps, we will assume that the mu and tau2 parameters are one-element vectors. And x is the data.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"function log_joint(; mu::Vector{Float64}, tau2::Vector{Float64}, x::Vector{Float64})\n # mu is the mean\n # tau2 is the variance\n # x is data\n\n # μ ~ Normal(0, 1)\n # τ² ~ InverseGamma(1, 1)\n # xᵢ ~ Normal(μ, √τ²)\n\n logp = 0.0\n mu = only(mu)\n tau2 = only(tau2)\n\n mu_prior = Normal(0, 1)\n logp += logpdf(mu_prior, mu)\n\n tau2_prior = InverseGamma(1, 1)\n logp += logpdf(tau2_prior, tau2)\n\n obs_prior = Normal(mu, sqrt(tau2))\n logp += sum(logpdf(obs_prior, xi) for xi in x)\n\n return logp\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"To make using LogDensityProblems interface, we create a simple type for this model.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"abstract type AbstractHierNormal end\n\nstruct HierNormal{Tdata<:NamedTuple} <: AbstractHierNormal\n data::Tdata\nend\n\nstruct ConditionedHierNormal{Tdata<:NamedTuple,Tconditioned_vars<:NamedTuple} <:\n AbstractHierNormal\n data::Tdata\n\n \" The variable to be conditioned on and its value\"\n conditioned_values::Tconditioned_vars\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"where ConditionedHierNormal is a type that represents the model conditioned on some variables, and","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"function AbstractPPL.condition(hn::HierNormal, conditioned_values::NamedTuple)\n return ConditionedHierNormal(hn.data, conditioned_values)\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"then we can simply write down the LogDensityProblems interface for this model.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"function LogDensityProblems.logdensity(\n hier_normal_model::ConditionedHierNormal{Tdata,Tconditioned_vars},\n params::AbstractVector,\n) where {Tdata,Tconditioned_vars}\n variable_to_condition = only(fieldnames(Tconditioned_vars))\n data = hier_normal_model.data\n conditioned_values = hier_normal_model.conditioned_values\n\n if variable_to_condition == :mu\n return log_joint(; mu=conditioned_values.mu, tau2=params, x=data.x)\n elseif variable_to_condition == :tau2\n return log_joint(; mu=params, tau2=conditioned_values.tau2, x=data.x)\n else\n error(\"Unsupported conditioning variable: $variable_to_condition\")\n end\nend\n\nfunction LogDensityProblems.capabilities(::HierNormal)\n return LogDensityProblems.LogDensityOrder{0}()\nend\n\nfunction LogDensityProblems.capabilities(::ConditionedHierNormal)\n return LogDensityProblems.LogDensityOrder{0}()\nend","category":"page"},{"location":"gibbs/#Implementing-A-Sampler-with-AbstractMCMC-Interface","page":"state Interface Functions","title":"Implementing A Sampler with AbstractMCMC Interface","text":"","category":"section"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"To illustrate the AbstractMCMC interface, we will first implement two very simple Metropolis-Hastings samplers: random walk and static proposal.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Although the interface doesn't force the sampler to implement Transition and State types, in practice, it has been the convention to do so.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Here we define some bare minimum types to represent the transitions and states.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"struct MHTransition{T}\n params::Vector{T}\nend\n\nstruct MHState{T}\n params::Vector{T}\n logp::Float64\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Next we define the state interface functions mentioned at the beginning of this section.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"# Interface 1: LogDensityProblems.logdensity\n# This function takes the logdensity function and the state (state is defined by the sampler package)\n# and returns the logdensity. It allows for optional recomputation of the log probability.\n# If recomputation is not needed, it returns the stored log probability from the state.\nfunction LogDensityProblems.logdensity(\n logdensity_model::AbstractMCMC.LogDensityModel, state::MHState; recompute_logp=true\n)\n logdensity_function = logdensity_model.logdensity\n return if recompute_logp\n AbstractMCMC.LogDensityProblems.logdensity(logdensity_function, state.params)\n else\n state.logp\n end\nend\n\n# Interface 2: Base.vec\n# This function takes a state and returns a vector of the parameter values stored in the state.\n# It is part of the interface for interacting with the state object.\nBase.vec(state::MHState) = state.params\n\n# Interface 3: (state::MHState)(logp::Float64)\n# This function allows the state to be updated with a new log probability.\n# ! this makes state into a Julia functor\n(state::MHState)(logp::Float64) = MHState(state.params, logp)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"It is very simple to implement the samplers according to the AbstractMCMC interface, where we can use LogDensityProblems.logdensity to easily read the log probability of the current state.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"\"\"\"\n RandomWalkMH{T} <: AbstractMCMC.AbstractSampler\n\nA random walk Metropolis-Hastings sampler with a normal proposal distribution. The field σ\nis the standard deviation of the proposal distribution.\n\"\"\"\nstruct RandomWalkMH{T} <: AbstractMHSampler\n σ::T\nend\n\n\"\"\"\n IndependentMH{T} <: AbstractMCMC.AbstractSampler\n\nA Metropolis-Hastings sampler with an independent proposal distribution.\n\"\"\"\nstruct IndependentMH{T} <: AbstractMHSampler\n proposal_dist::T\nend\n\n# the first step of the sampler\nfunction AbstractMCMC.step(\n rng::AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::AbstractMHSampler,\n args...;\n initial_params,\n kwargs...,\n)\n logdensity_function = logdensity_model.logdensity\n transition = MHTransition(initial_params)\n state = MHState(\n initial_params,\n only(LogDensityProblems.logdensity(logdensity_function, initial_params)),\n )\n\n return transition, state\nend\n\n@inline get_proposal_dist(sampler::RandomWalkMH, current_params::Vector{Float64}) =\n MvNormal(current_params, sampler.σ)\n@inline get_proposal_dist(sampler::IndependentMH, current_params::Vector{T}) where {T} =\n sampler.proposal_dist\n\n# the subsequent steps of the sampler\nfunction AbstractMCMC.step(\n rng::AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::AbstractMHSampler,\n state::MHState,\n args...;\n kwargs...,\n)\n logdensity_function = logdensity_model.logdensity\n current_params = state.params\n proposal_dist = get_proposal_dist(sampler, current_params)\n proposed_params = rand(rng, proposal_dist)\n logp_proposal = only(\n LogDensityProblems.logdensity(logdensity_function, proposed_params)\n )\n\n if log(rand(rng)) <\n compute_log_acceptance_ratio(sampler, state, proposed_params, logp_proposal)\n return MHTransition(proposed_params), MHState(proposed_params, logp_proposal)\n else\n return MHTransition(current_params), MHState(current_params, state.logp)\n end\nend\n\nfunction compute_log_acceptance_ratio(\n ::RandomWalkMH, state::MHState, ::Vector{Float64}, logp_proposal::Float64\n)\n return min(0, logp_proposal - state.logp)\nend\n\nfunction compute_log_acceptance_ratio(\n sampler::IndependentMH, state::MHState, proposal::Vector{T}, logp_proposal::Float64\n) where {T}\n return min(\n 0,\n logp_proposal - state.logp + logpdf(sampler.proposal_dist, state.params) -\n logpdf(sampler.proposal_dist, proposal),\n )\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"At last, we can proceed to implement a very simple Gibbs sampler.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"struct Gibbs{T<:NamedTuple} <: AbstractMCMC.AbstractSampler\n \"Maps variables to their samplers.\"\n sampler_map::T\nend\n\nstruct GibbsState{TraceNT<:NamedTuple,StateNT<:NamedTuple,SizeNT<:NamedTuple}\n \"Contains the values of all parameters up to the last iteration.\"\n trace::TraceNT\n\n \"Maps parameters to their sampler-specific MCMC states.\"\n mcmc_states::StateNT\n\n \"Maps parameters to their sizes.\"\n variable_sizes::SizeNT\nend\n\nstruct GibbsTransition{ValuesNT<:NamedTuple}\n \"Realizations of the parameters, this is considered a \\\"sample\\\" in the MCMC chain.\"\n values::ValuesNT\nend\n\n\"\"\"\n update_trace(trace::NamedTuple, gibbs_state::GibbsState)\n\nUpdate the trace with the values from the MCMC states of the sub-problems.\n\"\"\"\nfunction update_trace(\n trace::NamedTuple{trace_names}, gibbs_state::GibbsState{TraceNT,StateNT,SizeNT}\n) where {trace_names,TraceNT,StateNT,SizeNT}\n for parameter_variable in fieldnames(StateNT)\n sub_state = gibbs_state.mcmc_states[parameter_variable]\n sub_state_params_values = Base.vec(sub_state)\n reshaped_sub_state_params_values = reshape(\n sub_state_params_values, gibbs_state.variable_sizes[parameter_variable]\n )\n unflattened_sub_state_params = NamedTuple{(parameter_variable,)}((\n reshaped_sub_state_params_values,\n ))\n trace = merge(trace, unflattened_sub_state_params)\n end\n return trace\nend\n\nfunction error_if_not_fully_initialized(\n initial_params::NamedTuple{ParamNames}, sampler::Gibbs{<:NamedTuple{SamplerNames}}\n) where {ParamNames,SamplerNames}\n if Set(ParamNames) != Set(SamplerNames)\n throw(\n ArgumentError(\n \"initial_params must contain all parameters in the model, expected $(SamplerNames), got $(ParamNames)\",\n ),\n )\n end\nend\n\nfunction AbstractMCMC.step(\n rng::Random.AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::Gibbs{Tsamplingmap};\n initial_params::NamedTuple,\n kwargs...,\n) where {Tsamplingmap}\n error_if_not_fully_initialized(initial_params, sampler)\n\n model_parameter_names = fieldnames(Tsamplingmap)\n results = map(model_parameter_names) do parameter_variable\n sub_sampler = sampler.sampler_map[parameter_variable]\n\n variables_to_be_conditioned_on = setdiff(\n model_parameter_names, (parameter_variable,)\n )\n conditioning_variables_values = NamedTuple{Tuple(variables_to_be_conditioned_on)}(\n Tuple([initial_params[g] for g in variables_to_be_conditioned_on])\n )\n\n # LogDensityProblems' `logdensity` function expects a single vector of real numbers\n # `Gibbs` stores the parameters as a named tuple, thus we need to flatten the sub_problem_parameters_values\n # and unflatten after the sampling step\n flattened_sub_problem_parameters_values = vec(initial_params[parameter_variable])\n\n sub_logdensity_model = AbstractMCMC.LogDensityModel(\n AbstractPPL.condition(\n logdensity_model.logdensity, conditioning_variables_values\n ),\n )\n sub_state = last(\n AbstractMCMC.step(\n rng,\n sub_logdensity_model,\n sub_sampler;\n initial_params=flattened_sub_problem_parameters_values,\n kwargs...,\n ),\n )\n (sub_state, size(initial_params[parameter_variable]))\n end\n\n mcmc_states_tuple = first.(results)\n variable_sizes_tuple = last.(results)\n\n gibbs_state = GibbsState(\n initial_params,\n NamedTuple{Tuple(model_parameter_names)}(mcmc_states_tuple),\n NamedTuple{Tuple(model_parameter_names)}(variable_sizes_tuple),\n )\n\n trace = update_trace(NamedTuple(), gibbs_state)\n return GibbsTransition(trace), gibbs_state\nend\n\n# subsequent steps\nfunction AbstractMCMC.step(\n rng::Random.AbstractRNG,\n logdensity_model::AbstractMCMC.LogDensityModel,\n sampler::Gibbs{Tsamplingmap},\n gibbs_state::GibbsState;\n kwargs...,\n) where {Tsamplingmap}\n trace = gibbs_state.trace\n mcmc_states = gibbs_state.mcmc_states\n variable_sizes = gibbs_state.variable_sizes\n\n model_parameter_names = fieldnames(Tsamplingmap)\n mcmc_states = map(model_parameter_names) do parameter_variable\n sub_sampler = sampler.sampler_map[parameter_variable]\n sub_state = mcmc_states[parameter_variable]\n variables_to_be_conditioned_on = setdiff(\n model_parameter_names, (parameter_variable,)\n )\n conditioning_variables_values = NamedTuple{Tuple(variables_to_be_conditioned_on)}(\n Tuple([trace[g] for g in variables_to_be_conditioned_on])\n )\n cond_logdensity = AbstractPPL.condition(\n logdensity_model.logdensity, conditioning_variables_values\n )\n cond_logdensity_model = AbstractMCMC.LogDensityModel(cond_logdensity)\n\n logp = LogDensityProblems.logdensity(cond_logdensity_model, sub_state)\n sub_state = (sub_state)(logp)\n sub_state = last(\n AbstractMCMC.step(\n rng, cond_logdensity_model, sub_sampler, sub_state; kwargs...\n ),\n )\n trace = update_trace(trace, gibbs_state)\n sub_state\n end\n mcmc_states = NamedTuple{Tuple(model_parameter_names)}(mcmc_states)\n\n return GibbsTransition(trace), GibbsState(trace, mcmc_states, variable_sizes)\nend","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We are using NamedTuple to store the mapping between variables and samplers. The order will determine the order of the Gibbs sweeps. A limitation is that exactly one sampler for each variable is required, which means it is less flexible than Gibbs in Turing.jl.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"We uses the AbstractPPL.condition to devide the full model into smaller conditional probability problems. And each conditional probability problem corresponds to a sampler and corresponding state.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"The Gibbs sampler has the same interface as other samplers in AbstractMCMC (we don't implement the above state interface for GibbsState to keep it simple, but it can be implemented similarly).","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"The Gibbs sampler operates in two main phases:","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Initialization:\nSet up initial states for each conditional probability problem.\nIterative Sampling: For each iteration, the sampler performs a sweep over all conditional probability problems:\na. Condition on other variables:\nFix the values of all variables except the current one.\nb. Update current variable:\nRecompute the log probability of the current state, as other variables may have changed:\nUse LogDensityProblems.logdensity(cond_logdensity_model, sub_state) to get the new log probability.\nUpdate the state with sub_state = sub_state(logp) to incorporate the new log probability.\nPerform a sampling step for the current conditional probability problem:\nUse AbstractMCMC.step(rng, cond_logdensity_model, sub_sampler, sub_state; kwargs...) to generate a new state.\nUpdate the global trace:\nExtract parameter values from the new state using Base.vec(new_sub_state).\nIncorporate these values into the overall Gibbs state trace.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"This process allows the Gibbs sampler to iteratively update each variable while conditioning on the others, gradually exploring the joint distribution of all variables.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Now we can use the Gibbs sampler to sample from the hierarchical Normal model.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"First we generate some data,","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"N = 100 # Number of data points\nmu_true = 0.5 # True mean\ntau2_true = 2.0 # True variance\n\nx_data = rand(Normal(mu_true, sqrt(tau2_true)), N)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Then we can create a HierNormal model, with the data we just generated.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"hn = HierNormal((x=x_data,))","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Using Gibbs sampling allows us to use random walk MH for mu and prior MH for tau2, because tau2 has support only on positive real numbers.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"samples = sample(\n hn,\n Gibbs((\n mu=RandomWalkMH(0.3),\n tau2=IndependentMH(product_distribution([InverseGamma(1, 1)])),\n )),\n 10000;\n initial_params=(mu=[0.0], tau2=[1.0]),\n)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"Then we can extract the samples and compute the mean of the samples.","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"mu_samples = [sample.values.mu for sample in samples][20001:end]\ntau2_samples = [sample.values.tau2 for sample in samples][20001:end]\n\nmean(mu_samples)\nmean(tau2_samples)\n(mu_mean, tau2_mean)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"the result should looks like:","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"(4.995812149309413, 1.9372372289677886)","category":"page"},{"location":"gibbs/","page":"state Interface Functions","title":"state Interface Functions","text":"which is close to the true values (5, 2).","category":"page"},{"location":"#AbstractMCMC.jl","page":"Home","title":"AbstractMCMC.jl","text":"","category":"section"},{"location":"","page":"Home","title":"Home","text":"Abstract types and interfaces for Markov chain Monte Carlo methods.","category":"page"},{"location":"","page":"Home","title":"Home","text":"AbstractMCMC defines an interface for sampling and combining Markov chains. It comes with a default sampling algorithm that provides support of progress bars, parallel sampling (multithreaded and multicore), and user-provided callbacks out of the box. Typically developers only have to define the sampling step of their inference method in an iterator-like fashion to make use of this functionality. Additionally, the package defines an iterator and a transducer for sampling Markov chains based on the interface.","category":"page"}] }