Format

Author: FredericWantiez

src/container.jl

@@ -6,19 +6,17 @@ end
 
 const Particle = Trace
 
-function Trace(f, rng::Random.AbstractRNG)
-    trng = TracedRNG(rng)
-
+function Trace(f, rng::TracedRNG)
     ctask = let f = f
         Libtask.CTask() do
-            res = f(trng)
+            res = f(rng)
             Libtask.produce(nothing)
             return res
         end
     end
 
     # add backward reference
-    newtrace = Trace(f, ctask, trng)
+    newtrace = Trace(f, ctask, rng)
     addreference!(ctask.task, newtrace)
 
     return newtrace

src/rng.jl

@@ -8,9 +8,12 @@ struct TracedRNG{T} <: Random.AbstractRNG where {T<:Random.AbstractRNG}
     seed::Any
 end
 
-
 # Set seed manually, for init ?
-Random.seed!(rng::TracedRNG, seed) = Random.seed!(rng.rng, seed)
+function Random.seed!(rng::TracedRNG, seed)
+    rng.rng.seed = seed
+    return Random.seed!(rng.rng, seed)
+end
+
 # Reset the rng to the initial seed
 Random.seed!(rng::TracedRNG) = Random.seed!(rng.rng, rng.seed)
 

src/smc.jl

@@ -35,10 +35,12 @@ function AbstractMCMC.sample(
 )
     if !isempty(kwargs)
         @warn "keyword arguments $(keys(kwargs)) are not supported by `SMC`"
-	end
+    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)
+    ])
 
     # Perform particle sweep.
     logevidence = sweep!(rng, particles, sampler.resampler)
@@ -83,7 +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)
+    ])
 
     # Perform a particle sweep.
     logevidence = sweep!(rng, particles, sampler.resampler)

test/container.jl

@@ -22,7 +22,7 @@
 
         # Create particle container.
         logps = [0.0, -1.0, -2.0]
-		particles = [AdvancedPS.Trace(fpc(logp), Random.MersenneTwister()) for logp in logps]
+        particles = [AdvancedPS.Trace(fpc(logp), AdvancedPS.TracedRNG()) for logp in logps]
         pc = AdvancedPS.ParticleContainer(particles)
 
         # Initial state.
@@ -94,7 +94,7 @@
         end
 
         # Test task copy version of trace
-		tr = AdvancedPS.Trace(f2, Random.MersenneTwister())
+        tr = AdvancedPS.Trace(f2, AdvancedPS.TracedRNG())
 
         consume(tr.ctask)
         consume(tr.ctask)

test/runtests.jl

@@ -15,7 +15,7 @@ using Test
     @testset "SMC and PG tests" begin
         include("smc.jl")
     end
-    @testset "RNG tests" begin 
-        include("rng.jl") 
+    @testset "RNG tests" begin
+        include("rng.jl")
     end
 end