Add choiceproduct utility function.

Author: ztangent

docs/src/initialize.md

@@ -23,3 +23,5 @@ To reduce variance, particle filters can also be initialized via stratified samp
 pf_initialize(model::GenerativeFunction, model_args::Tuple,
               observations::ChoiceMap, strata, n_particles::Int)
 ```
+
+For convenience, strata can be generated using the [`choiceproduct`](@ref) function.

docs/src/update.md

@@ -54,3 +54,5 @@ All of the above methods can also be combined with stratified sampling. This can
 pf_update!(state::ParticleFilterView, new_args::Tuple, argdiffs::Tuple,
            observations::ChoiceMap, strata)
 ```
+
+For convenience, strata can be generated using the [`choiceproduct`](@ref) function.

docs/src/utils.md

@@ -55,4 +55,12 @@ var(::ParticleFilterView, addr)
 var(f::Function, ::ParticleFilterView, addr, addrs...)
 proportionmap(::ParticleFilterView, addr)
 proportionmap(f::Function, ::ParticleFilterView, addr, addrs...)
+```
+
+## Stratification
+
+To support the use of stratified sampling, the [`choiceproduct`](@ref) method can be used to conveniently generate a list of choicemap strata:
+
+```@docs
+choiceproduct
 ```

src/utils.jl

@@ -1,4 +1,5 @@
 ## Various utility functions ##
+export choiceproduct
 export get_log_norm_weights, get_norm_weights
 export effective_sample_size, get_ess
 export log_ml_estimate, get_lml_est
@@ -55,6 +56,49 @@ function stratified_map!(f::Function, n_total::Int, strata, args...;
     return nothing
 end
 
+"""
+    choiceproduct((addr, vals))
+    choiceproduct((addr, vals), choices::Tuple...)
+    choiceproduct(choices::Dict)
+
+Returns an iterator over `ChoiceMap`s given a tuple or sequence of tuples of
+the form  `(addr, vals)`, where `addr` specifies a choice address, and
+`vals` specifies a list of values for that address.
+    
+If multiple tuples are provided, the iterator will be a Cartesian product over
+the `(addr, vals)` pairs, where each resulting `ChoiceMap` contains all
+specified addresses. Instead of specifying multiple tuples, a dictionary mapping
+addresses to values can also be provided.
+
+# Examples
+
+This function can be used to conveniently generate `ChoiceMap`s for stratified
+sampling. For example, we can use `choiceproduct` instead of manually
+constructing a list of strata:
+
+```julia
+# Manual construction
+strata = [choicemap((:a,  1), (:b,  3)), choicemap((:a,  2), (:b,  3))]
+# Using choiceproduct
+strata = choiceproduct((:a, [1, 2]), (:b, [3]))
+```
+"""
+function choiceproduct(choices::Tuple...)
+    prod_iter = Iterators.product((((addr, v) for v in vals) for
+                                  (addr, vals) in choices)...)
+    return (choicemap(cs...) for cs in prod_iter)
+end
+
+function choiceproduct(choices::Dict)
+    prod_iter = Iterators.product((((addr, v) for v in vals) for
+                                  (addr, vals) in choices)...)
+    return (choicemap(cs...) for cs in prod_iter)
+end
+
+function choiceproduct((addr, vals)::Tuple)
+    return (choicemap((addr, v)) for v in vals)
+end
+
 lognorm(vs::AbstractVector) = vs .- logsumexp(vs)
 
 function softmax(vs::AbstractVector{T}) where {T <: Real}

test/utils.jl

@@ -8,4 +8,16 @@
     @test sum(weights) ≈ 1.0
     ess = get_ess(state)
     @test ess ≈ sum(weights)^2 / sum(weights .^ 2)
+
+    strata = choiceproduct((:a, [1, 2])) |> collect
+    @test choicemap((:a , 1)) in strata
+    @test choicemap((:a , 2)) in strata    
+
+    strata = choiceproduct((:a, [1, 2]), (:b, [3])) |> collect
+    @test choicemap((:a , 1), (:b, 3)) in strata
+    @test choicemap((:a , 2), (:b, 3)) in strata    
+
+    strata = choiceproduct(Dict(:a => [1, 2], :b => [3])) |> collect
+    @test choicemap((:a , 1), (:b, 3)) in strata
+    @test choicemap((:a , 2), (:b, 3)) in strata    
 end