Addressed review comments. Added docstrings where necessary. Corrected update_discard. Added test to test the discard functionality in a hierarchical model example.

Author: femtomc

src/modeling_library/switch/regenerate.jl

@@ -9,37 +9,15 @@ mutable struct SwitchRegenerateState{T}
     SwitchRegenerateState{T}(weight::Float64, score::Float64, noise::Float64, prev_trace::Trace) where T = new{T}(weight, score, noise, prev_trace)
 end
 
-@inline regenerate_recurse_merge(prev_choices::ChoiceMap, selection::EmptySelection) = prev_choices
-@inline regenerate_recurse_merge(prev_choices::ChoiceMap, selection::AllSelection) = choicemap()
-function regenerate_recurse_merge(prev_choices::ChoiceMap, selection::Selection)
-    prev_choice_value_iterator = get_values_shallow(prev_choices)
-    prev_choice_submap_iterator = get_submaps_shallow(prev_choices)
-    new_choices = choicemap()
-    for (key, value) in prev_choice_value_iterator
-        in(key, selection) && continue
-        set_value!(new_choices, key, value)
-    end
-    for (key, node1) in prev_choice_submap_iterator
-        if in(key, selection)
-            subsel = getindex(selection, key)
-            node = regenerate_recurse_merge(node1, subsel)
-            set_submap!(new_choices, key, node)
-        else
-            set_submap!(new_choices, key, node1)
-        end
-    end
-    return new_choices
-end
-
 function process!(gen_fn::Switch{C, N, K, T},
                   index::Int,
-                  index_argdiff::UnknownChange,
+                  index_argdiff::Diff,
                   args::Tuple,
                   kernel_argdiffs::Tuple,
                   selection::Selection, 
                   state::SwitchRegenerateState{T}) where {C, N, K, T}
     branch_fn = getfield(gen_fn.branches, index)
-    merged = regenerate_recurse_merge(get_choices(state.prev_trace), selection)
+    merged = get_selected(get_choices(state.prev_trace), complement(selection))
     new_trace, weight = generate(branch_fn, args, merged)
     retdiff = UnknownChange()
     weight -= project(state.prev_trace, complement(selection))

src/modeling_library/switch/update.jl

@@ -16,53 +16,79 @@ function update_recurse_merge(prev_choices::ChoiceMap, choices::ChoiceMap)
     choice_submap_iterator = get_submaps_shallow(choices)
     choice_value_iterator = get_values_shallow(choices)
     new_choices = DynamicChoiceMap()
-    for (key, value) in prev_choice_value_iterator
-        key in keys(choice_value_iterator) && continue
-        set_value!(new_choices, key, value)
+
+    # Add (address, value) to new_choices from prev_choices if address does not occur in choices.
+    for (address, value) in prev_choice_value_iterator
+        address in keys(choice_value_iterator) && continue
+        set_value!(new_choices, address, value)
     end
-    for (key, node1) in prev_choice_submap_iterator
-        if key in keys(choice_submap_iterator)
-            node2 = get_submap(choices, key)
+    
+    # Add (address, submap) to new_choices from prev_choices if address does not occur in choices. 
+    # If it does, enter a recursive call to update_recurse_merge.
+    for (address, node1) in prev_choice_submap_iterator
+        if address in keys(choice_submap_iterator)
+            node2 = get_submap(choices, address)
             node = update_recurse_merge(node1, node2)
-            set_submap!(new_choices, key, node)
+            set_submap!(new_choices, address, node)
         else
-            set_submap!(new_choices, key, node1)
+            set_submap!(new_choices, address, node1)
         end
     end
-    for (key, value) in choice_value_iterator
-        set_value!(new_choices, key, value)
+
+    # Add (address, value) from choices to new_choices. This is okay because we've excluded any conflicting addresses from the prev_choices above.
+    for (address, value) in choice_value_iterator
+        set_value!(new_choices, address, value)
     end
+
     sel, _ = zip(prev_choice_submap_iterator...)
     comp = complement(select(sel...))
-    for (key, node) in get_submaps_shallow(get_selected(choices, comp))
-        set_submap!(new_choices, key, node)
+    for (address, node) in get_submaps_shallow(get_selected(choices, comp))
+        set_submap!(new_choices, address, node)
     end
     return new_choices
 end
 
-function update_discard(prev_trace::Trace, choices::ChoiceMap, new_trace::Trace)
+@doc(
+"""
+update_recurse_merge(prev_choices::ChoiceMap, choices::ChoiceMap)
+
+Returns choices that are in constraints, merged with all choices in the previous trace that do not have the same address as some choice in the constraints."
+""", update_recurse_merge)
+
+function update_discard(prev_choices::ChoiceMap, choices::ChoiceMap, new_choices::ChoiceMap)
     discard = choicemap()
-    prev_choices = get_choices(prev_trace)
     for (k, v) in get_submaps_shallow(prev_choices)
-        isempty(get_submap(get_choices(new_trace), k)) && continue
-        isempty(get_submap(choices, k)) && continue
-        set_submap!(discard, k, v)
+        new_submap = get_submap(new_choices, k)
+        choices_submap = get_submap(choices, k)
+        sub_discard = update_discard(v, choices_submap, new_submap)
+        set_submap!(discard, k, sub_discard)
     end
     for (k, v) in get_values_shallow(prev_choices)
-        has_value(get_choices(new_trace), k) || continue
-        has_value(choices, k) || continue
-        set_value!(discard, k, v)
+        if (!has_value(new_choices, k) || has_value(choices, k))
+            set_value!(discard, k, v)
+        end
     end
     discard
 end
 
+@doc(
+"""
+update_discard(prev_choices::ChoiceMap, choices::ChoiceMap, new_choices::ChoiceMap)
+
+Returns choices from previous trace that:
+   1. have an address which does not appear in the new trace.
+   2. have an address which does appear in the constraints.
+""", update_discard)
+
+@inline update_discard(prev_trace::Trace, choices::ChoiceMap, new_trace::Trace) = update_discard(get_choices(prev_trace), choices, get_choices(new_trace))
+
 function process!(gen_fn::Switch{C, N, K, T},
-                  index::Int,
-                  index_argdiff::UnknownChange, # TODO: Diffed wrapper?
-                  args::Tuple,
-                  kernel_argdiffs::Tuple,
-                  choices::ChoiceMap, 
-                  state::SwitchUpdateState{T}) where {C, N, K, T, DV}
+        index::Int,
+        index_argdiff::UnknownChange,
+        args::Tuple,
+        kernel_argdiffs::Tuple,
+        choices::ChoiceMap, 
+        state::SwitchUpdateState{T}) where {C, N, K, T, DV}
 
     # Generate new trace.
     merged = update_recurse_merge(get_choices(state.prev_trace), choices)
@@ -81,12 +107,12 @@ function process!(gen_fn::Switch{C, N, K, T},
 end
 
 function process!(gen_fn::Switch{C, N, K, T},
-                  index::Int,
-                  index_argdiff::NoChange, # TODO: Diffed wrapper?
-                  args::Tuple,
-                  kernel_argdiffs::Tuple,
-                  choices::ChoiceMap, 
-                  state::SwitchUpdateState{T}) where {C, N, K, T}
+        index::Int,
+        index_argdiff::NoChange, # TODO: Diffed wrapper?
+        args::Tuple,
+        kernel_argdiffs::Tuple,
+        choices::ChoiceMap, 
+        state::SwitchUpdateState{T}) where {C, N, K, T}
 
     # Update trace.
     new_trace, weight, retdiff, discard = update(getfield(state.prev_trace, :branch), args, kernel_argdiffs, choices)
@@ -104,9 +130,9 @@ end
 @inline process!(gen_fn::Switch{C, N, K, T}, index::C, index_argdiff::Diff, args::Tuple, kernel_argdiffs::Tuple, choices::ChoiceMap, state::SwitchUpdateState{T}) where {C, N, K, T} = process!(gen_fn, getindex(gen_fn.cases, index), index_argdiff, args, kernel_argdiffs, choices, state)
 
 function update(trace::SwitchTrace{T},
-                args::Tuple, 
-                argdiffs::Tuple,
-                choices::ChoiceMap) where T
+        args::Tuple, 
+        argdiffs::Tuple,
+        choices::ChoiceMap) where T
     gen_fn = trace.gen_fn
     index, index_argdiff = args[1], argdiffs[1]
     state = SwitchUpdateState{T}(0.0, 0.0, 0.0, trace)

test/modeling_library/switch.jl

@@ -299,4 +299,42 @@
             @test isapprox(gradients[:x => :z], expected_choice_grad[1])
         end
     end
+
+    # ------------ (More complex) hierarchy to test discard ------------ #
+    
+    # Model chunk.
+    @gen (grad) function bang3((grad)(x::Float64), (grad)(y::Float64))
+        std::Float64 = 3.0
+        z = @trace(normal(x + y, std), :z)
+        q = @trace(bang2(z, y), :q)
+        return z
+    end
+    @gen (grad) function fuzz3((grad)(x::Float64), (grad)(y::Float64))
+        std::Float64 = 3.0
+        z = @trace(normal(x + 2 * y, std), :z)
+        m = @trace(normal(x + 3 * y, std), :m)
+        q = @trace(bang3(z, y), :q)
+        return z
+    end
+    sc3 = Switch(bang3, fuzz3)
+    @gen (grad) function bam3(s::Int)
+        x ~ sc3(s, 5.0, 3.0)
+        return x
+    end
+    # ----.
+    
+    @testset "update" begin
+        tr = simulate(bam3, (2, ))
+        old_sc = get_score(tr)
+        chm = choicemap((:x => :z, 5.0))
+        future_discarded = tr[:x => :z]
+        new_tr, w, rd, discard = update(tr, (2, ), (UnknownChange(), ), chm)
+        @test discard[:x => :z] == future_discarded
+        @test isapprox(old_sc, get_score(new_tr) - w)
+        chm = choicemap((:x => :z, 10.0))
+        future_discarded = tr[:x => :q => :q => :z]
+        new_tr, w, rd, discard = update(tr, (1, ), (UnknownChange(), ), chm)
+        @test discard[:x => :q => :q => :z] == future_discarded
+        @test isapprox(old_sc, get_score(new_tr) - w)
+    end
 end