Support passing of keyword args to rejuv. kernels.

Author: ztangent

src/rejuvenate.jl

@@ -2,44 +2,48 @@
 export pf_rejuvenate!, pf_move_accept!, pf_move_reweight!
 export move_reweight
 
+using Gen: check_observations
+
 """
     pf_rejuvenate!(state::ParticleFilterState, kern, kern_args::Tuple=(),
-                   n_iters::Int=1; method=:move)
+                   n_iters::Int=1; method=:move, kwargs...)
 
 Rejuvenates particles by repeated application of a kernel `kern`. `kern`
 should be a callable which takes a trace as its first argument, and returns
 a tuple with a trace as the first return value. `method` specifies the
 rejuvenation method: `:move` for MCMC moves without a reweighting step,
-and `:reweight` for rejuvenation with a reweighting step.
+and `:reweight` for rejuvenation with a reweighting step. Additional keyword
+arguments are passed to the kernel.
 """
 function pf_rejuvenate!(state::ParticleFilterView, kern, kern_args::Tuple=(),
-                        n_iters::Int=1; method::Symbol=:move)
+                        n_iters::Int=1; method::Symbol=:move, kwargs...)
     if method == :move
-        return pf_move_accept!(state, kern, kern_args, n_iters)
+        return pf_move_accept!(state, kern, kern_args, n_iters; kwargs...)
     elseif method == :reweight
-        return pf_move_reweight!(state, kern, kern_args, n_iters)
+        return pf_move_reweight!(state, kern, kern_args, n_iters; kwargs...)
     else
         error("Method $method not recognized.")
     end
 end
 
 """
     pf_move_accept!(state::ParticleFilterState, kern,
-                    kern_args::Tuple=(), n_iters::Int=1)
+                    kern_args::Tuple=(), n_iters::Int=1; kwargs...)
 
 Rejuvenates particles by repeated application of a MCMC kernel `kern`. `kern`
 should be a callable which takes a trace as its first argument, and returns
 a tuple `(trace, accept)`, where `trace` is the (potentially) new trace, and
 `accept` is true if the MCMC move was accepted. Subsequent arguments to `kern`
-can be supplied with `kern_args`. The kernel is repeatedly applied to each trace
-for `n_iters`.
+can be supplied with `kern_args` or `kwargs`. The kernel is repeatedly applied
+to each trace for `n_iters`.
 """
 function pf_move_accept!(state::ParticleFilterView,
-                         kern, kern_args::Tuple=(), n_iters::Int=1)
+                         kern, kern_args::Tuple=(), n_iters::Int=1;
+                         kwargs...)
     # Potentially rejuvenate each trace
     for (i, trace) in enumerate(state.traces)
         for k = 1:n_iters
-            trace, accept = kern(trace, kern_args...)
+            trace, accept = kern(trace, kern_args...; kwargs...)
             @debug "Accepted: $accept"
         end
         state.new_traces[i] = trace
@@ -50,28 +54,31 @@ end
 
 """
     pf_move_reweight!(state::ParticleFilterState, kern,
-                      kern_args::Tuple=(), n_iters::Int=1)
+                      kern_args::Tuple=(), n_iters::Int=1; kwargs...)
 
 Rejuvenates and reweights particles by repeated application of a reweighting
 kernel `kern`, as described in [1]. `kern` should be a callable which takes a
 trace as its first argument, and returns a tuple `(trace, rel_weight)`,
 where `trace` is the new trace, and `rel_weight` is the relative log-importance
-weight. Subsequent arguments to `kern` can be supplied with `kern_args`.
-The kernel is repeatedly applied to each trace for `n_iters`, and the weights
-accumulated accordingly. Both the [`move_reweight`](@ref) function and
+weight. Subsequent arguments to `kern` can be supplied with `kern_args` or
+`kwargs`. The kernel is repeatedly applied to each trace for `n_iters`, and the
+weights accumulated accordingly.
+
+Both the [`move_reweight`](@ref) function and
 [symmetric trace translators](https://www.gen.dev/stable/ref/trace_translators/)
 can serve as reweighting kernels. 
 
 [1] R. A. G. Marques and G. Storvik, "Particle move-reweighting strategies for
 online inference," Preprint series. Statistical Research Report, 2013.
 """
 function pf_move_reweight!(state::ParticleFilterView,
-                           kern, kern_args::Tuple=(), n_iters::Int=1)
+                           kern, kern_args::Tuple=(), n_iters::Int=1;
+                           kwargs...)
     # Move and reweight each trace
     for (i, trace) in enumerate(state.traces)
         weight = 0
         for k = 1:n_iters
-            trace, rel_weight = kern(trace, kern_args...)
+            trace, rel_weight = kern(trace, kern_args...; kwargs...)
             weight += rel_weight
             @debug "Rel. Weight: $rel_weight"
         end
@@ -83,15 +90,15 @@ function pf_move_reweight!(state::ParticleFilterView,
 end
 
 """
-    move_reweight(trace, selection)
-    move_reweight(trace, proposal, proposal_args)
-    move_reweight(trace, proposal, proposal_args, involution)
-    move_reweight(trace, proposal_fwd, args_fwd,
-                  proposal_bwd, args_bwd, involution)
+    move_reweight(trace, selection; kwargs...)
+    move_reweight(trace, proposal, proposal_args; kwargs...)
+    move_reweight(trace, proposal, proposal_args, involution; kwargs...)
+    move_reweight(trace, proposal_fwd, args_fwd, proposal_bwd, args_bwd,
+                  involution ; kwargs...)
 
-Move-reweight MCMC kernel, which takes in a `trace` and returns a new trace
-along with a relative importance weight. This can be used for rejuvenation
-within a particle filter, as described in [1].
+Move-reweight rejuvenation kernel, which takes in a `trace` and returns a
+new trace along with a relative importance weight. This can be used for
+rejuvenation within a particle filter, as described in [1].
 
 Several variants of `move_reweight` exist, differing in the complexity
 involved in proposing and re-weighting random choices:
@@ -108,18 +115,25 @@ involved in proposing and re-weighting random choices:
    adjusts the computation of the relative importance weight by scoring
    the backward choices under the backward proposal.
 
+Similar to `metropolis_hastings`, a `check` flag and `observations` choicemap
+can be provided as keyword arguments to ensure that observed choices are
+preserved in the new trace.
+
 [1] R. A. G. Marques and G. Storvik, "Particle move-reweighting strategies for
 online inference," Preprint series. Statistical Research Report, 2013.
 """
-function move_reweight(trace::Trace, selection::Selection)
+function move_reweight(trace::Trace, selection::Selection;
+                       check=false, observations=EmptyChoiceMap())
     args = get_args(trace)
     argdiffs = map((_) -> NoChange(), args)
     new_trace, rel_weight = regenerate(trace, args, argdiffs, selection)
+    check && check_observations(get_choices(new_trace), observations)
     return new_trace, rel_weight
 end
 
 function move_reweight(trace::Trace, proposal::GenerativeFunction,
-                       proposal_args::Tuple)
+                       proposal_args::Tuple;
+                       check=false, observations=EmptyChoiceMap())
     model_args = Gen.get_args(trace)
     argdiffs = map((_) -> NoChange(), model_args)
     fwd_choices, fwd_score, fwd_ret =
@@ -128,6 +142,7 @@ function move_reweight(trace::Trace, proposal::GenerativeFunction,
         update(trace, model_args, argdiffs, fwd_choices)
     bwd_score, bwd_ret =
         assess(proposal, (new_trace, proposal_args...), discard)
+    check && check_observations(get_choices(new_trace), observations)
     rel_weight = weight - fwd_score + bwd_score
     return new_trace, rel_weight
 end
@@ -140,19 +155,22 @@ function move_reweight(trace::Trace, proposal::GenerativeFunction,
         involution(trace, fwd_choices, fwd_ret, proposal_args)
     bwd_score, bwd_ret =
         assess(proposal, (new_trace, proposal_args...), bwd_choices)
+    check && check_observations(get_choices(new_trace), observations)
     rel_weight = weight - fwd_score + bwd_score
     return new_trace, rel_weight
 end
 
 function move_reweight(trace::Trace, proposal_fwd::GenerativeFunction,
                        args_fwd::Tuple, proposal_bwd::GenerativeFunction,
-                       args_bwd::Tuple, involution)
+                       args_bwd::Tuple, involution;
+                       check=false, observations=EmptyChoiceMap())
     fwd_choices, fwd_score, fwd_ret =
         propose(proposal_fwd, (trace, args_fwd...,))
     new_trace, bwd_choices, weight =
         involution(trace, fwd_choices, fwd_ret, args_fwd)
     bwd_score, bwd_ret =
         assess(proposal_bwd, (new_trace, args_bwd...), bwd_choices)
+    check && check_observations(get_choices(new_trace), observations)
     rel_weight = weight - fwd_score + bwd_score
     return new_trace, rel_weight
 end

test/runtests.jl

@@ -246,19 +246,21 @@ trace, _ = generate(line_model, (1,), observations)
 slope, outlier = trace[:slope], trace[out_addr]
 
 # Test selection variant
-expected_w(out_old, out_new, slope) =
+expected_w = (out_old, out_new, slope) -> begin
     logpdf(normal, 0, slope, out_new ? 10. : 1.) -
     logpdf(normal, 0, slope, out_old ? 10. : 1.)
+end
 trs_ws = [move_reweight(trace, select(out_addr)) for i in 1:100]
 @test all(w ≈ expected_w(outlier, tr[out_addr], slope) for (tr, w) in trs_ws)
 
 # Test proposal variant
-expected_w(out_old, out_new, slope) =
+expected_w = (out_old, out_new, slope) -> begin 
     logpdf(bernoulli, out_new, 0.1) - logpdf(bernoulli, out_old, 0.1) +
     logpdf(normal, 0, slope, out_new ? 10. : 1.) -
     logpdf(normal, 0, slope, out_old ? 10. : 1.) +
     (out_old == out_new ? 0.0 :
         logpdf(bernoulli, out_old, 0.9) - logpdf(bernoulli, out_old, 0.1))
+end
 @gen outlier_propose(tr, idx) = {:line => idx => :outlier} ~ bernoulli(0.9)
 trs_ws = [move_reweight(trace, outlier_propose, (1,)) for i in 1:100]
 @test all(w ≈ expected_w(outlier, tr[out_addr], slope) for (tr, w) in trs_ws)
@@ -271,7 +273,7 @@ logger = SimpleLogger(buffer, Logging.Debug)
 state = pf_initialize(line_model, (10,), generate_line(10, 1.), 100)
 old_traces = get_traces(state)
 with_logger(logger) do
-    pf_move_accept!(state, metropolis_hastings, (select(:slope),), 1)
+    pf_move_accept!(state, mh, (select(:slope),), 1; check=false)
 end
 
 # Extract acceptances from debug log
@@ -293,7 +295,7 @@ logger = SimpleLogger(buffer, Logging.Debug)
 state = pf_initialize(line_model, (10,), generate_line(10, 1.), 100)
 old_weights = copy(get_log_weights(state))
 with_logger(logger) do
-    pf_move_reweight!(state, move_reweight, (select(:slope),), 1)
+    pf_move_reweight!(state, move_reweight, (select(:slope),), 1; check=false)
 end
 new_weights = copy(get_log_weights(state))
 
@@ -315,9 +317,10 @@ state = pf_initialize(line_model, (10,), generate_line(10, 1.), 100)
 old_traces = get_traces(state)[1:50]
 old_weights = get_log_weights(state)[51:end]
 
+kern_args = (select(:slope),)
 with_logger(logger) do
-    pf_move_accept!(state[1:50], metropolis_hastings, (select(:slope),), 1)
-    pf_move_reweight!(state[51:end], move_reweight, (select(:slope),), 1)
+    pf_rejuvenate!(state[1:50], mh, kern_args, 1; method=:move)
+    pf_rejuvenate!(state[51:end], move_reweight, kern_args, 1; method=:reweight)
 end
 
 # Extract acceptances and relative weights from debug log