Correct replaying mechanism

Author: FredericWantiez

src/container.jl

@@ -32,7 +32,7 @@ Base.copy(trace::Trace) = Trace(trace.f, copy(trace.ctask), deepcopy(trace.rng))
 # step to the next observe statement and
 # return the log probability of the transition (or nothing if done)
 function advance!(t::Trace, isref::Bool)
-    isref ? reset_rng!(t.rng) : save_state!(t.rng)
+    isref ? load_state(t.rng) : save_state!(t.rng)
     inc_count!(t.rng)
 
     # Move to next step
@@ -100,10 +100,16 @@ mutable struct ParticleContainer{T<:Particle}
     vals::Vector{T}
     "Unnormalized logarithmic weights."
     logWs::Vector{Float64}
+    "Traced RNG"
+    rng::TracedRNG
 end
 
 function ParticleContainer(particles::Vector{<:Particle})
-    return ParticleContainer(particles, zeros(length(particles)))
+    return ParticleContainer(particles, zeros(length(particles)), TracedRNG())
+end
+
+function ParticleContainer(particles::Vector{<:Particle}, r::TracedRNG)
+    return ParticleContainer(particles, zeros(length(particles)), r)
 end
 
 Base.collect(pc::ParticleContainer) = pc.vals
@@ -130,7 +136,10 @@ function Base.copy(pc::ParticleContainer)
     # copy weights
     logWs = copy(pc.logWs)
 
-    return ParticleContainer(vals, logWs)
+    # Copy rng and states
+    rng = copy(pc.rng)
+
+    return ParticleContainer(vals, logWs, rng)
 end
 
 """
@@ -184,6 +193,21 @@ function effectiveSampleSize(pc::ParticleContainer)
     return inv(sum(abs2, Ws))
 end
 
+"""
+    update_keys!(pc::ParticleContainer)
+
+Create new unique keys for the particles in the ParticleContainer
+"""
+function update_keys!(pc::ParticleContainer)
+    # Update keys to new particle ids
+    for i in 1:length(pc)
+        pi = pc.vals[i]
+        k = split(pi.rng, 1)
+        seed!(pi.rng, k[1])
+        set_counter!(pi.rng.rng, pi.rng.count)
+    end
+end
+
 """
     resample_propagate!(rng, pc::ParticleContainer[, randcat = resample_systematic,
                         ref = nothing; weights = getweights(pc)])
@@ -227,9 +251,11 @@ function resample_propagate!(
             pi = particles[i]
             isref = pi === ref
             p = isref ? fork(pi, isref) : pi
-            children[j += 1] = p
 
-            seeds = split(pi.rng, ni)
+            seeds = split(p.rng, ni)
+            seed!(p.rng, seeds[1])
+
+            children[j += 1] = p
             # fork additional children
             for k in 2:ni
                 part = fork(p, isref)
@@ -264,6 +290,8 @@ function resample_propagate!(
 
     if ess ≤ resampler.threshold * length(pc)
         resample_propagate!(rng, pc, resampler.resampler, ref; weights=weights)
+    else
+        update_keys!(pc)
     end
 
     return pc

src/rng.jl

@@ -5,8 +5,8 @@ using Distributions
 import Base.rand
 import Random.seed!
 
-# Use Philox2x for now
-BASE_RNG = Philox2x
+# Default RNG type for when nothing is specified
+_BASE_RNG = Philox2x
 
 """
     TracedRNG{R,T}
@@ -15,14 +15,22 @@ Wrapped random number generator from Random123 to keep track of random streams d
 """
 mutable struct TracedRNG{T} <:
                Random.AbstractRNG where {T<:(Random123.AbstractR123{R} where {R})}
+    "Model step counter"
     count::Int
+    "Inner RNG"
     rng::T
+    "Array of keys"
     keys
-    counters
 end
 
+"""
+    TracedRNG(r::Random123.AbstractR123)
+
+Initialize TracedRNG with r as the inner RNG 
+"""
 function TracedRNG(r::Random123.AbstractR123)
-    return TracedRNG(1, r, typeof(r.key)[], typeof(r.ctr1)[])
+    set_counter!(r, 1)
+    return TracedRNG(1, r, typeof(r.key)[])
 end
 
 """
@@ -31,48 +39,50 @@ end
 Create a default TracedRNG
 """
 function TracedRNG()
-    r = BASE_RNG()
+    r = _BASE_RNG()
     return TracedRNG(r)
 end
 
-# Plug into Random
+# Connect to the Random API
 Random.rng_native_52(rng::TracedRNG{U}) where {U} = Random.rng_native_52(rng.rng)
 Base.rand(rng::TracedRNG{U}, ::Type{T}) where {U,T} = Base.rand(rng.rng, T)
 
 """
     split(r::TracedRNG, n::Integer)
 
-Split keys of the internal Philox2x into n distinct seeds
+Split inner RNG into n new TracedRNG
 """
 function split(r::TracedRNG{T}, n::Integer) where {T}
-    n == 1 && return [r.rng.key]
     return map(i -> hash(r.rng.key, convert(UInt, r.rng.ctr1 + i)), 1:n)
 end
 
 """
-    update_rng!(r::TracedRNG, seed::Number)
+    seed!(r::TracedRNG, seed::Number)
 
-Set the key of the wrapped Philox2x rng
+Set the key of the inner RNG as `seed`
 """
 function seed!(r::TracedRNG{T}, seed) where {T}
     return seed!(r.rng, seed)
 end
 
 """
-    reset_rng(r::TracedRNG, seed)
+    load_state(r::TracedRNG, seed)
 
-Reset the rng to the running model step
+Load state from current model iteration. Random streams are now replayed
 """
-function reset_rng!(rng::TracedRNG{T}) where {T}
+function load_state(rng::TracedRNG{T}) where {T}
     key = rng.keys[rng.count]
-    ctr = rng.counters[rng.count]
     Random.seed!(rng.rng, key)
-    return set_counter!(rng.rng, ctr)
+    return set_counter!(rng.rng, rng.count)
 end
 
+""" 
+    save_state!(r::TracedRNG)
+
+Track current key of the inner RNG
+"""
 function save_state!(r::TracedRNG{T}) where {T}
-    push!(r.keys, r.rng.key)
-    return push!(r.counters, r.rng.ctr1)
+    return push!(r.keys, r.rng.key)
 end
 
 Base.copy(r::TracedRNG{T}) where {T} = TracedRNG(r.count, copy(r.rng), copy(r.keys))

src/smc.jl

@@ -38,9 +38,9 @@ function AbstractMCMC.sample(
     end
 
     # Create a set of particles.
-    particles = ParticleContainer([
-        Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)
-    ])
+    particles = ParticleContainer(
+        [Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)], TracedRNG()
+    )
 
     # Perform particle sweep.
     logevidence = sweep!(rng, particles, sampler.resampler)
@@ -85,9 +85,9 @@ function AbstractMCMC.step(
     rng::Random.AbstractRNG, model::AbstractMCMC.AbstractModel, sampler::PG; kwargs...
 )
     # Create a new set of particles.
-    particles = ParticleContainer([
-        Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)
-    ])
+    particles = ParticleContainer(
+        [Trace(model, TracedRNG()) for _ in 1:(sampler.nparticles)], TracedRNG()
+    )
 
     # Perform a particle sweep.
     logevidence = sweep!(rng, particles, sampler.resampler)
@@ -115,7 +115,7 @@ function AbstractMCMC.step(
             Trace(model, TracedRNG())
         end
     end
-    particles = ParticleContainer(x)
+    particles = ParticleContainer(x, TracedRNG())
 
     # Perform a particle sweep.
     logevidence = sweep!(rng, particles, sampler.resampler, particles.vals[nparticles])

test/rng.jl

@@ -6,7 +6,7 @@
 
         rand(rng, Distributions.Normal())
 
-        AdvancedPS.reset_rng!(rng)
+        AdvancedPS.load_state(rng)
         new_vns = rand(rng, Distributions.Normal())
         @test new_vns ≈ vns
     end

test/smc.jl

@@ -150,6 +150,34 @@
         @test all(isone(p.trajectory.f.x) for p in chains_pg)
         @test mean(x.logevidence for x in chains_pg) ≈ -2 * log(2) atol = 0.01
     end
+    
+    @testset "Replay reference" begin
+        mutable struct Model <: AbstractMCMC.AbstractModel
+            a::Float64
+            b::Float64
+
+            Model() = new()
+        end
+        
+        function (m::Model)(rng)
+            m.a = rand(rng, Normal())
+            AdvancedPS.observe(Normal(), m.a)
+
+            m.b = rand(rng, Normal())
+            AdvancedPS.observe(Normal(), m.b)
+        end
+
+        pg = AdvancedPS.PG(1)
+        first, second = sample(Model(), pg, 2);
+
+        first_model = first.trajectory.f
+        second_model = second.trajectory.f
+
+        # Single Particle - must be replaying
+        @test first_model.a ≈ second_model.a
+        @test first_model.b ≈ second_model.b
+        @test first.logevidence ≈ second.logevidence
+    end
 end
 
 # @testset "pmmh.jl" begin