Progress in DPPL 0.37 compat for particle MCMC

Author: mhauru

src/mcmc/particle_mcmc.jl

@@ -33,10 +33,14 @@ end
 function AdvancedPS.advance!(
     trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false
 )
-    # Make sure we load/reset the rng in the new replaying mechanism
-    trace = Accessors.@set trace.model.f.varinfo = DynamicPPL.increment_num_produce!!(
-        trace.model.f.varinfo
+    # We want to increment num produce for the VarInfo stored in the trace. The trace is
+    # mutable, so we create a new model with the incremented VarInfo and set it in the trace
+    model = trace.model
+    model = Accessors.@set model.f.varinfo = DynamicPPL.increment_num_produce!!(
+        model.f.varinfo
     )
+    trace.model = model
+    # Make sure we load/reset the rng in the new replaying mechanism
     isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng)
     score = consume(trace.model.ctask)
     if score === nothing
@@ -55,10 +59,6 @@ function AdvancedPS.reset_model(trace::TracedModel)
     return Accessors.@set trace.varinfo = DynamicPPL.reset_num_produce!!(trace.varinfo)
 end
 
-function AdvancedPS.reset_logprob!(trace::TracedModel)
-    return Accessors.@set trace.model.varinfo = DynamicPPL.resetlogp!!(trace.model.varinfo)
-end
-
 function Libtask.TapedTask(taped_globals, model::TracedModel; kwargs...)
     return Libtask.TapedTask(
         taped_globals, model.evaluator[1], model.evaluator[2:end]...; kwargs...
@@ -390,78 +390,124 @@ function DynamicPPL.use_threadsafe_eval(
     return false
 end
 
-function trace_local_varinfo_maybe(varinfo)
-    try
-        trace = Libtask.get_taped_globals(Any).other
-        return (trace === nothing ? varinfo : trace.model.f.varinfo)::AbstractVarInfo
+"""
+    get_trace_local_varinfo_maybe(vi::AbstractVarInfo)
+
+Get the `Trace` local varinfo if one exists.
+
+If executed within a `TapedTask`, return the `varinfo` stored in the "taped globals" of the
+task, otherwise return `vi`.
+"""
+function get_trace_local_varinfo_maybe(varinfo::AbstractVarInfo)
+    trace = try
+        Libtask.get_taped_globals(Any).other
     catch e
-        # NOTE: this heuristic allows Libtask evaluating a model outside a `Trace`.
-        if e == KeyError(:task_variable)
-            return varinfo
-        else
-            rethrow(e)
-        end
+        e == KeyError(:task_variable) ? nothing : rethrow(e)
     end
+    return (trace === nothing ? varinfo : trace.model.f.varinfo)::AbstractVarInfo
 end
 
-function trace_local_rng_maybe(rng::Random.AbstractRNG)
-    try
-        return Libtask.get_taped_globals(Any).rng
+"""
+    get_trace_local_varinfo_maybe(rng::Random.AbstractRNG)
+
+Get the `Trace` local rng if one exists.
+
+If executed within a `TapedTask`, return the `rng` stored in the "taped globals" of the
+task, otherwise return `vi`.
+"""
+function get_trace_local_rng_maybe(rng::Random.AbstractRNG)
+    return try
+        Libtask.get_taped_globals(Any).rng
     catch e
-        # NOTE: this heuristic allows Libtask evaluating a model outside a `Trace`.
-        if e == KeyError(:task_variable)
-            return rng
-        else
-            rethrow(e)
-        end
+        e == KeyError(:task_variable) ? rng : rethrow(e)
     end
 end
 
-# TODO(DPPL0.37/penelopeysm) The whole tilde pipeline for particle MCMC needs to be
-# thoroughly fixed.
+"""
+    set_trace_local_varinfo_maybe(vi::AbstractVarInfo)
+
+Set the `Trace` local varinfo if executing within a `Trace`. Return `nothing`.
+
+If executed within a `TapedTask`, set the `varinfo` stored in the "taped globals" of the
+task. Otherwise do nothing.
+"""
+function set_trace_local_varinfo_maybe(vi::AbstractVarInfo)
+    # TODO(mhauru) This should be done in a try-catch block, as in the commented out code.
+    # However, Libtask currently can't handle this block.
+    trace = #try
+        Libtask.get_taped_globals(Any).other
+    # catch e
+    #     e == KeyError(:task_variable) ? nothing : rethrow(e)
+    # end
+    if trace !== nothing
+        model = trace.model
+        model = Accessors.@set model.f.varinfo = vi
+        trace.model = model
+    end
+    return nothing
+end
+
 function DynamicPPL.assume(
-    rng, ::Sampler{<:Union{PG,SMC}}, dist::Distribution, vn::VarName, _vi::AbstractVarInfo
+    rng, ::Sampler{<:Union{PG,SMC}}, dist::Distribution, vn::VarName, vi::AbstractVarInfo
 )
-    vi = trace_local_varinfo_maybe(_vi)
-    trng = trace_local_rng_maybe(rng)
+    arg_vi_id = objectid(vi)
+    vi = get_trace_local_varinfo_maybe(vi)
+    using_local_vi = objectid(vi) == arg_vi_id
+
+    trng = get_trace_local_rng_maybe(rng)
 
     if ~haskey(vi, vn)
         r = rand(trng, dist)
-        push!!(vi, vn, r, dist)
+        vi = push!!(vi, vn, r, dist)
     elseif DynamicPPL.is_flagged(vi, vn, "del")
         DynamicPPL.unset_flag!(vi, vn, "del") # Reference particle parent
         r = rand(trng, dist)
         vi[vn] = DynamicPPL.tovec(r)
+        # TODO(mhauru):
+        # The below is the only line that differs from assume called on SampleFromPrior.
+        # Could we just call assume on SampleFromPrior and then `setorder!!` after that?
         vi = DynamicPPL.setorder!!(vi, vn, DynamicPPL.get_num_produce(vi))
     else
         r = vi[vn]
     end
-    # TODO: call accumulate_assume?!
+
+    # TODO(mhauru) This get/set business is awful.
+    old_logp = DynamicPPL.getlogprior(vi)
+    vi = DynamicPPL.accumulate_assume!!(vi, r, 0, vn, dist)
+    vi = DynamicPPL.setlogprior!!(vi, old_logp)
+
+    # TODO(mhauru) Rather than this if-block, we should use try-catch within
+    # `set_trace_local_varinfo_maybe`. However, currently Libtask can't handle such a block,
+    # hence this.
+    if !using_local_vi
+        set_trace_local_varinfo_maybe(vi)
+    end
     return r, vi
 end
 
-# TODO(mhauru) Fix this.
-# function DynamicPPL.observe(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, value, vi)
-#     # NOTE: The `Libtask.produce` is now hit in `acclogp_observe!!`.
-#     return logpdf(dist, value), trace_local_varinfo_maybe(vi)
-# end
-
-function DynamicPPL.acclogp!!(
-    context::DynamicPPL.SamplingContext{<:Sampler{<:Union{PG,SMC}}},
-    varinfo::AbstractVarInfo,
-    logp,
+function DynamicPPL.tilde_observe!!(
+    ctx::DynamicPPL.SamplingContext{<:Sampler{<:Union{PG,SMC}}}, right, left, vn, vi
 )
-    varinfo_trace = trace_local_varinfo_maybe(varinfo)
-    return DynamicPPL.acclogp!!(DynamicPPL.childcontext(context), varinfo_trace, logp)
-end
+    arg_vi_id = objectid(vi)
+    vi = get_trace_local_varinfo_maybe(vi)
+    using_local_vi = objectid(vi) == arg_vi_id
+
+    # TODO(mhauru) This get/set business is awful.
+    old_logp = DynamicPPL.getloglikelihood(vi)
+    left, vi = DynamicPPL.tilde_observe!!(ctx.context, right, left, vn, vi)
+    new_loglikelihood = DynamicPPL.getloglikelihood(vi) - old_logp
+    vi = DynamicPPL.setloglikelihood!!(vi, old_logp)
+
+    # TODO(mhauru) Rather than this if-block, we should use try-catch within
+    # `set_trace_local_varinfo_maybe`. However, currently Libtask can't handle such a block,
+    # hence this.
+    if !using_local_vi
+        set_trace_local_varinfo_maybe(vi)
+    end
 
-# TODO(mhauru) Fix this.
-# function DynamicPPL.acclogp_observe!!(
-#     context::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, varinfo::AbstractVarInfo, logp
-# )
-#     Libtask.produce(logp)
-#     return trace_local_varinfo_maybe(varinfo)
-# end
+    Libtask.produce(new_loglikelihood)
+    return left, vi
+end
 
 # Convenient constructor
 function AdvancedPS.Trace(
@@ -483,7 +529,6 @@ end
 # called on `:invoke` expressions rather than `:call`s, but since those are implementation
 # details of the compiler, we set a bunch of methods as might_produce = true. We start with
 # `acclogp_observe!!` which is what calls `produce` and go up the call stack.
-# Libtask.might_produce(::Type{<:Tuple{typeof(DynamicPPL.acclogp_observe!!),Vararg}}) = true
 Libtask.might_produce(::Type{<:Tuple{typeof(DynamicPPL.tilde_observe!!),Vararg}}) = true
 Libtask.might_produce(::Type{<:Tuple{typeof(DynamicPPL.evaluate!!),Vararg}}) = true
 function Libtask.might_produce(