Merge pull request #5 from FredericWantiez/feature/split

Fix types, import and PR review

Author: FredericWantiez

src/AdvancedPS.jl

@@ -5,6 +5,7 @@ using Distributions: Distributions
 using Libtask: Libtask
 using Random: Random
 using StatsFuns: StatsFuns
+using Random123: Random123
 
 include("resampling.jl")
 include("rng.jl")

src/container.jl

@@ -1,7 +1,7 @@
-struct Trace{F}
+struct Trace{F,U,N,V<:Random123.AbstractR123{U}}
     f::F
     ctask::Libtask.CTask
-    rng::TracedRNG
+    rng::TracedRNG{U,N,V}
 end
 
 const Particle = Trace
@@ -60,7 +60,7 @@ end
 # Create new task and copy randomness
 function forkr(trace::Trace)
     newf = reset_model(trace.f)
-    set_counter!(trace.rng, 1)
+    Random123.set_counter!(trace.rng, 1)
 
     ctask = let f = trace.ctask.task.code
         Libtask.CTask() do
@@ -95,13 +95,13 @@ Data structure for particle filters
 - normalise!(pc::ParticleContainer)
 - consume(pc::ParticleContainer): return incremental likelihood
 """
-mutable struct ParticleContainer{T<:Particle}
+mutable struct ParticleContainer{T<:Particle,U,N,V<:Random123.AbstractR123{U}}
     "Particles."
     vals::Vector{T}
     "Unnormalized logarithmic weights."
     logWs::Vector{Float64}
     "Traced RNG to replay the resampling step"
-    rng::TracedRNG
+    rng::TracedRNG{U,N,V}
 end
 
 function ParticleContainer(particles::Vector{<:Particle})
@@ -205,7 +205,7 @@ function update_keys!(pc::ParticleContainer, ref::Union{Particle,Nothing}=nothin
     for i in 1:n
         pi = pc.vals[i]
         k = split(pi.rng.rng.key)
-        seed!(pi.rng, k[1])
+        Random.seed!(pi.rng, k[1])
     end
 end
 
@@ -255,13 +255,13 @@ function resample_propagate!(
             nseeds = isref ? ni - 1 : ni
 
             seeds = split(p.rng.rng.key, nseeds)
-            !isref && seed!(p.rng, seeds[1])
+            !isref && Random.seed!(p.rng, seeds[1])
 
             children[j += 1] = p
             # fork additional children
             for k in 2:ni
                 part = fork(p, isref)
-                seed!(part.rng, seeds[k])
+                Random.seed!(part.rng, seeds[k])
                 children[j += 1] = part
             end
         end

src/rng.jl

@@ -1,104 +1,84 @@
-using Random123
-using Random
-using Distributions
-
-import Base.rand
-import Random.seed!
-import Random123: set_counter!
-
 # Default RNG type for when nothing is specified
-_BASE_RNG = Philox2x
+const _BASE_RNG = Random123.Philox2x
 
 """
-    TracedRNG{R,T}
+    TracedRNG{R,N,T}
 
 Wrapped random number generator from Random123 to keep track of random streams during model evaluation
 """
-mutable struct TracedRNG{T} <:
-               Random.AbstractRNG where {T<:(Random123.AbstractR123{R} where {R})}
+mutable struct TracedRNG{R,N,T<:Random123.AbstractR123{R}} <: Random.AbstractRNG
     "Model step counter"
     count::Int
     "Inner RNG"
     rng::T
     "Array of keys"
-    keys
+    keys::Array{R,N}
 end
 
 """
-    TracedRNG(r::Random123.AbstractR123)
-
-Initialize TracedRNG with r as the inner RNG 
+    TracedRNG(r::Random123.AbstractR123=AdvancedPS._BASE_RNG())
+Create a `TracedRNG` with `r` as the inner RNG. 
 """
-function TracedRNG(r::Random123.AbstractR123)
-    set_counter!(r, 0)
+function TracedRNG(r::Random123.AbstractR123=_BASE_RNG())
+    Random123.set_counter!(r, 0)
     return TracedRNG(1, r, typeof(r.key)[])
 end
 
-"""
-    TracedRNG()
-
-Create a default TracedRNG
-"""
-function TracedRNG()
-    r = _BASE_RNG()
-    return TracedRNG(r)
-end
-
 # 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)
+Random.rng_native_52(rng::TracedRNG) = Random.rng_native_52(rng.rng)
+Base.rand(rng::TracedRNG, ::Type{T}) where {T} = Base.rand(rng.rng, T)
 
 """
-    split(key::Integer, n::Integer)
+    split(key::Integer, n::Integer=1)
 
-Split key into n new keys
+Split `key` into `n` new keys
 """
-function split(key::Integer, n::Integer=1) where {T}
-    return map(i -> hash(key, convert(UInt, i)), 1:n)
+function split(key::Integer, n::Integer=1)
+    return [hash(key, i) for i in UInt(1):UInt(n)]
 end
 
 """
-    load_state!(r::TracedRNG, seed)
+    load_state!(r::TracedRNG)
 
 Load state from current model iteration. Random streams are now replayed
 """
-function load_state!(rng::TracedRNG{T}) where {T}
+function load_state!(rng::TracedRNG)
     key = rng.keys[rng.count]
     Random.seed!(rng.rng, key)
-    return set_counter!(rng.rng, 0)
+    return Random123.set_counter!(rng.rng, 0)
 end
 
 """
     update_rng!(rng::TracedRNG)
 
-Set key and counter of inner RNG to key and the running model step
+Set key and counter of inner rng in `rng` to `key` and the running model step to 0
 """
-function seed!(rng::TracedRNG{T}, key) where {T}
-    seed!(rng.rng, key)
-    return set_counter!(rng.rng, 0)
+function Random.seed!(rng::TracedRNG, key)
+    Random.seed!(rng.rng, key)
+    return Random123.set_counter!(rng.rng, 0)
 end
 
 """ 
     save_state!(r::TracedRNG)
 
-Track current key of the inner RNG
+Add current key of the inner rng in `r` to `keys`.
 """
-function save_state!(r::TracedRNG{T}) where {T}
+function save_state!(r::TracedRNG)
     return push!(r.keys, r.rng.key)
 end
 
-Base.copy(r::TracedRNG{T}) where {T} = TracedRNG(r.count, copy(r.rng), copy(r.keys))
+Base.copy(r::TracedRNG) = TracedRNG(r.count, copy(r.rng), deepcopy(r.keys))
 
 """
-    set_count!(r::TracedRNG, n::Integer)
+    set_counter!(r::TracedRNG, n::Integer)
 
-Set the counter of the TracedRNG, used to keep track of the current model step
+Set the counter of the inner rng in `r`, used to keep track of the current model step
 """
-set_counter!(r::TracedRNG, n::Integer) = r.count = n
+Random123.set_counter!(r::TracedRNG, n::Integer) = r.count = n
 
 """
     inc_counter!(r::TracedRNG, n::Integer=1)
 
-Increase the model step counter by n
+Increase the model step counter by `n`
 """
 inc_counter!(r::TracedRNG, n::Integer=1) = r.count += n

test/rng.jl

@@ -18,7 +18,7 @@
 
         @test key ≠ new_key
 
-        AdvancedPS.seed!(rng, new_key)
+        Random.seed!(rng, new_key)
         @test rng.rng.key === new_key
     end
 end