Add weight correction for stratified sampling.

Author: ztangent

src/initialize.jl

@@ -97,10 +97,12 @@ function pf_initialize(
     V = dynamic ? Trace : U # Determine trace type for particle filter
     traces = Vector{V}(undef, n_particles)
     log_weights = Vector{Float64}(undef, n_particles)
+    n_strata = length(strata)
     # Generate traces in a stratified manner
     stratified_map!(n_particles, strata; layout=layout) do i, stratum
         constraints = merge(stratum, observations)
         (traces[i], log_weights[i]) = generate(model, model_args, constraints)
+        log_weights[i] += log(n_strata)
     end
     return ParticleFilterState{V}(traces, Vector{V}(undef, n_particles),
                                   log_weights, 0., collect(1:n_particles))
@@ -115,11 +117,12 @@ function pf_initialize(
     V = dynamic ? Trace : U # Determine trace type for particle filter
     traces = Vector{V}(undef, n_particles)
     log_weights = Vector{Float64}(undef, n_particles)
+    n_strata = length(strata)
     stratified_map!(n_particles, strata; layout=layout) do i, stratum
         (prop_choices, prop_weight, _) = propose(proposal, proposal_args)
         constraints = merge(stratum, observations, prop_choices)
         (traces[i], model_weight) = generate(model, model_args, constraints)
-        log_weights[i] = model_weight - prop_weight
+        log_weights[i] = model_weight - prop_weight + log(n_strata)
     end
     return ParticleFilterState{V}(traces, Vector{V}(undef, n_particles),
                                   log_weights, 0., collect(1:n_particles))

src/update.jl

@@ -195,14 +195,15 @@ function pf_update!(state::ParticleFilterView, new_args::Tuple,
                     layout=:interleaved)
     # Update traces in a stratified manner
     n_particles = length(state.traces)
+    n_strata = length(strata)
     stratified_map!(n_particles, strata; layout=layout) do i, stratum
         constraints = merge(stratum, observations)
         state.new_traces[i], increment, _, discard =
             update(state.traces[i], new_args, argdiffs, constraints)
         if !isempty(discard)
             error("Choices were updated or deleted: $discard")
         end
-        state.log_weights[i] += increment    
+        state.log_weights[i] += increment + log(n_strata)
     end
     update_refs!(state)
     return state
@@ -215,11 +216,12 @@ function pf_update!(state::ParticleFilterView, translator::TraceTranslator,
     observations = translator.new_observations
     # Update traces in a stratified manner
     n_particles = length(state.traces)
+    n_strata = length(strata)
     stratified_map!(n_particles, strata; layout=layout) do i, stratum
         translator.new_observations = merge(stratum, observations)
         state.new_traces[i], log_weight =
             translator(state.traces[i]; translator_args...)
-        state.log_weights[i] += log_weight
+        state.log_weights[i] += log_weight + log(n_strata)
     end
     update_refs!(state)
     return state

test/initialize.jl

@@ -3,6 +3,7 @@
 @testset "Initialize with default proposal" begin
     state = pf_initialize(line_model, (0,), choicemap(), 100)
     @test all(-2 <= tr[:slope] <= 2 for tr in get_traces(state))
+    @test all(w ≈ 0 for w in get_log_weights(state))
     state = pf_initialize(line_model, (1,), line_choicemap(1), 100)
     @test all(tr[:line => 1 => :y] == 0 for tr in get_traces(state))
     state = pf_initialize(line_model, (10,), line_choicemap(10), 100)
@@ -17,6 +18,7 @@ end
 @testset "Initialize with custom proposal" begin
     state = pf_initialize(line_model, (0,), choicemap(), line_propose, (0,), 100)
     @test all(tr[:slope] == 0 for tr in get_traces(state))
+    @test all(w ≈ log(1/5) for w in get_log_weights(state))
     state = pf_initialize(line_model, (1,), line_choicemap(1),
                           outlier_propose, ([1],), 100)
     @test all(tr[:line => 1 => :outlier] == false for tr in get_traces(state))
@@ -38,6 +40,9 @@ end
     slope_strata = (slope_choicemap(s) for s in -2:1:2)
     observations = line_choicemap(1)
     # Test contiguous stratification
+    state = pf_initialize(line_model, (0,), choicemap(),
+                          slope_strata, 100; layout=:contiguous)
+    @test all(w ≈ 0 for w in get_log_weights(state))
     state = pf_initialize(line_model, (1,), observations,
                           slope_strata, 100; layout=:contiguous)
     for (k, slope) in zip([20, 40, 60, 80, 100], -2:1:2)
@@ -46,6 +51,9 @@ end
         @test all(tr[:line => 1 => :y] == 0 for tr in traces)
     end
     # Test interleaved stratification
+    state = pf_initialize(line_model, (0,), choicemap(),
+                          slope_strata, 100; layout=:interleaved)
+    @test all(w ≈ 0 for w in get_log_weights(state))
     state = pf_initialize(line_model, (1,), observations,
                           slope_strata, 100; layout=:interleaved)
     for (k, slope) in zip([1, 2, 3, 4, 5], -2:1:2)
@@ -66,6 +74,9 @@ end
         @test all(tr[:slope] == slope for tr in traces)
         @test all(tr[:line => 1 => :outlier] == false for tr in traces)
         @test all(tr[:line => 1 => :y] == 0 for tr in traces)
+        expected_w = (logpdf(bernoulli, false, 0.1) +
+                      logpdf(normal, 0.0, slope, 1.0))
+        @test all(w ≈ expected_w for w in get_log_weights(state[(k-20+1):k]))
     end
     # Test interleaved stratification
     state = pf_initialize(line_model, (1,), observations, slope_strata,
@@ -75,6 +86,9 @@ end
         @test all(tr[:slope] == slope for tr in traces)
         @test all(tr[:line => 1 => :outlier] == false for tr in traces)
         @test all(tr[:line => 1 => :y] == 0 for tr in traces)
+        expected_w = (logpdf(bernoulli, false, 0.1) +
+                      logpdf(normal, 0.0, slope, 1.0))
+        @test all(w ≈ expected_w for w in get_log_weights(state[k:5:100]))
     end
 end
 

test/update.jl

@@ -2,9 +2,12 @@
 
 @testset "Update with default proposal" begin
     state = pf_initialize(line_model, (0,), choicemap(), 100)
-    state = pf_update!(state, (10,), (UnknownChange(),), line_choicemap(10))
-    @test all(tr[:line => 10 => :y] == 0 for tr in get_traces(state))
-    @test all(w != 0 for w in get_log_weights(state))
+    state = pf_update!(state, (1,), (UnknownChange(),), line_choicemap(1))
+    @test all(tr[:line => 1 => :y] == 0 for tr in get_traces(state))
+    outliers = [tr[:line => 1 => :outlier] for tr in get_traces(state)]
+    expected_ws = [logpdf(normal, 0.0, tr[:slope], o ? 10.0 : 1.0)
+                   for (o, tr) in zip(outliers, get_traces(state))]
+    @test all(get_log_weights(state) .≈ expected_ws)
 end
 
 @testset "Update with stratification" begin
@@ -17,17 +20,23 @@ end
     for (k, val) in zip([50, 100], [false, true])
         traces = get_traces(state[(k-50+1):k])
         @test all(tr[:line => 1 => :outlier] == val for tr in traces)
+        std = val ? 10.0 : 1.0
+        expected_ws = [(logpdf(bernoulli, val, 0.1) + log(2) +
+                        logpdf(normal, 0.0, tr[:slope], std)) for tr in traces]
+        @test all(get_log_weights(state[(k-50+1):k]) .≈ expected_ws)
     end
-    @test all(w != 0 for w in get_log_weights(state))
     # Test interleaved stratification
     state = pf_initialize(line_model, (0,), choicemap(), 100)
     state = pf_update!(state, (1,), (UnknownChange(),), observations,
                        outlier_strata; layout=:interleaved)
     for (k, val) in zip(1:2, [false, true])
         traces = get_traces(state[k:2:100])
         @test all(tr[:line => 1 => :outlier] == val for tr in traces)
+        std = val ? 10.0 : 1.0
+        expected_ws = [(logpdf(bernoulli, val, 0.1) + log(2) +
+                        logpdf(normal, 0.0, tr[:slope], std)) for tr in traces]
+        @test all(get_log_weights(state[k:2:100]) .≈ expected_ws)
     end
-    @test all(w != 0 for w in get_log_weights(state))
 end
 
 @gen outlier_propose(tr, idxs) =