Seed particles from top-level RNG (#41)

* Set seed from parent RNG

* Decouple seed and sampler

* No-allocation

* Update Project.toml

Co-authored-by: Hong Ge <[email protected]>

Author: FredericWantiez

Project.toml

@@ -14,7 +14,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 [compat]
 AbstractMCMC = "2, 3"
 Distributions = "0.23, 0.24, 0.25"
-Libtask = "0.6"
+Libtask = "0.6.7"
 Random123 = "1.3"
 StatsFuns = "0.9"
 julia = "1.3"

src/container.jl

@@ -94,8 +94,15 @@ function ParticleContainer(particles::Vector{<:Particle})
     return ParticleContainer(particles, zeros(length(particles)), TracedRNG())
 end
 
-function ParticleContainer(particles::Vector{<:Particle}, r::TracedRNG)
-    return ParticleContainer(particles, zeros(length(particles)), r)
+function ParticleContainer(particles::Vector{<:Particle}, rng::TracedRNG)
+    return ParticleContainer(particles, zeros(length(particles)), rng)
+end
+
+function ParticleContainer(
+    particles::Vector{<:Particle}, trng::TracedRNG, rng::Random.AbstractRNG
+)
+    pc = ParticleContainer(particles, trng)
+    return seed_from_rng!(pc, rng)
 end
 
 Base.collect(pc::ParticleContainer) = pc.vals
@@ -190,12 +197,35 @@ function update_keys!(pc::ParticleContainer, ref::Union{Particle,Nothing}=nothin
     n = ref === nothing ? nparticles : nparticles - 1
     for i in 1:n
         pi = pc.vals[i]
-        k = split(pi.rng.rng.key)
+        k = split(state(pi.rng.rng))
         Random.seed!(pi.rng, k[1])
     end
     return nothing
 end
 
+"""
+    seed_from_rng!(pc::ParticleContainer, rng::Random.AbstractRNG, ref::Union{Particle,Nothing}=nothing)
+
+Set seeds of particle rng from user-provided `rng`
+"""
+function seed_from_rng!(
+    pc::ParticleContainer{T,<:TracedRNG{R,N,<:Random123.AbstractR123{I}}},
+    rng::Random.AbstractRNG,
+    ref::Union{Particle,Nothing}=nothing,
+) where {T,R,N,I}
+    n = length(pc.vals)
+    nseeds = isnothing(ref) ? n : n - 1
+
+    sampler = Random.Sampler(rng, I)
+    for i in 1:nseeds
+        subrng = pc.vals[i].rng
+        Random.seed!(subrng, gen_seed(rng, subrng, sampler))
+    end
+    Random.seed!(pc.rng, gen_seed(rng, pc.rng, sampler))
+
+    return pc
+end
+
 """
     resample_propagate!(rng, pc::ParticleContainer[, randcat = resample_systematic,
                         ref = nothing; weights = getweights(pc)])

src/rng.jl

@@ -1,12 +1,12 @@
 # Default RNG type for when nothing is specified
-const _BASE_RNG = Random123.Philox2x
+const _BASE_RNG = Random123.Philox2x # Rng with state bigger than 1 are broken because of Split
 
 """
     TracedRNG{R,N,T}
 
 Wrapped random number generator from Random123 to keep track of random streams during model evaluation
 """
-mutable struct TracedRNG{R,N,T<:Random123.AbstractR123{R}} <: Random.AbstractRNG
+mutable struct TracedRNG{R,N,T<:Random123.AbstractR123} <: Random.AbstractRNG
     "Model step counter"
     count::Int
     "Inner RNG"
@@ -21,7 +21,7 @@ Create a `TracedRNG` with `r` as the inner RNG.
 """
 function TracedRNG(r::Random123.AbstractR123=_BASE_RNG())
     Random123.set_counter!(r, 0)
-    return TracedRNG(1, r, typeof(r.key)[])
+    return TracedRNG(1, r, Random123.seed_type(r)[])
 end
 
 # Connect to the Random API
@@ -50,7 +50,7 @@ function load_state!(rng::TracedRNG)
 end
 
 """
-    update_rng!(rng::TracedRNG)
+    Random.seed!(rng::TracedRNG, key)
 
 Set key and counter of inner rng in `rng` to `key` and the running model step to 0
 """
@@ -59,15 +59,33 @@ function Random.seed!(rng::TracedRNG, key)
     return Random123.set_counter!(rng.rng, 0)
 end
 
+"""
+    gen_seed(rng::Random.AbstractRNG, subrng::TracedRNG, sampler::Random.Sampler)
+
+Generate a `seed` for the subrng based on top-level `rng` and `sampler`
+"""
+function gen_seed(rng::Random.AbstractRNG, ::TracedRNG{<:Integer}, sampler::Random.Sampler)
+    return rand(rng, sampler)
+end
+
+function gen_seed(
+    rng::Random.AbstractRNG, ::TracedRNG{<:NTuple{N}}, sampler::Random.Sampler
+) where {N}
+    return Tuple(rand(rng, sampler, N))
+end
+
 """ 
     save_state!(r::TracedRNG)
 
 Add current key of the inner rng in `r` to `keys`.
 """
 function save_state!(r::TracedRNG)
-    return push!(r.keys, r.rng.key)
+    return push!(r.keys, state(r.rng))
 end
 
+state(rng::Random123.Philox2x) = rng.key
+state(rng::Random123.Philox4x) = (rng.key1, rng.key2)
+
 Base.copy(r::TracedRNG) = TracedRNG(r.count, copy(r.rng), deepcopy(r.keys))
 
 """

src/smc.jl

@@ -39,7 +39,7 @@ function AbstractMCMC.sample(
 
     # Create a set of particles.
     particles = ParticleContainer(
-        [Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)], TracedRNG()
+        [Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)], TracedRNG(), rng
     )
 
     # Perform particle sweep.
@@ -86,7 +86,7 @@ function AbstractMCMC.step(
 )
     # Create a new set of particles.
     particles = ParticleContainer(
-        [Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)], TracedRNG()
+        [Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)], TracedRNG(), rng
     )
 
     # Perform a particle sweep.
@@ -107,6 +107,7 @@ function AbstractMCMC.step(
 )
     # Create a new set of particles.
     nparticles = sampler.nparticles
+
     x = map(1:nparticles) do i
         if i == nparticles
             # Create reference trajectory.
@@ -115,10 +116,12 @@ function AbstractMCMC.step(
             Trace(model, TracedRNG())
         end
     end
-    particles = ParticleContainer(x, TracedRNG())
+
+    reference = x[end]
+    particles = ParticleContainer(x, TracedRNG(), rng)
 
     # Perform a particle sweep.
-    logevidence = sweep!(rng, particles, sampler.resampler, particles.vals[nparticles])
+    logevidence = sweep!(rng, particles, sampler.resampler, reference)
 
     # Pick a particle to be retained.
     newtrajectory = rand(rng, particles)

test/container.jl

@@ -121,4 +121,32 @@
         @test consume(tr.ctask) == 2
         @test consume(a.ctask) == 4
     end
+
+    @testset "seed container" begin
+        function dummy(rng) end
+
+        seed = 1
+        n = 3
+        rng = Random.MersenneTwister(seed)
+
+        particles = [AdvancedPS.Trace(dummy, AdvancedPS.TracedRNG()) for _ in 1:n]
+        pc = AdvancedPS.ParticleContainer(particles, AdvancedPS.TracedRNG())
+
+        AdvancedPS.seed_from_rng!(pc, rng)
+        old_seeds = vcat([part.rng.rng.key for part in pc.vals], [pc.rng.rng.key])
+
+        Random.seed!(rng, seed)
+        AdvancedPS.seed_from_rng!(pc, rng)
+        new_seeds = vcat([part.rng.rng.key for part in pc.vals], [pc.rng.rng.key])
+
+        # Check if we reset the seeds properly
+        @test old_seeds ≈ new_seeds
+
+        Random.seed!(rng, 2)
+        AdvancedPS.seed_from_rng!(pc, rng, pc.vals[n])
+        ref_seeds = vcat([part.rng.rng.key for part in pc.vals], [pc.rng.rng.key])
+
+        # Dont reset reference particle
+        @test ref_seeds[n] ≈ new_seeds[n]
+    end
 end

test/smc.jl

@@ -7,12 +7,12 @@
         sampler = AdvancedPS.SMC(15, 0.6)
         @test sampler.nparticles == 15
         @test sampler.resampler ===
-              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)
+            AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)
 
         sampler = AdvancedPS.SMC(20, AdvancedPS.resample_multinomial, 0.6)
         @test sampler.nparticles == 20
         @test sampler.resampler ===
-              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)
+            AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)
 
         sampler = AdvancedPS.SMC(25, AdvancedPS.resample_systematic)
         @test sampler.nparticles == 25
@@ -105,12 +105,12 @@
         sampler = AdvancedPS.PG(60, 0.6)
         @test sampler.nparticles == 60
         @test sampler.resampler ===
-              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)
+            AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)
 
         sampler = AdvancedPS.PG(80, AdvancedPS.resample_multinomial, 0.6)
         @test sampler.nparticles == 80
         @test sampler.resampler ===
-              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)
+            AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)
 
         sampler = AdvancedPS.PG(100, AdvancedPS.resample_systematic)
         @test sampler.nparticles == 100