Track Particle container

Author: FredericWantiez

src/container.jl

@@ -89,15 +89,23 @@ Data structure for particle filters
 - normalise!(pc::ParticleContainer)
 - consume(pc::ParticleContainer): return incremental likelihood
 """
-mutable struct ParticleContainer{T<:Particle}
+mutable struct ParticleContainer{T<:Particle,R<:Random.AbstractRNG}
     "Particles."
     vals::Vector{T}
     "Unnormalized logarithmic weights."
     logWs::Vector{Float64}
+    "TracedRNG to track the resampling step"
+    rng::TracedRNG{R}
 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}, rng::T
+) where {T<:Random.AbstractRNG}
+    return ParticleContainer(particles, zeros(length(particles)), TracedRNG(rng))
 end
 
 Base.collect(pc::ParticleContainer) = pc.vals
@@ -124,7 +132,7 @@ function Base.copy(pc::ParticleContainer)
     # copy weights
     logWs = copy(pc.logWs)
 
-    return ParticleContainer(vals, logWs)
+    return ParticleContainer(vals, logWs, pc.rng)
 end
 
 """

src/rng.jl

@@ -2,10 +2,11 @@
 Data structure to keep track of the history of the random stream
 produced by RNG.
 """
-struct TracedRNG{T} <: Random.AbstractRNG where {T<:Random.AbstractRNG}
+mutable struct TracedRNG{T} <: Random.AbstractRNG where {T<:Random.AbstractRNG}
     count::Base.RefValue{Int}
     rng::T
-    seed::Any
+    seed::Array
+    states::Array{T}
 end
 
 # Set seed manually, for init ?
@@ -18,7 +19,7 @@ end
 Random.seed!(rng::TracedRNG) = Random.seed!(rng.rng, rng.seed)
 
 TracedRNG() = TracedRNG(Random.MersenneTwister()) # Pick up an explicit RNG from Random
-TracedRNG(rng::Random.AbstractRNG) = TracedRNG(Ref(0), rng, rng.seed)
+TracedRNG(rng::Random.AbstractRNG) = TracedRNG(Ref(0), rng, rng.seed, [rng])
 TracedRNG(rng::Random._GLOBAL_RNG) = TracedRNG(Random.default_rng())
 
 # Intercept rand
@@ -28,6 +29,7 @@ Random.rng_native_52(r::TracedRNG) = UInt64
 function Base.rand(rng::TracedRNG, ::Type{T}) where {T}
     res = Base.rand(rng.rng, T)
     inc_count!(rng, length(res))
+    push!(rng.states, copy(rng.rng))
     return res
 end
 

src/smc.jl

@@ -38,12 +38,12 @@ 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)], rng
+    )
 
     # Perform particle sweep.
-    logevidence = sweep!(rng, particles, sampler.resampler)
+    logevidence = sweep!(particles.rng, particles, sampler.resampler)
 
     return SMCSample(collect(particles), getweights(particles), logevidence)
 end
@@ -85,12 +85,12 @@ 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)], rng
+    )
 
     # Perform a particle sweep.
-    logevidence = sweep!(rng, particles, sampler.resampler)
+    logevidence = sweep!(particles.rng, particles, sampler.resampler)
 
     # Pick a particle to be retained.
     trajectory = rand(rng, particles)
@@ -115,10 +115,10 @@ function AbstractMCMC.step(
             Trace(model, TracedRNG())
         end
     end
-    particles = ParticleContainer(x)
+    particles = ParticleContainer(x, rng)
 
     # Perform a particle sweep.
-    logevidence = sweep!(rng, particles, sampler.resampler)
+    logevidence = sweep!(particles.rng, particles, sampler.resampler)
 
     # Pick a particle to be retained.
     newtrajectory = rand(rng, particles)