Improve interface

Author: Tortar

Project.toml

@@ -1,6 +1,6 @@
 name = "StreamSampling"
 uuid = "ff63dad9-3335-55d8-95ec-f8139d39e468"
-version = "0.5.2"
+version = "0.6.0"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

docs/src/api.md

@@ -14,6 +14,7 @@ empty!
 value
 ordvalue
 nobs
+StreamSample
 itsample
 ```
 

docs/src/index.md

@@ -44,4 +44,23 @@ julia> value(rs)
  74
 ```
 
+If the total number of elements in the stream is known beforehand and the sampling is unweighted, it is
+also possible to iterate over a `StreamSample` like so
+
+```julia
+julia> using StreamSampling
+
+julia> iter = 1:100
+
+julia> ss = StreamSample{Int}(iter, 5, 100);
+
+julia> collect(ss)
+5-element Vector{Int64}:
+  7
+  9
+ 20
+ 49
+ 74
+```
+
 Consult the [API page](https://juliadynamics.github.io/StreamSampling.jl/stable/api) for more information about the package interface.

src/SamplingInterface.jl

@@ -1,7 +1,7 @@
 
 """
-    ReservoirSample([rng], T, method = AlgRSWRSKIP())
-    ReservoirSample([rng], T, n::Int, method = AlgL(); ordered = false)
+    ReservoirSample{T}([rng], method = AlgRSWRSKIP())
+    ReservoirSample{T}([rng], n::Int, method = AlgL(); ordered = false)
 
 Initializes a reservoir sample which can then be fitted with [`fit!`](@ref).
 The first signature represents a sample where only a single element is collected.
@@ -10,19 +10,21 @@ they were collected with [`ordvalue`](@ref).
 
 Look at the [`Sampling Algorithms`](@ref) section for the supported methods. 
 """
-function ReservoirSample(T, method::ReservoirAlgorithm = AlgRSWRSKIP())
-    return ReservoirSample(Random.default_rng(), T, method, MutSample())
+struct ReservoirSample{T} 1 === 1 end
+
+function ReservoirSample{T}(method::ReservoirAlgorithm = AlgRSWRSKIP()) where T
+    return ReservoirSample{T}(Random.default_rng(), method, MutSample())
 end
-function ReservoirSample(rng::AbstractRNG, T, method::ReservoirAlgorithm = AlgRSWRSKIP())
-    return ReservoirSample(rng, T, method, MutSample())
+function ReservoirSample{T}(rng::AbstractRNG, method::ReservoirAlgorithm = AlgRSWRSKIP()) where T
+    return ReservoirSample{T}(rng, method, MutSample())
 end
-Base.@constprop :aggressive function ReservoirSample(T, n::Integer, method::ReservoirAlgorithm=AlgL(); 
-        ordered = false)
-    return ReservoirSample(Random.default_rng(), T, n, method, MutSample(), ordered ? Ord() : Unord())
+Base.@constprop :aggressive function ReservoirSample{T}(n::Integer, method::ReservoirAlgorithm=AlgL(); 
+        ordered = false) where T
+    return ReservoirSample{T}(Random.default_rng(), n, method, MutSample(), ordered ? Ord() : Unord())
 end
-Base.@constprop :aggressive function ReservoirSample(rng::AbstractRNG, T, n::Integer, 
-        method::ReservoirAlgorithm=AlgL(); ordered = false)
-    return ReservoirSample(rng, T, n, method, MutSample(), ordered ? Ord() : Unord())
+Base.@constprop :aggressive function ReservoirSample{T}(rng::AbstractRNG, n::Integer, 
+        method::ReservoirAlgorithm=AlgL(); ordered = false) where T
+    return ReservoirSample{T}(rng, n, method, MutSample(), ordered ? Ord() : Unord())
 end
 
 """
@@ -86,7 +88,6 @@ function Base.merge!(::AbstractReservoirSample)
     error("Abstract Version")
 end
 
-
 """
     Base.merge(rs::AbstractReservoirSample...)
 
@@ -98,6 +99,31 @@ function OnlineStatsBase.merge(::AbstractReservoirSample)
     error("Abstract Version")
 end
 
+"""
+    StreamSample{T}([rng], iter, n, [N], method = AlgS())
+
+Initializes a stream sample, 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
+[`StreamSample`](@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)`.
+"""
+struct StreamSample{T} 1 === 1 end
+
+function StreamSample{T}(iter, n, N, method::StreamAlgorithm = AlgORDS()) where T
+    return StreamSample{T}(Random.default_rng(), iter, n, N, method)
+end
+function StreamSample{T}(iter, n, method::StreamAlgorithm = AlgORDS()) where T
+    return StreamSample{T}(Random.default_rng(), iter, n, length(iter), method)
+end
+function StreamSample{T}(rng::AbstractRNG, iter, n, method::StreamAlgorithm = AlgORDS()) where T
+    return StreamSample{T}(rng, iter, n, length(iter), method)
+end
+function StreamSample{T}(rng::AbstractRNG, iter, n, N, method::StreamAlgorithm = AlgORDS()) where T
+    return StreamSample{T}(rng, iter, n, N, method)
+end
+
 """
     itsample([rng], iter, method = AlgRSWRSKIP())
     itsample([rng], iter, wfunc, method = AlgWRSWRSKIP())
@@ -145,37 +171,38 @@ end
 Base.@constprop :aggressive function itsample(rng::AbstractRNG, iter, method = AlgRSWRSKIP();
         iter_type = infer_eltype(iter))
     if Base.IteratorSize(iter) isa Base.SizeUnknown
-        s = ReservoirSample(rng, iter_type, method, ImmutSample())
+        s = ReservoirSample{iter_type}(rng, method, ImmutSample())
         return update_all!(s, iter)
     else 
-        return sortedindices_sample(rng, iter)
+        return sorted_sample_single(rng, iter)
     end
 end
 Base.@constprop :aggressive function itsample(rng::AbstractRNG, iter, n::Int, method = AlgL(); 
         iter_type = infer_eltype(iter), ordered = false)
     if Base.IteratorSize(iter) isa Base.SizeUnknown
-        s = ReservoirSample(rng, iter_type, n, method, ImmutSample(), ordered ? Ord() : Unord())
+        s = ReservoirSample{iter_type}(rng, n, method, ImmutSample(), ordered ? Ord() : Unord())
         return update_all!(s, iter, ordered)
     else
-        replace = method isa AlgL || method isa AlgR ? NoReplace() : Replace()
-        sortedindices_sample(rng, iter, n, replace; iter_type, ordered)
+        m = method isa AlgL || method isa AlgR ? AlgORDS() : AlgORDSWR()
+        s = collect(StreamSample{iter_type}(rng, iter, n, length(iter), m))
+        return ordered ? s : shuffle!(rng, s)
     end
 end
 function itsample(rng::AbstractRNG, iter, wv::Function, method = AlgWRSWRSKIP(); iter_type = infer_eltype(iter))
-    s = ReservoirSample(rng, iter_type, method, ImmutSample())
+    s = ReservoirSample{iter_type}(rng, method, ImmutSample())
     return update_all!(s, iter, wv)
 end
 Base.@constprop :aggressive function itsample(rng::AbstractRNG, iter, wv::Function, n::Int, method = AlgAExpJ(); 
         iter_type = infer_eltype(iter), ordered = false)
-    s = ReservoirSample(rng, iter_type, n, method, ImmutSample(), ordered ? Ord() : Unord())
+    s = ReservoirSample{iter_type}(rng, n, method, ImmutSample(), ordered ? Ord() : Unord())
     return update_all!(s, iter, ordered, wv)
 end
 function itsample(rngs::Tuple, iters::Tuple, n::Int,; iter_types = infer_eltype.(iters))
     n_it = length(iters)
     vs = Vector{Vector{Union{iter_types...}}}(undef, n_it)
     ps = Vector{Float64}(undef, n_it)
     Threads.@threads for i in 1:n_it
-        s = ReservoirSample(rngs[i], iter_types[i], n, AlgRSWRSKIP(), ImmutSample(), Unord())
+        s = ReservoirSample{iter_types[i]}(rngs[i], n, AlgRSWRSKIP(), ImmutSample(), Unord())
         vs[i], ps[i] = update_all_p!(s, iters[i])
     end
     ps /= sum(ps)
@@ -186,7 +213,7 @@ function itsample(rngs::Tuple, iters::Tuple, wfuncs::Tuple, n::Int; iter_types =
     vs = Vector{Vector{Union{iter_types...}}}(undef, n_it)
     ps = Vector{Float64}(undef, n_it)
     Threads.@threads for i in 1:n_it
-        s = ReservoirSample(rngs[i], iter_types[i], n, AlgWRSWRSKIP(), ImmutSample(), Unord())
+        s = ReservoirSample{iter_types[i]}(rngs[i], n, AlgWRSWRSKIP(), ImmutSample(), Unord())
         vs[i], ps[i] = update_all_p!(s, iters[i], wfuncs[i])
     end
     ps /= sum(ps)

src/SamplingUtils.jl

@@ -13,9 +13,6 @@ function infer_eltype(itr)
     ifelse(T2 !== Union{} && T2 <: T1, T2, T1)
 end
 
-get_sorted_indices(rng, n, N, ::Replace) = SortedRandRangeIter(rng, 1:N, n)
-get_sorted_indices(rng, n, N, ::NoReplace) = sort!(sample(rng, 1:N, n; replace=false))
-
 struct SortedRandRangeIter{R}
     rng::R
     range::UnitRange{Int}

src/SortedSamplingMulti.jl

@@ -1,35 +1,43 @@
 
-"""
-    sortedindices_sample(rng, iter)
-    sortedindices_sample(rng, iter, n; replace = false, ordered = false)
-
-Algorithm which generates sorted random indices used to retrieve the sample 
-from the iterable. The number of elements in the iterable needs to be known 
-before starting the sampling.
-"""
-function sortedindices_sample(rng, iter, n::Int, replace; 
-        iter_type = infer_eltype(iter), ordered = false)
-    N = length(iter)
-    if N <= n
-        reservoir = collect(iter)
-        replace isa Replace && return sample(rng, reservoir, n, ordered=ordered)
-        return ordered ? reservoir : shuffle!(rng, reservoir)
+struct SampleMultiAlgORD{T,R,I,D} <: AbstractStreamSample
+    rng::R
+    it::I
+    n::Int
+    inds::D
+    function SampleMultiAlgORD{T}(rng::R, it::I, n, inds::D) where {T,R,I,D}
+        return new{T,R,I,D}(rng, it, n, inds)
     end
-    reservoir = Vector{iter_type}(undef, n)
-    indices = get_sorted_indices(rng, n, N, replace)
-    curr_idx, state_idx = iterate(indices)
+end
+
+function StreamSample{T}(rng::AbstractRNG, iter, n, N, ::AlgORDSWR) where T
+    return SampleMultiAlgORD{T}(rng, iter, n, SortedRandRangeIter(rng, 1:N, n))
+end
+function StreamSample{T}(rng::AbstractRNG, iter, n, N, ::AlgORDS) where T
+    return SampleMultiAlgORD{T}(rng, iter, min(n, N), sort!(sample(rng, 1:N, min(n, N); replace=false)))
+end
+
+@inline function Base.iterate(s::SampleMultiAlgORD)
+    indices, iter = s.inds, s.it
+    curr_idx, state_idx = iterate(indices)::Tuple
     el, state_el = iterate(iter)::Tuple
     for _ in 1:curr_idx-1
         el, state_el = iterate(iter, state_el)::Tuple
     end
-    reservoir[1] = el
-    @inbounds for i in 2:n
-        next_idx, state_idx = iterate(indices, state_idx)
-        for _ in 1:next_idx-curr_idx
-            el, state_el = iterate(iter, state_el)::Tuple
-        end
-        reservoir[i] = el
-        curr_idx = next_idx
+    return (el, (el, state_el, curr_idx, state_idx))
+end
+@inline function Base.iterate(s::SampleMultiAlgORD, state)
+    el, state_el, curr_idx, state_idx = state
+    indices, iter = s.inds, s.it
+    it_indices = iterate(indices, state_idx)
+    it_indices === nothing && return nothing
+    next_idx, state_idx = it_indices
+    for _ in 1:next_idx-curr_idx
+        el, state_el = iterate(iter, state_el)::Tuple
     end
-    return ordered ? reservoir : shuffle!(rng, reservoir)
+    return (el, (el, state_el, next_idx, state_idx))
 end
+
+Base.IteratorEltype(::SampleMultiAlgORD) = Base.HasEltype()
+Base.eltype(::SampleMultiAlgORD{T}) where T = T
+Base.IteratorSize(::SampleMultiAlgORD) = Base.HasLength()
+Base.length(s::SampleMultiAlgORD) = s.n

src/SortedSamplingSingle.jl

@@ -1,5 +1,5 @@
 
-function sortedindices_sample(rng, iter; kwargs...)
+function sorted_sample_single(rng, iter)
     k = rand(rng, 1:length(iter))
     for (i, el) in enumerate(iter)
         i == k && return el

src/StreamSampling.jl

@@ -16,7 +16,7 @@ using Random
 using StatsBase
 
 export fit!, merge!, value, ordvalue, nobs, itsample
-export AbstractReservoirSample, ReservoirSample
+export AbstractReservoirSample, ReservoirSample, StreamSample
 export AlgL, AlgR, AlgRSWRSKIP, AlgARes, AlgAExpJ, AlgWRSWRSKIP
 
 struct ImmutSample end
@@ -25,56 +25,76 @@ struct MutSample end
 struct Ord end
 struct Unord end
 
-struct Replace end
-struct NoReplace end
-
+abstract type AbstractStreamSample end
 abstract type AbstractReservoirSample <: OnlineStat{Any} end
 
-abstract type ReservoirAlgorithm end
+abstract type StreamAlgorithm end
+abstract type ReservoirAlgorithm <: StreamAlgorithm end
+
+"""
+Implements random sampling without replacement. To be used with [`StreamSample`](@ref)
+or [`itsample`](@ref).
+"""
+struct AlgORDS <: StreamAlgorithm end
+
+"""
+Implements random stream sampling with replacement. To be used with [`StreamSample`](@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 sampling without replacement. 
+Implements random reservoir sampling without replacement. To be used with [`ReservoirSample`](@ref)
+or [`itsample`](@ref).
 
 Adapted from algorithm R described in "Random sampling with a reservoir, J. S. Vitter, 1985".
 """
 struct AlgR <: ReservoirAlgorithm end
 
 """
-Implements random sampling without replacement.
+Implements random reservoir sampling without replacement. To be used with [`ReservoirSample`](@ref)
+or [`itsample`](@ref).
 
 Adapted from algorithm L described in "Random sampling with a reservoir, J. S. Vitter, 1985".
 """
 struct AlgL <: ReservoirAlgorithm end
 
 """
-Implements random sampling with replacement.
+Implements random reservoir sampling with replacement. To be used with [`ReservoirSample`](@ref)
+or [`itsample`](@ref).
 
-Adapted fron algorithm RSWR_SKIP described in "Reservoir-based Random Sampling with Replacement from 
+Adapted fron algorithm RSWR-SKIP described in "Reservoir-based Random Sampling with Replacement from 
 Data Stream, B. Park et al., 2008".
 """
 struct AlgRSWRSKIP <: ReservoirAlgorithm end
 
 """
-Implements weighted random sampling without replacement.
+Implements weighted random reservoir sampling without replacement. To be used with [`ReservoirSample`](@ref)
+or [`itsample`](@ref).
 
-Adapted from algorithm A-Res described in "Weighted random sampling with a reservoir, 
-P. S. Efraimidis et al., 2006".
+Adapted from algorithm A-Res described in "Weighted random sampling with a reservoir, P. S. Efraimidis
+et al., 2006".
 """
 struct AlgARes <: ReservoirAlgorithm end
 
 """
-Implements weighted random sampling without replacement.
+Implements weighted random reservoir sampling without replacement. To be used with [`ReservoirSample`](@ref)
+or [`itsample`](@ref).
 
-Adapted from algorithm A-ExpJ described in "Weighted random sampling with a reservoir, 
-P. S. Efraimidis et al., 2006".
+Adapted from algorithm A-ExpJ described in "Weighted random sampling with a reservoir, P. S. Efraimidis
+et al., 2006".
 """
 struct AlgAExpJ <: ReservoirAlgorithm end
 
 """
-Implements weighted random sampling with replacement.
+Implements weighted random reservoir sampling with replacement. To be used with [`ReservoirSample`](@ref)
+or [`itsample`](@ref).
 
-Adapted from algorithm WRSWR_SKIP described in "A Skip-based Algorithm for Weighted Reservoir 
-Sampling with Replacement, A. Meligrana, 2024". 
+Adapted from algorithm WRSWR-SKIP described in "Weighted Reservoir Sampling with Replacement from Multiple
+Data Streams, A. Meligrana, 2024". 
 """
 struct AlgWRSWRSKIP <: ReservoirAlgorithm end