add multithreading to importance sampling

Author: marcoct

src/inference/importance.jl

@@ -17,40 +17,75 @@ The second variant uses a custom proposal distribution defined by the given gene
 All addresses of random choices sampled by the proposal should also be sampled by the model function.
 Setting `verbose=true` prints a progress message every sample.
 """
-function importance_sampling(model::GenerativeFunction{T,U}, model_args::Tuple,
-                             observations::ChoiceMap,
-                             num_samples::Int, verbose=false) where {T,U}
+function importance_sampling(
+        model::GenerativeFunction{T,U}, model_args::Tuple,
+        observations::ChoiceMap, num_samples::Int;
+        verbose=false, multithreaded=false) where {T,U}
     traces = Vector{U}(undef, num_samples)
     log_weights = Vector{Float64}(undef, num_samples)
-    for i=1:num_samples
-        verbose && println("sample: $i of $num_samples")
-        (traces[i], log_weights[i]) = generate(model, model_args, observations)
+    if multithreaded
+        Threads.@threads for i in 1:num_samples
+            importance_sampling_iter!(traces, log_weights, model, model_args, observations, i, verbose)
+        end
+    else
+        for i=1:num_samples
+            importance_sampling_iter!(traces, log_weights, model, model_args, observations, i, verbose)
+        end
     end
     log_total_weight = logsumexp(log_weights)
     log_ml_estimate = log_total_weight - log(num_samples)
     log_normalized_weights = log_weights .- log_total_weight
     return (traces, log_normalized_weights, log_ml_estimate)
 end
 
-function importance_sampling(model::GenerativeFunction{T,U}, model_args::Tuple,
-                             observations::ChoiceMap,
-                             proposal::GenerativeFunction, proposal_args::Tuple,
-                             num_samples::Int, verbose=false) where {T,U}
+function importance_sampling_iter!(
+        traces::Vector, log_weights::Vector{Float64},
+        model::GenerativeFunction, model_args::Tuple,
+        observations::ChoiceMap, i::Int, verbose::Bool)
+    (traces[i], log_weights[i]) = generate(model, model_args, observations)
+    verbose && Core.println("sample: $i of $num_samples completed in thread $(Threads.threadid())")
+    return nothing
+end
+
+function importance_sampling(
+        model::GenerativeFunction{T,U}, model_args::Tuple,
+        observations::ChoiceMap, proposal::GenerativeFunction, proposal_args::Tuple,
+        num_samples::Int; verbose=false, multithreaded=false) where {T,U}
     traces = Vector{U}(undef, num_samples)
     log_weights = Vector{Float64}(undef, num_samples)
-    for i=1:num_samples
-        verbose && println("sample: $i of $num_samples")
-        (proposed_choices, proposal_weight, _) = propose(proposal, proposal_args)
-        constraints = merge(observations, proposed_choices)
-        (traces[i], model_weight) = generate(model, model_args, constraints)
-        log_weights[i] = model_weight - proposal_weight
+    if multithreaded
+        Threads.@threads for i=1:num_samples
+            importance_sampling_iter!(
+                traces, log_weights, model, model_args,
+                observations, proposal, proposal_args, i, verbose)
+        end
+    else
+        for i=1:num_samples
+            importance_sampling_iter!(
+                traces, log_weights, model, model_args,
+                observations, proposal, proposal_args, i, verbose)
+        end
     end
     log_total_weight = logsumexp(log_weights)
     log_ml_estimate = log_total_weight - log(num_samples)
     log_normalized_weights = log_weights .- log_total_weight
     return (traces, log_normalized_weights, log_ml_estimate)
 end
 
+function importance_sampling_iter!(
+        traces::Vector, log_weights::Vector{Float64},
+        model::GenerativeFunction, model_args::Tuple,
+        observations::ChoiceMap,
+        proposal::GenerativeFunction, proposal_args::Tuple,
+        i::Int, verbose::Bool)
+    (proposed_choices, proposal_weight, _) = propose(proposal, proposal_args)
+    constraints = merge(observations, proposed_choices)
+    (traces[i], model_weight) = generate(model, model_args, constraints)
+    log_weights[i] = model_weight - proposal_weight
+    verbose && Core.println("sample $i of $num_samples completed in thread $(Threads.threadid())")
+    return nothing
+end
+
 """
     (trace, lml_est) = importance_resampling(model::GenerativeFunction,
         model_args::Tuple, observations::ChoiceMap, num_samples::Int,