Use nice functions

Author: penelopeysm

src/mcmc/particle_mcmc.jl

@@ -4,6 +4,32 @@
 
 ### AdvancedPS models and interface
 
+"""
+    set_all_del!(vi::AbstractVarInfo)
+
+Set the "del" flag for all variables in the VarInfo `vi`, thus marking them for
+resampling.
+"""
+function set_all_del!(vi::AbstractVarInfo)
+    for vn in keys(vi)
+        DynamicPPL.set_flag!(vi, vn, "del")
+    end
+    return nothing
+end
+
+"""
+    unset_all_del!(vi::AbstractVarInfo)
+
+Unset the "del" flag for all variables in the VarInfo `vi`, thus preventing
+them from being resampled.
+"""
+function unset_all_del!(vi::AbstractVarInfo)
+    for vn in keys(vi)
+        DynamicPPL.unset_flag!(vi, vn, "del")
+    end
+    return nothing
+end
+
 struct TracedModel{S<:AbstractSampler,V<:AbstractVarInfo,M<:Model,E<:Tuple} <:
        AdvancedPS.AbstractGenericModel
     model::M
@@ -58,9 +84,7 @@ function AdvancedPS.delete_retained!(trace::TracedModel)
     # 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!.
-    for vn in keys(trace.varinfo)
-        DynamicPPL.set_flag!(trace.varinfo, vn, "del")
-    end
+    set_all_del!(trace.varinfo)
     return trace
 end
 
@@ -190,9 +214,7 @@ function DynamicPPL.initialstep(
     # Reset the VarInfo.
     vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator())
     vi = DynamicPPL.reset_num_produce!!(vi)
-    for vn in keys(vi)
-        DynamicPPL.set_flag!(vi, vn, "del")
-    end
+    set_all_del!(vi)
     vi = DynamicPPL.resetlogp!!(vi)
     vi = DynamicPPL.empty!!(vi)
 
@@ -323,9 +345,7 @@ function DynamicPPL.initialstep(
     vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator())
     # Reset the VarInfo before new sweep
     vi = DynamicPPL.reset_num_produce!!(vi)
-    for vn in keys(vi)
-        DynamicPPL.set_flag!(vi, vn, "del")
-    end
+    set_all_del!(vi)
     vi = DynamicPPL.resetlogp!!(vi)
 
     # Create a new set of particles
@@ -350,9 +370,7 @@ function DynamicPPL.initialstep(
     # 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.
-    for vn in keys(_vi)
-        DynamicPPL.unset_flag!(_vi, vn, "del")
-    end
+    unset_all_del!(_vi)
     transition = PGTransition(model, _vi, logevidence)
 
     return transition, PGState(_vi, reference.rng)
@@ -371,9 +389,7 @@ function AbstractMCMC.step(
     reference = AdvancedPS.forkr(AdvancedPS.Trace(model, spl, vi, state.rng))
 
     # For all other particles, do not retain the variables but resample them.
-    for vn in keys(vi)
-        DynamicPPL.set_flag!(vi, vn, "del")
-    end
+    set_all_del!(vi)
 
     # Create a new set of particles.
     num_particles = spl.alg.nparticles
@@ -400,9 +416,7 @@ function AbstractMCMC.step(
     # 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.
-    for vn in keys(_vi)
-        DynamicPPL.unset_flag!(_vi, vn, "del")
-    end
+    unset_all_del!(_vi)
     transition = PGTransition(model, _vi, logevidence)
 
     return transition, PGState(_vi, newreference.rng)