Update pMCMC implementation as per discussion

Author: penelopeysm

src/mcmc/particle_mcmc.jl

@@ -11,6 +11,12 @@ Set the "del" flag for all variables in the VarInfo `vi`, thus marking them for
 resampling.
 """
 function set_all_del!(vi::AbstractVarInfo)
+    # TODO(penelopeysm): Instead of being a 'del' flag on the VarInfo, we
+    # could either:
+    # - keep a boolean 'resample' flag on the trace, or
+    # - modify the model context appropriately.
+    # However, this refactoring will have to wait until InitContext is
+    # merged into DPPL.
     for vn in keys(vi)
         DynamicPPL.set_flag!(vi, vn, "del")
     end
@@ -59,37 +65,30 @@ end
 function AdvancedPS.advance!(
     trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false
 )
-    # 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)
     return score
 end
 
 function AdvancedPS.delete_retained!(trace::TracedModel)
-    # TODO(DPPL0.37/penelopeysm): Explain this a bit better.
-    #
     # This method is called if, during a CSMC update, we perform a resampling
     # and choose the reference particle as the trajectory to carry on from.
     # In such a case, we need to ensure that when we continue sampling (i.e.
     # the next time we hit tilde_assume), we don't use the values in the 
     # reference particle but rather sample new values.
-    # In this implementation, we indiscriminately set the 'del' flag for all
-    # variables in the VarInfo. This is slightly overkill: it is not necessary
-    # to set the 'del' flag for variables that were already sampled. However,
-    # it allows us to avoid using DynamicPPL.set_retained_vns_del!.
+    #
+    # Here, we indiscriminately set the 'del' flag for all variables in the
+    # VarInfo. This is slightly overkill: it is not necessary to set the 'del'
+    # flag for variables that were already sampled. However, it allows us to
+    # avoid keeping track of which variables were sampled, which leads to many
+    # simplifications in the VarInfo data structure.
     set_all_del!(trace.varinfo)
     return trace
 end
 
 function AdvancedPS.reset_model(trace::TracedModel)
-    return Accessors.@set trace.varinfo = DynamicPPL.reset_num_produce!!(trace.varinfo)
+    return trace
 end
 
 function Libtask.TapedTask(taped_globals, model::TracedModel; kwargs...)
@@ -213,7 +212,6 @@ function DynamicPPL.initialstep(
 )
     # Reset the VarInfo.
     vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator())
-    vi = DynamicPPL.reset_num_produce!!(vi)
     set_all_del!(vi)
     vi = DynamicPPL.resetlogp!!(vi)
     vi = DynamicPPL.empty!!(vi)
@@ -344,7 +342,6 @@ function DynamicPPL.initialstep(
 )
     vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator())
     # Reset the VarInfo before new sweep
-    vi = DynamicPPL.reset_num_produce!!(vi)
     set_all_del!(vi)
     vi = DynamicPPL.resetlogp!!(vi)
 
@@ -366,11 +363,6 @@ function DynamicPPL.initialstep(
 
     # Compute the first transition.
     _vi = reference.model.f.varinfo
-    # Unset any 'del' flags before we actually construct the transition.
-    # This is necessary because the model will be re-evaluated and we
-    # want to make sure we do use the values in the reference particle
-    # instead of resampling them.
-    unset_all_del!(_vi)
     transition = PGTransition(model, _vi, logevidence)
 
     return transition, PGState(_vi, reference.rng)
@@ -382,10 +374,10 @@ function AbstractMCMC.step(
     # Reset the VarInfo before new sweep.
     vi = state.vi
     vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator())
-    vi = DynamicPPL.reset_num_produce!!(vi)
     vi = DynamicPPL.resetlogp!!(vi)
 
     # Create reference particle for which the samples will be retained.
+    unset_all_del!(vi)
     reference = AdvancedPS.forkr(AdvancedPS.Trace(model, spl, vi, state.rng))
 
     # For all other particles, do not retain the variables but resample them.
@@ -412,11 +404,6 @@ function AbstractMCMC.step(
 
     # Compute the transition.
     _vi = newreference.model.f.varinfo
-    # Unset any 'del' flags before we actually construct the transition.
-    # This is necessary because the model will be re-evaluated and we
-    # want to make sure we do use the values in the reference particle
-    # instead of resampling them.
-    unset_all_del!(_vi)
     transition = PGTransition(model, _vi, logevidence)
 
     return transition, PGState(_vi, newreference.rng)
@@ -499,12 +486,11 @@ function DynamicPPL.assume(
         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))
+        # Could we just call assume on SampleFromPrior with a specific rng?
+        r = rand(trng, dist)
+        vi[vn] = DynamicPPL.tovec(r)
     else
         r = vi[vn]
     end
@@ -546,8 +532,6 @@ function AdvancedPS.Trace(
     rng::AdvancedPS.TracedRNG,
 )
     newvarinfo = deepcopy(varinfo)
-    newvarinfo = DynamicPPL.reset_num_produce!!(newvarinfo)
-
     tmodel = TracedModel(model, sampler, newvarinfo, rng)
     newtrace = AdvancedPS.Trace(tmodel, rng)
     return newtrace