#512 - Fix missing addresses in discard during update

[yifr]

@mlb2251 and I ran into a bug @fsaad described in #512. We tested his fix on our own example and it works! Let us know if there's anything more we need to do to get this merged. Thanks!

[ztangent]

Thanks for looking into this and making the PR! The fix looks good, I think we should just add a test case (e.g. the MWE that @fsaad gave in #512) in the tests for update, and then we can merge.

The tests for the dynamic DSL implementation of update can be found here:

Gen.jl/test/dsl/dynamic_dsl.jl

Lines 86 to 184 in 95c797f

	@testset "update" begin
	@gen function bar()
	@trace(normal(0, 1), :a)
	end
	@gen function baz()
	@trace(normal(0, 1), :b)
	end
	@gen function foo()
	if @trace(bernoulli(0.4), :branch)
	@trace(normal(0, 1), :x)
	@trace(bar(), :u)
	else
	@trace(normal(0, 1), :y)
	@trace(baz(), :v)
	end
	end
	# get a trace which follows the first branch
	constraints = choicemap()
	constraints[:branch] = true
	(trace,) = generate(foo, (), constraints)
	x = get_choices(trace)[:x]
	a = get_choices(trace)[:u => :a]
	# force to follow the second branch
	y = 1.123
	b = -2.1
	constraints = choicemap()
	constraints[:branch] = false
	constraints[:y] = y
	constraints[:v => :b] = b
	(new_trace, weight, retdiff, discard) = update(trace,
	(), (), constraints)
	# test discard
	@test get_value(discard, :branch) == true
	@test get_value(discard, :x) == x
	@test get_value(discard, :u => :a) == a
	@test length(collect(get_values_shallow(discard))) == 2
	@test length(collect(get_submaps_shallow(discard))) == 1
	# test new trace
	new_assignment = get_choices(new_trace)
	@test get_value(new_assignment, :branch) == false
	@test get_value(new_assignment, :y) == y
	@test get_value(new_assignment, :v => :b) == b
	@test length(collect(get_values_shallow(new_assignment))) == 2
	@test length(collect(get_submaps_shallow(new_assignment))) == 1
	# test score and weight
	prev_score = (
	logpdf(bernoulli, true, 0.4) +
	logpdf(normal, x, 0, 1) +
	logpdf(normal, a, 0, 1))
	expected_new_score = (
	logpdf(bernoulli, false, 0.4) +
	logpdf(normal, y, 0, 1) +
	logpdf(normal, b, 0, 1))
	expected_weight = expected_new_score - prev_score
	@test isapprox(expected_new_score, get_score(new_trace))
	@test isapprox(expected_weight, weight)
	# test retdiff
	@test retdiff === UnknownChange()
	# Addresses under the :data namespace will be visited,
	# but nothing there will be discarded.
	@gen function loopy()
	a = @trace(normal(0, 1), :a)
	for i=1:5
	@trace(normal(a, 1), :data => i)
	end
	end
	# Get an initial trace
	constraints = choicemap()
	constraints[:a] = 0
	for i=1:5
	constraints[:data => i] = 0
	end
	(trace,) = generate(loopy, (), constraints)
	# Update a
	constraints = choicemap()
	constraints[:a] = 1
	(new_trace, weight, retdiff, discard) = update(trace,
	(), (), constraints)
	# Test discard, score, weight, retdiff
	@test get_value(discard, :a) == 0
	prev_score = logpdf(normal, 0, 0, 1) * 6
	expected_new_score = logpdf(normal, 1, 0, 1) + 5 * logpdf(normal, 0, 1, 1)
	expected_weight = expected_new_score - prev_score
	@test isapprox(expected_new_score, get_score(new_trace))
	@test isapprox(expected_weight, weight)
	@test retdiff === UnknownChange()
	end

[yifr]

Sorry for the delay! I added in @fsaad's suggested test in the update test suite. Let me know if it looks good, or needs any changes!

[ztangent]

Looks good to me!