predict implemented for multiple chains (#1421)

* made predict work for multiple chains

* export predict from Turing

* added appropriate tests

* fixed tests

* fixed tests and implementation

* added rng

* Update src/inference/Inference.jl

Co-authored-by: David Widmann <[email protected]>

* version bump

Co-authored-by: David Widmann <[email protected]>

Author: torfjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.14.5"
+version = "0.14.6"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/Turing.jl

@@ -68,6 +68,7 @@ end
 ###########
 # `using` statements for stuff to re-export
 using DynamicPPL: elementwise_loglikelihoods, generated_quantities, logprior, logjoint
+using StatsBase: predict
 
 # Turing essentials - modelling macros and inference algorithms
 export  @model,                 # modelling
@@ -120,6 +121,7 @@ export  @model,                 # modelling
         filldist,
         arraydist,
 
+        predict,
         elementwise_loglikelihoods,
         genereated_quantities,
         logprior,

src/inference/Inference.jl

@@ -599,7 +599,7 @@ DynamicPPL.inspace(vn::VarName, spl::Sampler) = inspace(vn, getspace(spl.alg))
 
 """
 
-    predict(model::Model, chain::MCMCChains.Chains; include_all=false)
+    predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)
 
 Execute `model` conditioned on each sample in `chain`, and return the resulting `Chains`.
 
@@ -637,7 +637,7 @@ julia> chain_lin_reg = sample(m_train, NUTS(100, 0.65), 200);
 
 julia> m_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)), σ);
 
-julia> predictions = Turing.Inference.predict(m_test, chain_lin_reg)
+julia> predictions = predict(m_test, chain_lin_reg)
 Object of type Chains, with data of type 100×2×1 Array{Float64,3}
 
 Iterations        = 1:100
@@ -667,20 +667,30 @@ julia> sum(abs2, ys_test - ys_pred) ≤ 0.1
 true
 ```
 """
-function predict(model::Turing.Model, chain::MCMCChains.Chains; include_all = false)
+function predict(model::Model, chain::MCMCChains.Chains; kwargs...)
+    return predict(Random.GLOBAL_RNG, model, chain; kwargs...)
+end
+function predict(rng::AbstractRNG, model::Model, chain::MCMCChains.Chains; include_all = false)
     spl = DynamicPPL.SampleFromPrior()
 
     # Sample transitions using `spl` conditioned on values in `chain`
-    transitions = transitions_from_chain(model, chain; sampler = spl)
+    transitions = [
+        transitions_from_chain(rng, model, chain[:, :, chn_idx]; sampler = spl)
+        for chn_idx = 1:size(chain, 3)
+    ]
 
     # Let the Turing internals handle everything else for you
-    chain_result = AbstractMCMC.bundle_samples(
-        Distributions.GLOBAL_RNG,
-        model,
-        spl,
-        length(chain),
-        transitions,
-        MCMCChains.Chains
+    chain_result = reduce(
+        MCMCChains.chainscat, [
+            AbstractMCMC.bundle_samples(
+                rng,
+                model,
+                spl,
+                length(chain),
+                transitions[chn_idx],
+                MCMCChains.Chains
+            ) for chn_idx = 1:size(chain, 3)
+        ]
     )
 
     parameter_names = if include_all
@@ -695,6 +705,7 @@ end
 """
 
     transitions_from_chain(
+        [rng::AbstractRNG,]
         model::Model, 
         chain::MCMCChains.Chains; 
         sampler = DynamicPPL.SampleFromPrior()
@@ -741,6 +752,14 @@ julia> [first(t.θ.x) for t in transitions] # extract samples for `x`
 ```
 """
 function transitions_from_chain(
+    model::Turing.Model,
+    chain::MCMCChains.Chains;
+    kwargs...
+)
+    return transitions_from_chain(Random.GLOBAL_RNG, model, chain; kwargs...)
+end
+function transitions_from_chain(
+    rng::AbstractRNG,
     model::Turing.Model,
     chain::MCMCChains.Chains;
     sampler = DynamicPPL.SampleFromPrior()
@@ -774,7 +793,7 @@ function transitions_from_chain(
             end
         end
         # Execute `model` on the parameters set in `vi` and sample those with `"del"` flag using `sampler`
-        model(vi, sampler)
+        model(rng, vi, sampler)
 
         # Convert `VarInfo` into `NamedTuple` and save
         theta = DynamicPPL.tonamedtuple(vi)

test/inference/utilities.jl

@@ -27,17 +27,50 @@ using Random
     chain_lin_reg = sample(m_lin_reg, NUTS(100, 0.65), 200);
 
     # Predict on two last indices
-    m_lin_reg_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)));
+    m_lin_reg_test = linear_reg(xs_test, fill(missing, length(ys_test)));    
     predictions = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)
 
     ys_pred = vec(mean(Array(group(predictions, :y)); dims = 1))
 
     @test sum(abs2, ys_test - ys_pred) ≤ 0.1
 
+    # Ensure that `rng` is respected
+    predictions1 = let rng = MersenneTwister(42)
+        predict(rng, m_lin_reg_test, chain_lin_reg[1:2])
+    end
+    predictions2 = let rng = MersenneTwister(42)
+        predict(rng, m_lin_reg_test, chain_lin_reg[1:2])
+    end
+    @test all(Array(predictions1) .== Array(predictions2))
+
     # Predict on two last indices for vectorized
     m_lin_reg_test = linear_reg_vec(xs_test, missing);
     predictions_vec = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)
     ys_pred_vec = vec(mean(Array(group(predictions_vec, :y)); dims = 1))
 
     @test sum(abs2, ys_test - ys_pred_vec) ≤ 0.1
+
+    # Multiple chains
+    chain_lin_reg = sample(m_lin_reg, NUTS(100, 0.65), MCMCThreads(), 200, 2);
+    m_lin_reg_test = linear_reg(xs_test, fill(missing, length(ys_test)));
+    predictions = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)
+
+    @test size(chain_lin_reg, 3) == size(predictions, 3)
+
+    for chain_idx in MCMCChains.chains(chain_lin_reg)
+        ys_pred = vec(mean(Array(group(predictions[:, :, chain_idx], :y)); dims = 1))
+        @test sum(abs2, ys_test - ys_pred) ≤ 0.1
+    end
+
+    # Predict on two last indices for vectorized
+    m_lin_reg_test = linear_reg_vec(xs_test, missing);
+    predictions_vec = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)
+
+    for chain_idx in MCMCChains.chains(chain_lin_reg)
+        ys_pred_vec = vec(mean(
+            Array(group(predictions_vec[:, :, chain_idx], :y));
+            dims = 1
+        ))
+        @test sum(abs2, ys_test - ys_pred_vec) ≤ 0.1
+    end
 end