Merge pull request #4 from FredericWantiez/feature/split

Fix naming

Author: FredericWantiez

src/container.jl

@@ -32,8 +32,8 @@ 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 ? load_state(t.rng) : save_state!(t.rng)
-    inc_count!(t.rng)
+    isref ? load_state!(t.rng) : save_state!(t.rng)
+    inc_counter!(t.rng)
 
     # Move to next step
     return Libtask.consume(t.ctask)
@@ -60,7 +60,7 @@ end
 # Create new task and copy randomness
 function forkr(trace::Trace)
     newf = reset_model(trace.f)
-    set_count!(trace.rng, 1)
+    set_counter!(trace.rng, 1)
 
     ctask = let f = trace.ctask.task.code
         Libtask.CTask() do
@@ -100,7 +100,7 @@ mutable struct ParticleContainer{T<:Particle}
     vals::Vector{T}
     "Unnormalized logarithmic weights."
     logWs::Vector{Float64}
-    "Traced RNG"
+    "Traced RNG to replay the resampling step"
     rng::TracedRNG
 end
 
@@ -204,8 +204,8 @@ 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, 1)
-        update_rng!(pi.rng, k[1])
+        k = split(pi.rng.rng.key)
+        seed!(pi.rng, k[1])
     end
 end
 
@@ -252,15 +252,16 @@ function resample_propagate!(
             pi = particles[i]
             isref = pi === ref
             p = isref ? fork(pi, isref) : pi
+            nseeds = isref ? ni - 1 : ni
 
-            seeds = split(p.rng, ni)
-            !isref && update_rng!(p.rng, seeds[1])
+            seeds = split(p.rng.rng.key, nseeds)
+            !isref && seed!(p.rng, seeds[1])
 
             children[j += 1] = p
             # fork additional children
             for k in 2:ni
                 part = fork(p, isref)
-                update_rng!(part.rng, seeds[k])
+                seed!(part.rng, seeds[k])
                 children[j += 1] = part
             end
         end

src/rng.jl

@@ -4,6 +4,7 @@ using Distributions
 
 import Base.rand
 import Random.seed!
+import Random123: set_counter!
 
 # Default RNG type for when nothing is specified
 _BASE_RNG = Philox2x
@@ -29,7 +30,7 @@ end
 Initialize TracedRNG with r as the inner RNG 
 """
 function TracedRNG(r::Random123.AbstractR123)
-    set_counter!(r, 1)
+    set_counter!(r, 0)
     return TracedRNG(1, r, typeof(r.key)[])
 end
 
@@ -48,42 +49,33 @@ 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(key::Integer, n::Integer)
 
-Split inner RNG into n new TracedRNG
+Split key into n new keys
 """
-function split(r::TracedRNG{T}, n::Integer) where {T}
-    return map(i -> hash(r.rng.key, convert(UInt, r.rng.ctr1 + i)), 1:n)
+function split(key::Integer, n::Integer=1) where {T}
+    return map(i -> hash(key, convert(UInt, i)), 1:n)
 end
 
 """
-    seed!(r::TracedRNG, seed::Number)
-
-Set the key of the inner RNG as `seed`
-"""
-function seed!(r::TracedRNG{T}, seed) where {T}
-    return seed!(r.rng, seed)
-end
-
-"""
-    load_state(r::TracedRNG, seed)
+    load_state!(r::TracedRNG, seed)
 
 Load state from current model iteration. Random streams are now replayed
 """
-function load_state(rng::TracedRNG{T}) where {T}
+function load_state!(rng::TracedRNG{T}) where {T}
     key = rng.keys[rng.count]
     Random.seed!(rng.rng, key)
-    return set_counter!(rng.rng, rng.count)
+    return set_counter!(rng.rng, 0)
 end
 
 """
     update_rng!(rng::TracedRNG)
 
 Set key and counter of inner RNG to key and the running model step
 """
-function update_rng!(rng::TracedRNG{T}, key) where {T}
-    seed!(rng, key)
-    return set_counter!(rng.rng, rng.count)
+function seed!(rng::TracedRNG{T}, key) where {T}
+    seed!(rng.rng, key)
+    return set_counter!(rng.rng, 0)
 end
 
 """ 
@@ -102,7 +94,11 @@ Base.copy(r::TracedRNG{T}) where {T} = TracedRNG(r.count, copy(r.rng), copy(r.ke
 
 Set the counter of the TracedRNG, used to keep track of the current model step
 """
-set_count!(r::TracedRNG, n::Integer) = r.count = n
+set_counter!(r::TracedRNG, n::Integer) = r.count = n
+
+"""
+    inc_counter!(r::TracedRNG, n::Integer=1)
 
-inc_count!(r::TracedRNG, n::Integer) = r.count += n
-inc_count!(r::TracedRNG) = inc_count!(r, 1)
+Increase the model step counter by n
+"""
+inc_counter!(r::TracedRNG, n::Integer=1) = r.count += n

test/rng.jl

@@ -6,8 +6,19 @@
 
         rand(rng, Distributions.Normal())
 
-        AdvancedPS.load_state(rng)
+        AdvancedPS.load_state!(rng)
         new_vns = rand(rng, Distributions.Normal())
         @test new_vns ≈ vns
     end
+
+    @testset "split" begin
+        rng = AdvancedPS.TracedRNG()
+        key = rng.rng.key
+        new_key, = AdvancedPS.split(key, 1)
+
+        @test key ≠ new_key
+
+        AdvancedPS.seed!(rng, new_key)
+        @test rng.rng.key === new_key
+    end
 end