add algorithm for weighted stream sampling

Author: Tortar

src/SamplingInterface.jl

@@ -128,9 +128,26 @@ is assumed to have a `eltype` of `T`. The methods implemented in
 [`StreamSampler`](@ref) require the knowledge of the total number
 of elements in the stream `N`, if not provided it is assumed to be
 available by calling `length(iter)`.
+
+-----
+
+    StreamSampler{T}([rng], iter, wfunc, n, W, method = AlgORDWSWR())
+
+Initializes a weigthed stream sampler, which can then be iterated over
+to return the sampling elements of the iterable `iter` which
+is assumed to have a `eltype` of `T`. The methods implemented in
+[`StreamSampler`](@ref) for weighted streams require the knowledge
+of the total weight of the stream `W` and a weight function `wfunc`
+specifying how to map an element to its weight. 
 """
 struct StreamSampler{T} 1 === 1 end
 
+function StreamSampler{T}(iter, wfunc::Function, n, W, method::StreamAlgorithm = AlgORDWSWR()) where T
+    return StreamSampler{T}(Random.default_rng(), iter, wfunc, n, W, method)
+end
+function StreamSampler{T}(rng::AbstractRNG, iter, wfunc::Function, n, W, method::StreamAlgorithm = AlgORDWSWR()) where T
+    return StreamSampler{T}(rng, iter, wfunc, n, W, method)
+end
 function StreamSampler{T}(iter, n, N, method::StreamAlgorithm = AlgD()) where T
     return StreamSampler{T}(Random.default_rng(), iter, n, N, method)
 end
@@ -159,10 +176,11 @@ If the iterator is empty, it returns `nothing`.
     itsample([rng], iter, wfunc, n::Int, method = AlgAExpJ(); ordered = false)
 
 Return a vector of `n` random elements of the iterator, 
-optionally specifying a `rng` (which defaults to `Random.default_rng()`)
-a weight function `wfunc` and a `method`. `ordered` dictates whether an 
-ordered sample (also called a sequential sample, i.e. a sample where items 
-appear in the same order as in `iter`) must be collected.
+optionally specifying a `rng` (which defaults to `Random.default_rng()`),
+a weight function `wfunc` specifying how to map an element to its weight
+and a `method`. `ordered` dictates whether an ordered sample (also called a 
+sequential sample, i.e. a sample where items  appear in the same order as in
+`iter`) must be collected.
 
 If the iterator has less than `n` elements, in the case of sampling without
 replacement, it returns a vector of those elements.

src/SortedSamplingMulti.jl

@@ -1,4 +1,50 @@
 
+struct MultiAlgWeightedORDSampler{T,R,I,F}
+    rng::R
+    it::I
+    f::F
+    n::Int
+    W::Float64
+    function MultiAlgWeightedORDSampler{T}(rng::R, it::I, f::F, n, W) where {T,R,I,F}
+        return new{T,R,I,F}(rng, it, f, n, W)
+    end
+end
+
+@inline function Base.iterate(s::MultiAlgWeightedORDSampler)
+    local el, state_el
+    w, curx, k = 0.0, 0.0, 0
+    for i in s.n:-1:1
+        curx += (1-exp(-randexp(s.rng)/i))*(1-curx)
+        while w < curx * s.W
+            nstate = k == 0 ? iterate(s.it) : iterate(s.it, state_el)
+            nstate == nothing && return nothing
+            k += 1
+            el, state_el = nstate
+            w += s.f(el)
+        end
+        return (el, (el, w, state_el, curx, i-1))
+    end
+end
+@inline function Base.iterate(s::MultiAlgWeightedORDSampler, state)
+    state[end] == 0 && return nothing
+    el, w, state_el, curx, n = state
+    for i in n:-1:1
+        curx += (1-exp(-randexp(s.rng)/i))*(1-curx)
+        while w < curx * s.W
+            nstate = iterate(s.it, state_el)
+            nstate == nothing && return nothing
+            el, state_el = nstate
+            w += s.f(el)
+        end
+        return (el, (el, w, state_el, curx, i-1))
+    end
+end
+
+Base.IteratorEltype(::MultiAlgWeightedORDSampler) = Base.HasEltype()
+Base.eltype(::MultiAlgWeightedORDSampler{T}) where T = T
+Base.IteratorSize(::MultiAlgWeightedORDSampler) = Base.HasLength()
+Base.length(s::MultiAlgWeightedORDSampler) = s.n
+
 struct MultiAlgORDSampler{T,R,I,D} <: AbstractStreamSampler
     rng::R
     it::I
@@ -9,6 +55,9 @@ struct MultiAlgORDSampler{T,R,I,D} <: AbstractStreamSampler
     end
 end
 
+function StreamSampler{T}(rng::AbstractRNG, iter, wfunc::Function, n, W, ::AlgORDWSWR) where T
+    return MultiAlgWeightedORDSampler{T}(rng, iter, wfunc, n, W)
+end
 function StreamSampler{T}(rng::AbstractRNG, iter, n, N, ::AlgD) where T
     return MultiAlgORDSampler{T}(rng, iter, min(n, N), SeqSampleIter(rng, N, min(n, N)))
 end

src/StreamSampling.jl

@@ -19,7 +19,7 @@ using StatsBase
 
 export fit!, merge!, value, ordvalue, nobs, itsample
 export AbstractReservoirSampler, ReservoirSampler, StreamSampler
-export AlgL, AlgR, AlgRSWRSKIP, AlgARes, AlgAExpJ, AlgWRSWRSKIP, AlgD, AlgORDSWR
+export AlgL, AlgR, AlgRSWRSKIP, AlgARes, AlgAExpJ, AlgWRSWRSKIP, AlgD, AlgORDSWR, AlgORDWSWR
 
 struct ImmutSampler end
 struct MutSampler end
@@ -34,24 +34,6 @@ abstract type AbstractWeightedReservoirSampler <: AbstractReservoirSampler end
 abstract type StreamAlgorithm end
 abstract type ReservoirAlgorithm <: StreamAlgorithm end
 
-"""
-Implements random stream sampling without replacement. To be used with [`StreamSampler`](@ref)
-or [`itsample`](@ref).
-
-Adapted from algorithm D described in "An Efficient Algorithm for Sequential Random Sampling,
-J. S. Vitter, 1987".
-"""
-struct AlgD <: StreamAlgorithm end
-
-"""
-Implements random stream sampling with replacement. To be used with [`StreamSampler`](@ref)
-or [`itsample`](@ref).
-
-Adapted from algorithm 4 described in "Generating Sorted Lists of Random Numbers, J. L. Bentley
-et al., 1980".
-"""
-struct AlgORDSWR <: StreamAlgorithm end
-
 """
 Implements random reservoir sampling without replacement. To be used with [`ReservoirSampler`](@ref)
 or [`itsample`](@ref).
@@ -104,6 +86,32 @@ Replacement, A. Meligrana, 2024".
 """
 struct AlgWRSWRSKIP <: ReservoirAlgorithm end
 
+"""
+Implements random stream sampling without replacement. To be used with [`StreamSampler`](@ref)
+or [`itsample`](@ref).
+
+Adapted from algorithm D described in "An Efficient Algorithm for Sequential Random Sampling,
+J. S. Vitter, 1987".
+"""
+struct AlgD <: StreamAlgorithm end
+
+"""
+Implements random stream sampling with replacement. To be used with [`StreamSampler`](@ref)
+or [`itsample`](@ref).
+
+Adapted from algorithm 4 described in "Generating Sorted Lists of Random Numbers, J. L. Bentley
+et al., 1980".
+"""
+struct AlgORDSWR <: StreamAlgorithm end
+
+"""
+Implements weighted random stream sampling with replacement. To be used with [`StreamSampler`](@ref).
+
+Adapted from algorithm 3 described in "An asymptotically optimal, online algorithm for weighted random
+sampling with replacement, M. Startek, 2016".
+"""
+struct AlgORDWSWR <: StreamAlgorithm end
+
 include("SamplingUtils.jl")
 include("SamplingInterface.jl")
 include("SortedSamplingSingle.jl")