[no ci] initial updates for InitContext

Author: penelopeysm

Project.toml

@@ -64,7 +64,7 @@ Distributions = "0.25.77"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.37.2"
+DynamicPPL = "0.38"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3, 1"
 Libtask = "0.9.3"
@@ -90,3 +90,6 @@ julia = "1.10.8"
 [extras]
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+
+[sources]
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "breaking"}

src/mcmc/abstractmcmc.jl

@@ -1,9 +1,9 @@
 # TODO: Implement additional checks for certain samplers, e.g.
 # HMC not supporting discrete parameters.
 function _check_model(model::DynamicPPL.Model)
-    # TODO(DPPL0.38/penelopeysm): use InitContext
-    spl_model = DynamicPPL.contextualize(model, DynamicPPL.SamplingContext(model.context))
-    return DynamicPPL.check_model(spl_model, VarInfo(); error_on_failure=true)
+    new_context = DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext())
+    new_model = DynamicPPL.contextualize(model, new_context)
+    return DynamicPPL.check_model(new_model, VarInfo(); error_on_failure=true)
 end
 function _check_model(model::DynamicPPL.Model, alg::InferenceAlgorithm)
     return _check_model(model)

src/mcmc/emcee.jl

@@ -46,21 +46,16 @@ function AbstractMCMC.step(
 
     # Sample from the prior
     n = spl.alg.ensemble.n_walkers
-    vis = [VarInfo(rng, model, SampleFromPrior()) for _ in 1:n]
+    vis = [VarInfo(rng, model) for _ in 1:n]
 
     # Update the parameters if provided.
     if initial_params !== nothing
-        length(initial_params) == n ||
-            throw(ArgumentError("initial parameters have to be specified for each walker"))
-        vis = map(vis, initial_params) do vi, init
-            # TODO(DPPL0.38/penelopeysm) This whole thing can be replaced with init!!
-            vi = DynamicPPL.initialize_parameters!!(vi, init, model)
-
-            # Update log joint probability.
-            spl_model = DynamicPPL.contextualize(
-                model, DynamicPPL.SamplingContext(rng, SampleFromPrior(), model.context)
-            )
-            last(DynamicPPL.evaluate!!(spl_model, vi))
+        if !(initial_params isa AbstractVector && length(initial_params) == n)
+            err_msg = "initial_params for `Emcee` must be a vector of initialisation strategies, with length equal to the number of walkers ($n)"
+            throw(ArgumentError(err_msg))
+        end
+        vis = map(vis, initial_params) do vi, strategy
+            DynamicPPL.init!!(rng, model, vi, strategy)
         end
     end
 

src/mcmc/ess.jl

@@ -82,23 +82,12 @@ EllipticalSliceSampling.isgaussian(::Type{<:ESSPrior}) = true
 
 # Only define out-of-place sampling
 function Base.rand(rng::Random.AbstractRNG, p::ESSPrior)
-    varinfo = p.varinfo
-    # TODO: Surely there's a better way of doing this now that we have `SamplingContext`?
-    # TODO(DPPL0.38/penelopeysm): This can be replaced with `init!!(p.model,
-    # p.varinfo, PriorInit())` after TuringLang/DynamicPPL.jl#984. The reason
-    # why we had to use the 'del' flag before this was because
-    # SampleFromPrior() wouldn't overwrite existing variables.
-    # The main problem I'm rather unsure about is ESS-within-Gibbs. The
-    # current implementation I think makes sure to only resample the variables
-    # that 'belong' to the current ESS sampler. InitContext on the other hand
-    # would resample all variables in the model (??) Need to think about this
-    # carefully.
-    vns = keys(varinfo)
-    for vn in vns
-        set_flag!(varinfo, vn, "del")
-    end
-    p.model(rng, varinfo)
-    return varinfo[:]
+    # TODO(penelopeysm/DPPL 0.38) The main problem I'm rather unsure about is
+    # ESS-within-Gibbs. The current implementation I think makes sure to only resample the
+    # variables that 'belong' to the current ESS sampler. InitContext on the other hand
+    # would resample all variables in the model (??) Need to think about this carefully.
+    _, vi = DynamicPPL.init!!(p.model, p.varinfo, DynamicPPL.InitFromPrior())
+    return vi[:]
 end
 
 # Mean of prior distribution

src/mcmc/gibbs.jl

@@ -47,7 +47,7 @@ A context used in the implementation of the Turing.jl Gibbs sampler.
 There will be one `GibbsContext` for each iteration of a component sampler.
 
 `target_varnames` is a a tuple of `VarName`s that the current component sampler
-is sampling. For those `VarName`s, `GibbsContext` will just pass `tilde_assume`
+is sampling. For those `VarName`s, `GibbsContext` will just pass `tilde_assume!!`
 calls to its child context. For other variables, their values will be fixed to
 the values they have in `global_varinfo`.
 
@@ -140,7 +140,7 @@ function is_target_varname(context::GibbsContext, vns::AbstractArray{<:VarName})
 end
 
 # Tilde pipeline
-function DynamicPPL.tilde_assume(context::GibbsContext, right, vn, vi)
+function DynamicPPL.tilde_assume!!(context::GibbsContext, right, vn, vi)
     child_context = DynamicPPL.childcontext(context)
 
     # Note that `child_context` may contain `PrefixContext`s -- in which case
@@ -175,7 +175,7 @@ function DynamicPPL.tilde_assume(context::GibbsContext, right, vn, vi)
 
     return if is_target_varname(context, vn)
         # Fall back to the default behavior.
-        DynamicPPL.tilde_assume(child_context, right, vn, vi)
+        DynamicPPL.tilde_assume!!(child_context, right, vn, vi)
     elseif has_conditioned_gibbs(context, vn)
         # This branch means that a different sampler is supposed to handle this
         # variable. From the perspective of this sampler, this variable is
@@ -191,9 +191,10 @@ function DynamicPPL.tilde_assume(context::GibbsContext, right, vn, vi)
         # presumably a new variable that should be sampled from its prior. We need to add
         # this new variable to the global `varinfo` of the context, but not to the local one
         # being used by the current sampler.
-        value, new_global_vi = DynamicPPL.tilde_assume(
-            child_context,
-            DynamicPPL.SampleFromPrior(),
+        value, new_global_vi = DynamicPPL.tilde_assume!!(
+            # child_context might be a PrefixContext so we have to be careful to not
+            # overwrite it.
+            DynamicPPL.setleafcontext(child_context, DynamicPPL.InitContext()),
             right,
             vn,
             get_global_varinfo(context),
@@ -204,7 +205,7 @@ function DynamicPPL.tilde_assume(context::GibbsContext, right, vn, vi)
 end
 
 # As above but with an RNG.
-function DynamicPPL.tilde_assume(
+function DynamicPPL.tilde_assume!!(
     rng::Random.AbstractRNG, context::GibbsContext, sampler, right, vn, vi
 )
     # See comment in the above, rng-less version of this method for an explanation.
@@ -215,7 +216,7 @@ function DynamicPPL.tilde_assume(
         # This branch means that that `sampler` is supposed to handle
         # this variable. We can thus use its default behaviour, with
         # the 'local' sampler-specific VarInfo.
-        DynamicPPL.tilde_assume(rng, child_context, sampler, right, vn, vi)
+        DynamicPPL.tilde_assume!!(rng, child_context, sampler, right, vn, vi)
     elseif has_conditioned_gibbs(context, vn)
         # This branch means that a different sampler is supposed to handle this
         # variable. From the perspective of this sampler, this variable is
@@ -231,10 +232,10 @@ function DynamicPPL.tilde_assume(
         # presumably a new variable that should be sampled from its prior. We need to add
         # this new variable to the global `varinfo` of the context, but not to the local one
         # being used by the current sampler.
-        value, new_global_vi = DynamicPPL.tilde_assume(
-            rng,
-            child_context,
-            DynamicPPL.SampleFromPrior(),
+        value, new_global_vi = DynamicPPL.tilde_assume!!(
+            # child_context might be a PrefixContext so we have to be careful to not
+            # overwrite it.
+            DynamicPPL.setleafcontext(child_context, DynamicPPL.InitContext(rng)),
             right,
             vn,
             get_global_varinfo(context),

src/mcmc/mh.jl

@@ -329,13 +329,11 @@ function propose!!(
     prev_trans = AMH.Transition(vt, prev_state.logjoint_internal, false)
 
     # Make a new transition.
-    spl_model = DynamicPPL.contextualize(
-        model, DynamicPPL.SamplingContext(rng, spl, model.context)
-    )
+    model = DynamicPPL.setleafcontext(model, MHContext(rng))
     densitymodel = AMH.DensityModel(
         Base.Fix1(
             LogDensityProblems.logdensity,
-            DynamicPPL.LogDensityFunction(spl_model, DynamicPPL.getlogjoint_internal, vi),
+            DynamicPPL.LogDensityFunction(model, DynamicPPL.getlogjoint_internal, vi),
         ),
     )
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
@@ -366,13 +364,11 @@ function propose!!(
     prev_trans = AMH.Transition(vals, prev_state.logjoint_internal, false)
 
     # Make a new transition.
-    spl_model = DynamicPPL.contextualize(
-        model, DynamicPPL.SamplingContext(rng, spl, model.context)
-    )
+    model = DynamicPPL.setleafcontext(model, MHContext(rng))
     densitymodel = AMH.DensityModel(
         Base.Fix1(
             LogDensityProblems.logdensity,
-            DynamicPPL.LogDensityFunction(spl_model, DynamicPPL.getlogjoint_internal, vi),
+            DynamicPPL.LogDensityFunction(model, DynamicPPL.getlogjoint_internal, vi),
         ),
     )
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
@@ -410,13 +406,20 @@ function AbstractMCMC.step(
     return Transition(model, new_state.varinfo, nothing), new_state
 end
 
-####
-#### Compiler interface, i.e. tilde operators.
-####
-function DynamicPPL.assume(
-    rng::Random.AbstractRNG, spl::Sampler{<:MH}, dist::Distribution, vn::VarName, vi
+struct MHContext{R} <: DynamicPPL.AbstractContext
+    rng::R
+end
+DynamicPPL.NodeTrait(::MHContext) = DynamicPPL.IsLeaf()
+
+function DynamicPPL.tilde_assume!!(
+    context::MHContext, right::Distribution, vn::VarName, vi::AbstractVarInfo
 )
-    # Just defer to `SampleFromPrior`.
-    retval = DynamicPPL.assume(rng, SampleFromPrior(), dist, vn, vi)
-    return retval
+    # Allow MH to sample new variables from the prior if it's not already present in the
+    # VarInfo.
+    dispatch_ctx = if haskey(vi, vn)
+        DynamicPPL.DefaultContext()
+    else
+        DynamicPPL.InitContext(context.rng, DynamicPPL.InitFromPrior())
+    end
+    return DynamicPPL.tilde_assume!!(dispatch_ctx, right, vn, vi)
 end

src/mcmc/particle_mcmc.jl

@@ -446,6 +446,9 @@ function set_trace_local_varinfo_maybe(vi::AbstractVarInfo)
     return nothing
 end
 
+# TODO(penelopeysm / DPPL 0.38): Figure this out
+struct ParticleMCMCContext <: DynamicPPL.AbstractContext end
+
 function DynamicPPL.assume(
     rng, ::Sampler{<:Union{PG,SMC}}, dist::Distribution, vn::VarName, vi::AbstractVarInfo
 )

src/mcmc/prior.jl

@@ -12,11 +12,6 @@ function AbstractMCMC.step(
     state=nothing;
     kwargs...,
 )
-    # TODO(DPPL0.38/penelopeysm): replace with init!!
-    sampling_model = DynamicPPL.contextualize(
-        model, DynamicPPL.SamplingContext(rng, DynamicPPL.SampleFromPrior(), model.context)
-    )
-    vi = VarInfo()
     vi = DynamicPPL.setaccs!!(
         vi,
         (
@@ -25,6 +20,6 @@ function AbstractMCMC.step(
             DynamicPPL.LogLikelihoodAccumulator(),
         ),
     )
-    _, vi = DynamicPPL.evaluate!!(sampling_model, vi)
+    _, vi = DynamicPPL.init!!(model, vi, DynamicPPL.InitFromPrior())
     return Transition(model, vi, nothing; reevaluate=false), nothing
 end

src/optimisation/Optimisation.jl

@@ -507,10 +507,9 @@ function estimate_mode(
     kwargs...,
 )
     if check_model
-        spl_model = DynamicPPL.contextualize(
-            model, DynamicPPL.SamplingContext(model.context)
-        )
-        DynamicPPL.check_model(spl_model, DynamicPPL.VarInfo(); error_on_failure=true)
+        new_context = DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext())
+        new_model = DynamicPPL.contextualize(model, new_context)
+        DynamicPPL.check_model(new_model, DynamicPPL.VarInfo(); error_on_failure=true)
     end
 
     constraints = ModeEstimationConstraints(lb, ub, cons, lcons, ucons)

test/mcmc/external_sampler.jl

@@ -136,40 +136,6 @@ function initialize_mh_rw(model)
     return AdvancedMH.RWMH(MvNormal(Zeros(d), 0.1 * I))
 end
 
-# TODO: Should this go somewhere else?
-# Convert a model into a `Distribution` to allow usage as a proposal in AdvancedMH.jl.
-struct ModelDistribution{M<:DynamicPPL.Model,V<:DynamicPPL.VarInfo} <:
-       ContinuousMultivariateDistribution
-    model::M
-    varinfo::V
-end
-function ModelDistribution(model::DynamicPPL.Model)
-    return ModelDistribution(model, DynamicPPL.VarInfo(model))
-end
-
-Base.length(d::ModelDistribution) = length(d.varinfo[:])
-function Distributions._logpdf(d::ModelDistribution, x::AbstractVector)
-    return logprior(d.model, DynamicPPL.unflatten(d.varinfo, x))
-end
-function Distributions._rand!(
-    rng::Random.AbstractRNG, d::ModelDistribution, x::AbstractVector{<:Real}
-)
-    model = d.model
-    varinfo = deepcopy(d.varinfo)
-    for vn in keys(varinfo)
-        DynamicPPL.set_flag!(varinfo, vn, "del")
-    end
-    DynamicPPL.evaluate!!(model, varinfo, DynamicPPL.SamplingContext(rng))
-    x .= varinfo[:]
-    return x
-end
-
-function initialize_mh_with_prior_proposal(model)
-    return AdvancedMH.MetropolisHastings(
-        AdvancedMH.StaticProposal(ModelDistribution(model))
-    )
-end
-
 function test_initial_params(
     model, sampler, initial_params=DynamicPPL.VarInfo(model)[:]; kwargs...
 )
@@ -268,28 +234,6 @@ end
                 @test isapprox(logpdf.(Normal(), chn[:x]), chn[:lp])
             end
         end
-
-        # NOTE: Broken because MH doesn't really follow the `logdensity` interface, but calls
-        # it with `NamedTuple` instead of `AbstractVector`.
-        # @testset "MH with prior proposal" begin
-        #     @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
-        #         sampler = initialize_mh_with_prior_proposal(model);
-        #         sampler_ext = DynamicPPL.Sampler(externalsampler(sampler; unconstrained=false))
-        #         @testset "initial_params" begin
-        #             test_initial_params(model, sampler_ext)
-        #         end
-        #         @testset "inference" begin
-        #             DynamicPPL.TestUtils.test_sampler(
-        #                 [model],
-        #                 sampler_ext,
-        #                 10_000;
-        #                 discard_initial=1_000,
-        #                 rtol=0.2,
-        #                 sampler_name="AdvancedMH"
-        #             )
-        #         end
-        #     end
-        # end
     end
 end