Merge pull request #19 from TuringLang/resampling_rng

Add RNG to resampling methods

Project.toml

@@ -1,11 +1,12 @@
 name = "AdvancedPS"
 uuid = "576499cb-2369-40b2-a588-c64705576edc"
 authors = ["TuringLang"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [compat]

src/AdvancedPS.jl

@@ -2,6 +2,7 @@ module AdvancedPS
 
 import Distributions
 import Libtask
+import Random
 import StatsFuns
 
 include("resampling.jl")

src/container.jl

@@ -166,8 +166,8 @@ function effectiveSampleSize(pc::ParticleContainer)
 end
 
 """
- resample_propagate!(pc::ParticleContainer[, randcat = resample, ref = nothing;
- weights = getweights(pc)])
+ resample_propagate!(rng, pc::ParticleContainer[, randcat = resample_systematic,
+ ref = nothing; weights = getweights(pc)])
 
 Resample and propagate the particles in `pc`.
 
@@ -176,8 +176,9 @@ of the particle `weights`. For Particle Gibbs sampling, one can provide a refere
 `ref` that is ensured to survive the resampling step.
 """
 function resample_propagate!(
+ rng::Random.AbstractRNG,
  pc::ParticleContainer,
- randcat = resample,
+ randcat = resample_systematic,
  ref::Union{Particle, Nothing} = nothing;
  weights = getweights(pc)
 )
@@ -187,7 +188,7 @@ function resample_propagate!(
  # sample ancestor indices
  n = length(pc)
  nresamples = ref === nothing ? n : n - 1
- indx = randcat(weights, nresamples)
+ indx = randcat(rng, weights, nresamples)
 
  # count number of children for each particle
  num_children = zeros(Int, n)
@@ -230,6 +231,7 @@ function resample_propagate!(
 end
 
 function resample_propagate!(
+ rng::Random.AbstractRNG,
  pc::ParticleContainer,
  resampler::ResampleWithESSThreshold,
  ref::Union{Particle,Nothing} = nothing;
@@ -239,7 +241,7 @@ function resample_propagate!(
  ess = inv(sum(abs2, weights))
 
  if ess ≤ resampler.threshold * length(pc)
- resample_propagate!(pc, resampler.resampler, ref; weights = weights)
+ resample_propagate!(rng, pc, resampler.resampler, ref; weights = weights)
  end
 
  pc
@@ -292,7 +294,7 @@ function reweight!(pc::ParticleContainer)
 end
 
 """
- sweep!(pc::ParticleContainer, resampler)
+ sweep!(rng, pc::ParticleContainer, resampler)
 
 Perform a particle sweep and return an unbiased estimate of the log evidence.
 
@@ -303,11 +305,11 @@ The resampling steps use the given `resampler`.
 Del Moral, P., Doucet, A., & Jasra, A. (2006). Sequential monte carlo samplers.
 Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68(3), 411-436.
 """
-function sweep!(pc::ParticleContainer, resampler)
+function sweep!(rng::Random.AbstractRNG, pc::ParticleContainer, resampler)
  # Initial step:
 
  # Resample and propagate particles.
- resample_propagate!(pc, resampler)
+ resample_propagate!(rng, pc, resampler)
 
  # Compute the current normalizing constant ``Z₀`` of the unnormalized logarithmic
  # weights.
@@ -317,7 +319,7 @@ function sweep!(pc::ParticleContainer, resampler)
  logZ0 = logZ(pc)
 
  # Reweight the particles by including the first observation ``y₁``.
- isdone = reweight!(pc)
+ isdone = reweight!(rng, pc)
 
  # Compute the normalizing constant ``Z₁`` after reweighting.
  logZ1 = logZ(pc)
@@ -328,14 +330,14 @@ function sweep!(pc::ParticleContainer, resampler)
  # For observations ``y₂, …, yₜ``:
  while !isdone
  # Resample and propagate particles.
- resample_propagate!(pc, resampler)
+ resample_propagate!(rng, pc, resampler)
 
  # Compute the current normalizing constant ``Z₀`` of the unnormalized logarithmic
  # weights.
  logZ0 = logZ(pc)
 
  # Reweight the particles by including the next observation ``yₜ``.
- isdone = reweight!(pc)
+ isdone = reweight!(rng, pc)
 
  # Compute the normalizing constant ``Z₁`` after reweighting.
  logZ1 = logZ(pc)

src/resampling.jl

@@ -12,33 +12,33 @@ struct ResampleWithESSThreshold{R, T<:Real}
  threshold::T
 end
 
-function ResampleWithESSThreshold(resampler = resample)
+function ResampleWithESSThreshold(resampler = resample_systematic)
  ResampleWithESSThreshold(resampler, 0.5)
 end
 
 # More stable, faster version of rand(Categorical)
-function randcat(p::AbstractVector{<:Real})
+function randcat(rng::Random.AbstractRNG, p::AbstractVector{<:Real})
  T = eltype(p)
- r = rand(T)
+ r = rand(rng, T)
+ cp = p[1]
  s = 1
- for j in eachindex(p)
- r -= p[j]
- if r <= zero(T)
- s = j
- break
- end
+ n = length(p)
+ while cp <= r && s < n
+ @inbounds cp += p[s += 1]
  end
  return s
 end
 
 function resample_multinomial(
+ rng::Random.AbstractRNG,
  w::AbstractVector{<:Real},
  num_particles::Integer = length(w),
 )
- return rand(Distributions.sampler(Distributions.Categorical(w)), num_particles)
+ return rand(rng, Distributions.sampler(Distributions.Categorical(w)), num_particles)
 end
 
 function resample_residual(
+ rng::Random.AbstractRNG,
  w::AbstractVector{<:Real},
  num_particles::Integer = length(weights),
 )
@@ -57,19 +57,19 @@ function resample_residual(
  end
  residuals[j] = x - floor_x
  end
- 
+
  # sampling from residuals
  if i <= num_particles
  residuals ./= sum(residuals)
- rand!(Distributions.Categorical(residuals), view(indices, i:num_particles))
+ rand!(rng, Distributions.Categorical(residuals), view(indices, i:num_particles))
  end
- 
+
  return indices
 end
 
 
 """
- resample_stratified(weights, n)
+ resample_stratified(rng, weights, n)
 
 Return a vector of `n` samples `x₁`, ..., `xₙ` from the numbers 1, ..., `length(weights)`,
 generated by stratified resampling.
@@ -80,7 +80,11 @@ are selected according to the multinomial distribution defined by the normalized
 i.e., `xᵢ = j` if and only if
 ``uᵢ \\in [\\sum_{s=1}^{j-1} weights_{s}, \\sum_{s=1}^{j} weights_{s})``.
 """
-function resample_stratified(weights::AbstractVector{<:Real}, n::Integer = length(weights))
+function resample_stratified(
+ rng::Random.AbstractRNG,
+ weights::AbstractVector{<:Real},
+ n::Integer = length(weights),
+)
  # check input
  m = length(weights)
  m > 0 || error("weight vector is empty")
@@ -93,7 +97,7 @@ function resample_stratified(weights::AbstractVector{<:Real}, n::Integer = lengt
  sample = 1
  @inbounds for i in 1:n
  # sample next `u` (scaled by `n`)
- u = oftype(v, i - 1 + rand())
+ u = oftype(v, i - 1 + rand(rng))
 
  # as long as we have not found the next sample
  while v < u
@@ -114,7 +118,7 @@ function resample_stratified(weights::AbstractVector{<:Real}, n::Integer = lengt
 end
 
 """
- resample_systematic(weights, n)
+ resample_systematic(rng, weights, n)
 
 Return a vector of `n` samples `x₁`, ..., `xₙ` from the numbers 1, ..., `length(weights)`,
 generated by systematic resampling.
@@ -125,14 +129,18 @@ numbers `u₁`, ..., `uₙ` where ``uₖ = (u + k − 1) / n``. Based on these n
 normalized `weights`, i.e., `xᵢ = j` if and only if
 ``uᵢ \\in [\\sum_{s=1}^{j-1} weights_{s}, \\sum_{s=1}^{j} weights_{s})``.
 """
-function resample_systematic(weights::AbstractVector{<:Real}, n::Integer = length(weights))
+function resample_systematic(
+ rng::Random.AbstractRNG,
+ weights::AbstractVector{<:Real},
+ n::Integer = length(weights),
+)
  # check input
  m = length(weights)
  m > 0 || error("weight vector is empty")
 
  # pre-calculations
  @inbounds v = n * weights[1]
- u = oftype(v, rand())
+ u = oftype(v, rand(rng))
 
  # find all samples
  samples = Array{Int}(undef, n)
@@ -158,6 +166,3 @@ function resample_systematic(weights::AbstractVector{<:Real}, n::Integer = lengt
 
  return samples
 end
-
-# Default resampling scheme
-const resample = resample_systematic

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

test/container.jl

@@ -47,7 +47,7 @@
  @test AdvancedPS.logZ(pc) ≈ log(sum(exp, 2 .* logps))
 
  # Resample and propagate particles.
- AdvancedPS.resample_propagate!(pc)
+ AdvancedPS.resample_propagate!(Random.GLOBAL_RNG, pc)
  @test pc.logWs == zeros(3)
  @test AdvancedPS.getweights(pc) == fill(1/3, 3)
  @test all(AdvancedPS.getweight(pc, i) == 1/3 for i in 1:3)

test/resampling.jl

@@ -1,17 +1,16 @@
 @testset "resampling.jl" begin
  D = [0.3, 0.4, 0.3]
  num_samples = Int(1e6)
+ rng = Random.GLOBAL_RNG
 
- resSystematic = AdvancedPS.resample_systematic(D, num_samples )
- resStratified = AdvancedPS.resample_stratified(D, num_samples )
- resMultinomial= AdvancedPS.resample_multinomial(D, num_samples )
- resResidual = AdvancedPS.resample_residual(D, num_samples )
- AdvancedPS.resample(D)
- resSystematic2= AdvancedPS.resample(D, num_samples )
+ resSystematic = AdvancedPS.resample_systematic(rng, D, num_samples)
+ resStratified = AdvancedPS.resample_stratified(rng, D, num_samples)
+ resMultinomial= AdvancedPS.resample_multinomial(rng, D, num_samples)
+ resResidual = AdvancedPS.resample_residual(rng, D, num_samples)
+ AdvancedPS.resample_systematic(rng, D)
 
  @test sum(resSystematic .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples
- @test sum(resSystematic2 .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples
  @test sum(resStratified .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples
  @test sum(resMultinomial .== 2) ≈ (num_samples * 0.4) atol=1e-2*num_samples
  @test sum(resResidual .== 2) ≈ (num_samples * 0.4) atol=1e-2*num_samples
-end
+end

test/runtests.jl

@@ -1,5 +1,6 @@
 using AdvancedPS
 using Libtask
+using Random
 using Test
 
 @testset "AdvancedPS.jl" begin