Move abstractmcmc.jl to external_sampler.jl

Author: penelopeysm

src/mcmc/Inference.jl

@@ -123,7 +123,6 @@ end
 # Default Transition #
 ######################
 # Default
-# Extended in contrib/inference/abstractmcmc.jl
 getstats(t) = nothing
 
 abstract type AbstractTransition end
@@ -359,7 +358,6 @@ end
 # Concrete algorithm implementations. #
 #######################################
 
-include("abstractmcmc.jl")
 include("ess.jl")
 include("hmc.jl")
 include("mh.jl")

src/mcmc/abstractmcmc.jl

@@ -65,3 +65,101 @@ function externalsampler(
 )
     return ExternalSampler(sampler, adtype, Val(unconstrained))
 end
+
+struct TuringState{S,M,V,C}
+    state::S
+    ldf::DynamicPPL.LogDensityFunction{M,V,C}
+end
+
+state_to_turing(f::DynamicPPL.LogDensityFunction, state) = TuringState(state, f)
+function transition_to_turing(f::DynamicPPL.LogDensityFunction, transition)
+    # TODO: We should probably rename this `getparams` since it returns something
+    # very different from `Turing.Inference.getparams`.
+    θ = getparams(f.model, transition)
+    varinfo = DynamicPPL.unflatten(f.varinfo, θ)
+    return Transition(f.model, varinfo, transition)
+end
+
+function varinfo(state::TuringState)
+    θ = getparams(state.ldf.model, state.state)
+    # TODO: Do we need to link here first?
+    return DynamicPPL.unflatten(state.ldf.varinfo, θ)
+end
+varinfo(state::AbstractVarInfo) = state
+
+# NOTE: Only thing that depends on the underlying sampler.
+# Something similar should be part of AbstractMCMC at some point:
+# https://github.com/TuringLang/AbstractMCMC.jl/pull/86
+getparams(::DynamicPPL.Model, transition::AdvancedHMC.Transition) = transition.z.θ
+function getparams(model::DynamicPPL.Model, state::AdvancedHMC.HMCState)
+    return getparams(model, state.transition)
+end
+getstats(transition::AdvancedHMC.Transition) = transition.stat
+
+getparams(::DynamicPPL.Model, transition::AdvancedMH.Transition) = transition.params
+
+# TODO: Do we also support `resume`, etc?
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler_wrapper::Sampler{<:ExternalSampler};
+    initial_state=nothing,
+    initial_params=nothing,
+    kwargs...,
+)
+    alg = sampler_wrapper.alg
+    sampler = alg.sampler
+
+    # Initialise varinfo with initial params and link the varinfo if needed.
+    varinfo = DynamicPPL.VarInfo(model)
+    if requires_unconstrained_space(alg)
+        if initial_params !== nothing
+            # If we have initial parameters, we need to set the varinfo before linking.
+            varinfo = DynamicPPL.link(DynamicPPL.unflatten(varinfo, initial_params), model)
+            # Extract initial parameters in unconstrained space.
+            initial_params = varinfo[:]
+        else
+            varinfo = DynamicPPL.link(varinfo, model)
+        end
+    end
+
+    # Construct LogDensityFunction
+    f = DynamicPPL.LogDensityFunction(model, varinfo; adtype=alg.adtype)
+
+    # Then just call `AbstractMCMC.step` with the right arguments.
+    if initial_state === nothing
+        transition_inner, state_inner = AbstractMCMC.step(
+            rng, AbstractMCMC.LogDensityModel(f), sampler; initial_params, kwargs...
+        )
+    else
+        transition_inner, state_inner = AbstractMCMC.step(
+            rng,
+            AbstractMCMC.LogDensityModel(f),
+            sampler,
+            initial_state;
+            initial_params,
+            kwargs...,
+        )
+    end
+    # Update the `state`
+    return transition_to_turing(f, transition_inner), state_to_turing(f, state_inner)
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler_wrapper::Sampler{<:ExternalSampler},
+    state::TuringState;
+    kwargs...,
+)
+    sampler = sampler_wrapper.alg.sampler
+    f = state.ldf
+
+    # Then just call `AdvancedHMC.step` with the right arguments.
+    transition_inner, state_inner = AbstractMCMC.step(
+        rng, AbstractMCMC.LogDensityModel(f), sampler, state.state; kwargs...
+    )
+
+    # Update the `state`
+    return transition_to_turing(f, transition_inner), state_to_turing(f, state_inner)
+end

src/mcmc/external_sampler.jl

@@ -1,4 +1,4 @@
-module AbstractMCMCTests
+module ExternalSamplerTests
 
 using AbstractMCMC: AbstractMCMC
 using AdvancedMH: AdvancedMH

test/mcmc/abstractmcmc.jl test/mcmc/external_sampler.jltest/mcmc/abstractmcmc.jl renamed to test/mcmc/external_sampler.jl

@@ -54,7 +54,7 @@ end
             @timeit_include("mcmc/hmc.jl")
             @timeit_include("mcmc/Inference.jl")
             @timeit_include("mcmc/sghmc.jl")
-            @timeit_include("mcmc/abstractmcmc.jl")
+            @timeit_include("mcmc/external_sampler.jl")
             @timeit_include("mcmc/mh.jl")
             @timeit_include("ext/dynamichmc.jl")
             @timeit_include("mcmc/repeat_sampler.jl")