Add chain_number keyword argument when performing multi-chain sampling

Author: penelopeysm

Project.toml

@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probabilistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "5.7.2"
+version = "5.7.3"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"

src/sample.jl

@@ -549,6 +549,7 @@ function mcmcsample(
                                 else
                                     initial_state[chainidx]
                                 end,
+                                chain_number=chainidx,
                                 kwargs...,
                             )
                         end
@@ -669,7 +670,7 @@ function mcmcsample(
             Distributed.@async begin
                 try
                     function sample_chain(
-                        seed, initial_params, initial_state, child_progress
+                        seed, initial_params, initial_state, child_progress, chainidx
                     )
                         # Seed a new random number generator with the pre-made seed.
                         Random.seed!(rng, seed)
@@ -683,6 +684,7 @@ function mcmcsample(
                             progress=child_progress,
                             initial_params=initial_params,
                             initial_state=initial_state,
+                            chain_number=chainidx,
                             kwargs...,
                         )
 
@@ -696,6 +698,7 @@ function mcmcsample(
                         _initial_params,
                         _initial_state,
                         child_progresses,
+                        1:nchains;
                     )
                 finally
                     if progress == :overall
@@ -755,6 +758,7 @@ function mcmcsample(
             progressname=string(progressname, " (Chain ", i, " of ", nchains, ")"),
             initial_params=initial_params,
             initial_state=initial_state,
+            chain_number=i,
             kwargs...,
         )
     end

test/sample.jl

@@ -679,6 +679,19 @@
         @test all(chain[i].b == ref_chain[i].b for i in 1:N)
     end
 
+    @testset "chain_number keyword argument" begin
+        @testset for m in [MCMCSerial(), MCMCThreads(), MCMCDistributed()]
+            # check that the `chain_number` keyword argument is passed to the callback
+            chain_numbers = Int[]
+            function callback(args...; kwargs...)
+                @test haskey(kwargs, :chain_number)
+                return push!(chain_numbers, kwargs[:chain_number])
+            end
+            chain = sample(MyModel(), MySampler(), m, 10, 4; callback=callback)
+            @test sort(chain_numbers) == repeat(1:4; inner=10)
+        end
+    end
+
     @testset "Sample vector of `NamedTuple`s" begin
         chain = sample(MyModel(), MySampler(), 1_000; chain_type=Vector{NamedTuple})
         # Check output type