Create benchmark_comparison_stream_WWR.jl (#108)

Author: Tortar

benchmark/benchmark_comparison_non_stream.jl renamed to benchmark/benchmark_comparison_non_stream_WWR.jl

@@ -19,28 +19,28 @@ function weighted_reservoir_sample_seq(rng, a, ws, n)
     w_sum = sum(view_w_f_n)
     reservoir = sample(rng, @view(a[1:m]), Weights(view_w_f_n, w_sum), n)
     length(a) <= n && return reservoir, w_sum
-    w_skip = skip(rng, w_sum, n)
+    w_skip = @inline skip(rng, w_sum, n)
     @inbounds for i in n+1:length(a)
         w_el = ws[i]
         w_sum += w_el
         if w_sum > w_skip
             p = w_el/w_sum
             q = 1-p
-            z = q^(n-4)
+            z = exp((n-4)*log1p(-p))
             t = rand(rng, Uniform(z*q*q*q*q,1.0))
-            k = choose(n, p, q, t, z)
-            for j in 1:k
+            k = @inline choose(n, p, q, t, z)
+            @inbounds for j in 1:k
                 r = rand(rng, j:n)
-                @inbounds reservoir[r], reservoir[j] = reservoir[j], a[i]
+                reservoir[r], reservoir[j] = reservoir[j], a[i]
             end 
-            w_skip = skip(rng, w_sum, n)
+            w_skip = @inline skip(rng, w_sum, n)
         end
     end
     return reservoir, w_sum
 end
 
 function skip(rng, w_sum::AbstractFloat, n)
-    k = rand(rng)^(1/n)
+    k = exp(-randexp(rng)/n)
     return w_sum/k
 end
 

benchmark/benchmark_comparison_stream_WWR.jl

@@ -0,0 +1,78 @@
+
+using StreamSampling
+using StatsBase
+using OnlineStatsBase
+using HybridStructs
+using DataStructures
+using Random
+
+struct AlgAExpJWR end
+
+struct SampleMultiAlgAExpJWR{B, R, T} <: AbstractReservoirSample
+    n::Int
+    seen_k::Int
+    w_sum::Float64
+    rng::R
+    value::B
+    value_prev::Vector{T}
+    weights::Vector{Float64}
+end
+
+function StreamSampling.ReservoirSample{T}(rng::AbstractRNG, n::Integer, ::AlgAExpJWR,
+        ::StreamSampling.ImmutSample, ::StreamSampling.Unord) where T
+    value = BinaryHeap(Base.By(first, DataStructures.FasterForward()), Tuple{Float64,T}[])
+    sizehint!(value, n)
+    v = Vector{T}(undef, n)
+    w = Vector{Float64}(undef, n)
+    return SampleMultiAlgAExpJWR(n, 0, 0.0, rng, value, v, w)
+end
+
+@inline function OnlineStatsBase._fit!(s::SampleMultiAlgAExpJWR, el, w)
+    n = s.n
+    s = @inline update_state!(s, w)
+    if s.seen_k <= n
+        @inbounds s.value_prev[s.seen_k] = el
+        @inbounds s.weights[s.seen_k] = w
+        if s.seen_k === n
+            for x in sample(s.rng, s.value_prev, Weights(s.weights, s.w_sum), n)
+                push!(s.value, (skip_single(s.rng, s.w_sum), x))
+            end
+            empty!(s.value_prev)
+            empty!(s.weights)
+        end
+    else
+        while first(s.value)[1] <= s.w_sum
+            pop!(s.value)
+            push!(s.value, (skip_single(s.rng, s.w_sum), el))
+        end
+    end
+    return s
+end
+
+skip_single(rng, n) = n/rand(rng)
+
+function update_state!(s::SampleMultiAlgAExpJWR, w)
+    @update s.seen_k += 1
+    @update s.w_sum += w
+    return s
+end
+
+function OnlineStatsBase.value(s::SampleMultiAlgAExpJWR)
+    return shuffle!(s.rng, last.(s.value.valtree))
+end
+
+a = Iterators.filter(x -> x != 1, 1:10^8)
+wv_const(x) = 1.0
+wv_incr(x) = Float64(x)
+wv_decr(x) = 1/x
+wvs = (wv_decr, wv_const, wv_incr)
+
+for wv in wvs
+    for m in (AlgWRSWRSKIP(), AlgAExpJWR())
+        for sz in [10^i for i in 0:7]
+            b = @benchmark itsample($a, $wv, $sz, $m) seconds=10
+            println(wv, " ", m, " ", sz, " ", median(b.times))
+        end
+    end
+end
+