Remove Sampler and move its interface to Turing

Author: penelopeysm

Project.toml

@@ -92,4 +92,4 @@ DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 
 [sources]
-DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "breaking"}
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/no-sampler"}

ext/TuringDynamicHMCExt.jl

@@ -44,10 +44,10 @@ struct DynamicNUTSState{L,V<:DynamicPPL.AbstractVarInfo,C,M,S}
     stepsize::S
 end
 
-function DynamicPPL.initialstep(
+function Turing.Inference.initialstep(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
-    spl::DynamicPPL.Sampler{<:DynamicNUTS},
+    spl::DynamicNUTS,
     vi::DynamicPPL.AbstractVarInfo;
     kwargs...,
 )
@@ -59,7 +59,7 @@ function DynamicPPL.initialstep(
 
     # Define log-density function.
     ℓ = DynamicPPL.LogDensityFunction(
-        model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.alg.adtype
+        model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.adtype
     )
 
     # Perform initial step.
@@ -80,14 +80,14 @@ end
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
-    spl::DynamicPPL.Sampler{<:DynamicNUTS},
+    spl::DynamicNUTS,
     state::DynamicNUTSState;
     kwargs...,
 )
     # Compute next sample.
     vi = state.vi
     ℓ = state.logdensity
-    steps = DynamicHMC.mcmc_steps(rng, spl.alg.sampler, state.metric, ℓ, state.stepsize)
+    steps = DynamicHMC.mcmc_steps(rng, spl.sampler, state.metric, ℓ, state.stepsize)
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
     # Create next sample and state.

src/Turing.jl

@@ -160,6 +160,8 @@ export
     maximum_a_posteriori,
     maximum_likelihood,
     MAP,
-    MLE
+    MLE,
+    # Chain save/resume
+    loadstate
 
 end

src/mcmc/Inference.jl

@@ -22,7 +22,6 @@ using DynamicPPL:
     getsym,
     getdist,
     Model,
-    Sampler,
     DefaultContext
 using Distributions, Libtask, Bijectors
 using DistributionsAD: VectorOfMultivariate
@@ -50,8 +49,7 @@ import Random
 import MCMCChains
 import StatsBase: predict
 
-export InferenceAlgorithm,
-    Hamiltonian,
+export Hamiltonian,
     StaticHamiltonian,
     AdaptiveHamiltonian,
     MH,
@@ -71,15 +69,16 @@ export InferenceAlgorithm,
     RepeatSampler,
     Prior,
     predict,
-    externalsampler
+    externalsampler,
+    init_strategy,
+    loadstate
 
-###############################################
-# Abstract interface for inference algorithms #
-###############################################
-
-const TURING_CHAIN_TYPE = MCMCChains.Chains
+#########################################
+# Generic AbstractMCMC methods dispatch #
+#########################################
 
-include("algorithm.jl")
+const DEFAULT_CHAIN_TYPE = MCMCChains.Chains
+include("abstractmcmc.jl")
 
 ####################
 # Sampler wrappers #
@@ -312,8 +311,8 @@ getlogevidence(transitions, sampler, state) = missing
 # Default MCMCChains.Chains constructor.
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:Transition},
-    model::AbstractModel,
-    spl::Union{Sampler{<:InferenceAlgorithm},RepeatSampler},
+    model::DynamicPPL.Model,
+    spl::AbstractSampler,
     state,
     chain_type::Type{MCMCChains.Chains};
     save_state=false,
@@ -374,8 +373,8 @@ end
 
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:Transition},
-    model::AbstractModel,
-    spl::Union{Sampler{<:InferenceAlgorithm},RepeatSampler},
+    model::DynamicPPL.Model,
+    spl::AbstractSampler,
     state,
     chain_type::Type{Vector{NamedTuple}};
     kwargs...,
@@ -416,7 +415,7 @@ function group_varnames_by_symbol(vns)
     return d
 end
 
-function save(c::MCMCChains.Chains, spl::Sampler, model, vi, samples)
+function save(c::MCMCChains.Chains, spl::AbstractSampler, model, vi, samples)
     nt = NamedTuple{(:sampler, :model, :vi, :samples)}((spl, model, deepcopy(vi), samples))
     return setinfo(c, merge(nt, c.info))
 end
@@ -435,18 +434,12 @@ include("sghmc.jl")
 include("emcee.jl")
 include("prior.jl")
 
-#################################################
-# Generic AbstractMCMC methods dispatch #
-#################################################
-
-include("abstractmcmc.jl")
-
 ################
 # Typing tools #
 ################
 
 function DynamicPPL.get_matching_type(
-    spl::Sampler{<:Union{PG,SMC}}, vi, ::Type{TV}
+    spl::Union{PG,SMC}, vi, ::Type{TV}
 ) where {T,N,TV<:Array{T,N}}
     return Array{T,N}
 end

src/mcmc/abstractmcmc.jl

@@ -4,38 +4,98 @@ function _check_model(model::DynamicPPL.Model)
     new_model = DynamicPPL.setleafcontext(model, DynamicPPL.InitContext())
     return DynamicPPL.check_model(new_model, VarInfo(); error_on_failure=true)
 end
-function _check_model(model::DynamicPPL.Model, alg::InferenceAlgorithm)
+function _check_model(model::DynamicPPL.Model, ::AbstractSampler)
     return _check_model(model)
 end
 
+"""
+    Turing.Inference.init_strategy(spl::AbstractSampler)
+
+Get the default initialization strategy for a given sampler `spl`, i.e. how initial
+parameters for sampling are chosen if not specified by the user. By default, this is
+`InitFromPrior()`, which samples initial parameters from the prior distribution.
+"""
+init_strategy(::AbstractSampler) = DynamicPPL.InitFromPrior()
+
+"""
+    _convert_initial_params(initial_params)
+
+Convert `initial_params` to a `DynamicPPl.AbstractInitStrategy` if it is not already one, or
+throw a useful error message.
+"""
+_convert_initial_params(initial_params::DynamicPPL.AbstractInitStrategy) = initial_params
+function _convert_initial_params(nt::NamedTuple)
+    @info "Using a NamedTuple for `initial_params` will be deprecated in a future release. Please use `InitFromParams(namedtuple)` instead."
+    return InitFromParams(nt)
+end
+function _convert_initial_params(d::AbstractDict{<:VarName})
+    @info "Using a Dict for `initial_params` will be deprecated in a future release. Please use `InitFromParams(dict)` instead."
+    return InitFromParams(d)
+end
+function _convert_initial_params(::AbstractVector{<:Real})
+    errmsg = "`initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or ideally a `DynamicPPL.AbstractInitStrategy`. Using a vector of parameters for `initial_params` is no longer supported. Please see https://turinglang.org/docs/usage/sampling-options/#specifying-initial-parameters for details on how to update your code."
+    throw(ArgumentError(errmsg))
+end
+function _convert_initial_params(@nospecialize(_::Any))
+    errmsg = "`initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or a `DynamicPPL.AbstractInitStrategy`."
+    throw(ArgumentError(errmsg))
+end
+
+"""
+    default_varinfo(rng, model, sampler)
+
+Return a default varinfo object for the given `model` and `sampler`.
+The default method for this returns a NTVarInfo (i.e. 'typed varinfo').
+"""
+function default_varinfo(
+    rng::Random.AbstractRNG, model::DynamicPPL.Model, ::AbstractSampler
+)
+    # Note that in `AbstractMCMC.step`, the values in the varinfo returned here are
+    # immediately overwritten by a subsequent call to `init!!`. The reason why we
+    # _do_ create a varinfo with parameters here (as opposed to simply returning
+    # an empty `typed_varinfo(VarInfo())`) is to avoid issues where pushing to an empty
+    # typed VarInfo would fail. This can happen if two VarNames have different types
+    # but share the same symbol (e.g. `x.a` and `x.b`).
+    # TODO(mhauru) Fix push!! to work with arbitrary lens types, and then remove the arguments
+    # and return an empty VarInfo instead.
+    return DynamicPPL.typed_varinfo(VarInfo(rng, model))
+end
+
 #########################################
 # Default definitions for the interface #
 #########################################
 
-const DEFAULT_CHAIN_TYPE = MCMCChains.Chains
-
 function AbstractMCMC.sample(
-    model::AbstractModel, alg::InferenceAlgorithm, N::Integer; kwargs...
+    model::DynamicPPL.Model, spl::AbstractSampler, N::Integer; kwargs...
 )
-    return AbstractMCMC.sample(Random.default_rng(), model, alg, N; kwargs...)
+    return AbstractMCMC.sample(Random.default_rng(), model, spl, N; kwargs...)
 end
 
 function AbstractMCMC.sample(
     rng::AbstractRNG,
-    model::AbstractModel,
-    alg::InferenceAlgorithm,
+    model::DynamicPPL.Model,
+    spl::AbstractSampler,
     N::Integer;
+    initial_params=init_strategy(spl),
     check_model::Bool=true,
     chain_type=DEFAULT_CHAIN_TYPE,
     kwargs...,
 )
-    check_model && _check_model(model, alg)
-    return AbstractMCMC.sample(rng, model, Sampler(alg), N; chain_type, kwargs...)
+    check_model && _check_model(model, spl)
+    return AbstractMCMC.sample(
+        rng,
+        model,
+        spl,
+        N;
+        initial_params=_convert_initial_params(initial_params),
+        chain_type,
+        kwargs...,
+    )
 end
 
 function AbstractMCMC.sample(
-    model::AbstractModel,
-    alg::InferenceAlgorithm,
+    model::DynamicPPL.Model,
+    alg::AbstractSampler,
     ensemble::AbstractMCMC.AbstractMCMCEnsemble,
     N::Integer,
     n_chains::Integer;
@@ -47,18 +107,53 @@ function AbstractMCMC.sample(
 end
 
 function AbstractMCMC.sample(
-    rng::AbstractRNG,
-    model::AbstractModel,
-    alg::InferenceAlgorithm,
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    spl::AbstractSampler,
     ensemble::AbstractMCMC.AbstractMCMCEnsemble,
     N::Integer,
     n_chains::Integer;
     chain_type=DEFAULT_CHAIN_TYPE,
     check_model::Bool=true,
+    initial_params=init_strategy(spl),
     kwargs...,
 )
-    check_model && _check_model(model, alg)
+    check_model && _check_model(model, spl)
     return AbstractMCMC.sample(
-        rng, model, Sampler(alg), ensemble, N, n_chains; chain_type, kwargs...
+        rng,
+        model,
+        spl,
+        ensemble,
+        N,
+        n_chains;
+        chain_type,
+        initial_params=map(_convert_initial_params, initial_params),
+        kwargs...,
     )
 end
+
+loadstate(c::MCMCChains.Chains) = c.info.samplerstate
+
+# TODO(penelopeysm): Remove initialstep and generalise MCMC sampling procedures
+function initialstep end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    spl::AbstractSampler;
+    initial_params,
+    kwargs...,
+)
+    # Generate the default varinfo. Note that any parameters inside this varinfo
+    # will be immediately overwritten by the next call to `init!!`.
+    vi = default_varinfo(rng, model, spl)
+
+    # Fill it with initial parameters. Note that, if `InitFromParams` is used, the
+    # parameters provided must be in unlinked space (when inserted into the
+    # varinfo, they will be adjusted to match the linking status of the
+    # varinfo).
+    _, vi = DynamicPPL.init!!(rng, model, vi, initial_params)
+
+    # Call the actual function that does the first step.
+    return initialstep(rng, model, spl, vi; initial_params, kwargs...)
+end

src/mcmc/algorithm.jl

@@ -13,7 +13,7 @@ Foreman-Mackey, D., Hogg, D. W., Lang, D., & Goodman, J. (2013).
 emcee: The MCMC Hammer. Publications of the Astronomical Society of the
 Pacific, 125 (925), 306. https://doi.org/10.1086/670067
 """
-struct Emcee{E<:AMH.Ensemble} <: InferenceAlgorithm
+struct Emcee{E<:AMH.Ensemble} <: AbstractSampler
     ensemble::E
 end
 
@@ -33,23 +33,20 @@ end
 
 # Utility function to tetrieve the number of walkers
 _get_n_walkers(e::Emcee) = e.ensemble.n_walkers
-_get_n_walkers(spl::Sampler{<:Emcee}) = _get_n_walkers(spl.alg)
 
 # Because Emcee expects n_walkers initialisations, we need to override this
-function DynamicPPL.init_strategy(spl::Sampler{<:Emcee})
+function Turing.Inference.init_strategy(spl::Emcee)
     return fill(DynamicPPL.InitFromPrior(), _get_n_walkers(spl))
 end
-# TODO(penelopeysm / DPPL 0.38) This is type piracy (!!!) The function
-# `_convert_initial_params` will be moved to Turing soon, and this piracy SHOULD be removed
-# in https://github.com/TuringLang/Turing.jl/pull/2689, PLEASE make sure it is!
-function DynamicPPL._convert_initial_params(
+# We also have to explicitly allow this or else it will error...
+function Turing.Inference._convert_initial_params(
     x::AbstractVector{<:DynamicPPL.AbstractInitStrategy}
 )
     return x
 end
 
 function AbstractMCMC.step(
-    rng::Random.AbstractRNG, model::Model, spl::Sampler{<:Emcee}; initial_params, kwargs...
+    rng::Random.AbstractRNG, model::Model, spl::Emcee; initial_params, kwargs...
 )
     # Sample from the prior
     n = _get_n_walkers(spl)
@@ -83,7 +80,7 @@ function AbstractMCMC.step(
 end
 
 function AbstractMCMC.step(
-    rng::AbstractRNG, model::Model, spl::Sampler{<:Emcee}, state::EmceeState; kwargs...
+    rng::AbstractRNG, model::Model, spl::Emcee, state::EmceeState; kwargs...
 )
     # Generate a log joint function.
     vi = state.vi
@@ -95,7 +92,7 @@ function AbstractMCMC.step(
     )
 
     # Compute the next states.
-    t, states = AbstractMCMC.step(rng, densitymodel, spl.alg.ensemble, state.states)
+    t, states = AbstractMCMC.step(rng, densitymodel, spl.ensemble, state.states)
 
     # Compute the next transition and state.
     transition = map(states) do _state
@@ -110,7 +107,7 @@ end
 function AbstractMCMC.bundle_samples(
     samples::Vector{<:Vector},
     model::AbstractModel,
-    spl::Sampler{<:Emcee},
+    spl::Emcee,
     state::EmceeState,
     chain_type::Type{MCMCChains.Chains};
     save_state=false,