[no ci] initial updates for InitContext

Author: penelopeysm

HISTORY.md

@@ -1,5 +1,11 @@
 # 0.41.0
 
+## DynamicPPL 0.38
+
+Lorem ipsum dynamicppl sit amet
+
+## Initial step in MCMC sampling
+
 HMC and NUTS samplers no longer take an extra single step before starting the chain.
 This means that if you do not discard any samples at the start, the first sample will be the initial parameters (which may be user-provided).
 

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,19 @@ 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{<:DynamicPPL.AbstractInitStrategy} &&
+            length(initial_params) == n
+        )
+            err_msg = "initial_params for `Emcee` must be a vector of `DynamicPPL.AbstractInitStrategy`, 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<:AbstractRNG} <: 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

@@ -508,10 +508,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)