Implement GibbsConditional

HISTORY.md

@@ -1,3 +1,11 @@
+# 0.41.2
+
+Add `GibbsConditional`, a "sampler" that can be used to provide analytically known conditional posteriors in a Gibbs sampler.
+
+In Gibbs sampling, some variables are sampled with a component sampler, while holding other variables conditioned to their current values. Usually one e.g. takes turns sampling one variable with HMC and the other with a particle sampler. However, sometimes the posterior distribution of one variable is known analytically, given the conditioned values of other variables. `GibbsConditional` provides a way to implement these analytically known conditional posteriors and use them as component samplers for Gibbs. See the docstring of `GibbsConditional` for details.
+
+Note that `GibbsConditional` used to exist in Turing.jl until v0.36, at which it was removed when the whole Gibbs sampler was rewritten. This reintroduces the same functionality, though with a slightly different interface.
+
 # 0.41.1
 
 The `ModeResult` struct returned by `maximum_a_posteriori` and `maximum_likelihood` can now be wrapped in `InitFromParams()`.

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.41.1"
+version = "0.41.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/src/api.md

@@ -63,6 +63,7 @@ even though [`Prior()`](@ref) is actually defined in the `Turing.Inference` modu
 | `Emcee`              | [`Turing.Inference.Emcee`](@ref)              | Affine-invariant ensemble sampler                                   |
 | `ESS`                | [`Turing.Inference.ESS`](@ref)                | Elliptical slice sampling                                           |
 | `Gibbs`              | [`Turing.Inference.Gibbs`](@ref)              | Gibbs sampling                                                      |
+| `GibbsConditional`   | [`Turing.Inference.GibbsConditional`](@ref)   | Gibbs sampling with analytical conditional posterior distributions  |
 | `HMC`                | [`Turing.Inference.HMC`](@ref)                | Hamiltonian Monte Carlo                                             |
 | `SGLD`               | [`Turing.Inference.SGLD`](@ref)               | Stochastic gradient Langevin dynamics                               |
 | `SGHMC`              | [`Turing.Inference.SGHMC`](@ref)              | Stochastic gradient Hamiltonian Monte Carlo                         |

src/Turing.jl

@@ -102,6 +102,7 @@ export
     Emcee,
     ESS,
     Gibbs,
+    GibbsConditional,
     HMC,
     SGLD,
     SGHMC,

src/mcmc/Inference.jl

@@ -56,6 +56,7 @@ export Hamiltonian,
     ESS,
     Emcee,
     Gibbs,      # classic sampling
+    GibbsConditional,  # conditional sampling
     HMC,
     SGLD,
     PolynomialStepsize,
@@ -430,6 +431,7 @@ include("mh.jl")
 include("is.jl")
 include("particle_mcmc.jl")
 include("gibbs.jl")
+include("gibbs_conditional.jl")
 include("sghmc.jl")
 include("emcee.jl")
 include("prior.jl")

src/mcmc/gibbs_conditional.jl

@@ -0,0 +1,166 @@
+"""
+    GibbsConditional(get_cond_dists)
+
+A Gibbs component sampler that samples variables according to user-provided analytical
+conditional posterior distributions.
+
+When using Gibbs sampling, sometimes one may know the analytical form of the posterior for
+a given variable, given the conditioned values of the other variables. In such cases one can
+use `GibbsConditional` as a component sampler to to sample from these known conditionals
+directly, avoiding any MCMC methods. One does so with
+
+```julia
+sampler = Gibbs(
+    (@varname(var1), @varname(var2)) => GibbsConditional(get_cond_dists),
+    other samplers go here...
+)
+```
+
+Here `get_cond_dists(c::Dict{<:VarName})` should be a function that takes a `Dict` mapping
+the conditioned variables (anything other than `var1` and `var2`) to their values, and
+returns the conditional posterior distributions for `var1` and `var2`. You may, of course,
+have any number of variables being sampled as a block in this manner, we only use two as an
+example. The return value of `get_cond_dists` should be one of the following:
+- A single `Distribution`, if only one variable is being sampled.
+- An `AbstractDict{<:VarName,<:Distribution}` that maps the variables being sampled to their
+  conditional posteriors E.g. `Dict(@varname(var1) => dist1, @varname(var2) => dist2)`.
+- A `NamedTuple` of `Distribution`s, which is like the `AbstractDict` case but can be used
+  if all the variable names are single `Symbol`s, and may be more performant. E.g.
+  `(; var1=dist1, var2=dist2)`.
+
+# Examples
+
+```julia
+# Define a model
+@model function inverse_gdemo(x)
+    precision ~ Gamma(2, inv(3))
+    std = sqrt(1 / precision)
+    m ~ Normal(0, std)
+    for i in 1:length(x)
+        x[i] ~ Normal(m, std)
+    end
+end
+
+# Define analytical conditionals
+function cond_precision(c)
+    a = 2.0
+    b = inv(3)
+    m = c[@varname(m)]
+    x = c[@varname(x)]
+    n = length(x)
+    a_new = a + (n + 1) / 2
+    b_new = b + sum((x[i] - m)^2 for i in 1:n) / 2 + m^2 / 2
+    return Gamma(a_new, 1 / b_new)
+end
+
+function cond_m(c)
+    precision = c[@varname(precision)]
+    x = c[@varname(x)]
+    n = length(x)
+    m_mean = sum(x) / (n + 1)
+    m_var = 1 / (precision * (n + 1))
+    return Normal(m_mean, sqrt(m_var))
+end
+
+# Sample using GibbsConditional
+model = inverse_gdemo([1.0, 2.0, 3.0])
+chain = sample(model, Gibbs(
+    :precision => GibbsConditional(cond_precision),
+    :m => GibbsConditional(cond_m)
+), 1000)
+```
+"""
+struct GibbsConditional{C} <: AbstractSampler
+    get_cond_dists::C
+end
+
+isgibbscomponent(::GibbsConditional) = true
+
+"""
+    build_variable_dict(model::DynamicPPL.Model)
+
+Traverse the context stack of `model` and build a `Dict` of all the variable values that are
+set in GibbsContext, ConditionContext, or FixedContext.
+"""
+function build_variable_dict(model::DynamicPPL.Model)
+    context = model.context
+    cond_nt = DynamicPPL.conditioned(context)
+    fixed_nt = DynamicPPL.fixed(context)
+    # TODO(mhauru) Can we avoid invlinking all the time? Note that this causes a model
+    # evaluation, which may be expensive.
+    global_vi = DynamicPPL.invlink(get_gibbs_global_varinfo(context), model)
+    return merge(
+        DynamicPPL.values_as(global_vi, Dict),
+        Dict(
+            (DynamicPPL.VarName{sym}() => val for (sym, val) in pairs(cond_nt))...,
+            (DynamicPPL.VarName{sym}() => val for (sym, val) in pairs(fixed_nt))...,
+            (DynamicPPL.VarName{sym}() => val for (sym, val) in pairs(model.args))...,
+        ),
+    )
+end
+
+function get_gibbs_global_varinfo(context::DynamicPPL.AbstractContext)
+    return if context isa GibbsContext
+        get_global_varinfo(context)
+    elseif DynamicPPL.NodeTrait(context) isa DynamicPPL.IsParent
+        get_gibbs_global_varinfo(DynamicPPL.childcontext(context))
+    else
+        msg = """No GibbsContext found in context stack. Are you trying to use \
+            GibbsConditional outside of Gibbs?
+            """
+        throw(ArgumentError(msg))
+    end
+end
+
+function initialstep(
+    ::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    ::GibbsConditional,
+    vi::DynamicPPL.AbstractVarInfo;
+    kwargs...,
+)
+    state = DynamicPPL.is_transformed(vi) ? DynamicPPL.invlink(vi, model) : vi
+    # Since GibbsConditional is only used within Gibbs, it does not need to return a
+    # transition.
+    return nothing, state
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler::GibbsConditional,
+    state::DynamicPPL.AbstractVarInfo;
+    kwargs...,
+)
+    # Get all the conditioned variable values from the model context. This is assumed to
+    # include a GibbsContext as part of the context stack.
+    condvals = build_variable_dict(model)
+    conddists = sampler.get_cond_dists(condvals)
+
+    # We support three different kinds of return values for `sample.get_cond_dists`, to make
+    # life easier for the user.
+    if conddists isa AbstractDict
+        for (vn, dist) in conddists
+            state = setindex!!(state, rand(rng, dist), vn)
+        end
+    elseif conddists isa NamedTuple
+        for (vn_sym, dist) in pairs(conddists)
+            vn = VarName{vn_sym}()
+            state = setindex!!(state, rand(rng, dist), vn)
+        end
+    else
+        # Single variable case
+        vn = only(keys(state))
+        state = setindex!!(state, rand(rng, conddists), vn)
+    end
+
+    # Since GibbsConditional is only used within Gibbs, it does not need to return a
+    # transition.
+    return nothing, state
+end
+
+function setparams_varinfo!!(
+    ::DynamicPPL.Model, ::GibbsConditional, ::Any, params::DynamicPPL.AbstractVarInfo
+)
+    return params
+end

test/mcmc/gibbs.jl

@@ -496,7 +496,7 @@ end
 
     @testset "dynamic model with analytical posterior" begin
         # A dynamic model where b ~ Bernoulli determines the dimensionality
-        # When b=0: single parameter θ₁ 
+        # When b=0: single parameter θ₁
         # When b=1: two parameters θ₁, θ₂ where we observe their sum
         @model function dynamic_bernoulli_normal(y_obs=2.0)
             b ~ Bernoulli(0.3)
@@ -575,7 +575,7 @@ end
         #       end
         #   end
         #   sample(f(), Gibbs(:a => PG(10), :x => MH()), 1000)
-        # 
+        #
         # because the number of observations in each particle depends on the value
         # of `a`.
         #