add Rand objects

Author: rfourquet

src/RandomExtensions.jl

@@ -1,6 +1,6 @@
 module RandomExtensions
 
-export Combine, Uniform, Normal, Exponential, CloseOpen
+export Combine, Uniform, Normal, Exponential, CloseOpen, Rand
 
 import Random: Sampler, rand, rand!
 
@@ -13,6 +13,7 @@ using SparseArrays: sprand, sprandn
 include("distributions.jl")
 include("sampling.jl")
 include("containers.jl")
+include("iteration.jl")
 
 
 ## updated rand docstring (TODO: replace Base's one)

src/iteration.jl

@@ -0,0 +1,36 @@
+# iterating over on-the-fly generated random values
+
+## Rand
+
+struct Rand{R<:AbstractRNG,S<:Sampler}
+    rng::R
+    sp::S
+end
+
+# X can be an explicit Distribution, or an implicit one like 1:10
+Rand(rng::AbstractRNG, X) = Rand(rng, Sampler(rng, X))
+Rand(rng::AbstractRNG, ::Type{X}=Float64) where {X} = Rand(rng, Sampler(rng, X))
+
+Rand(X) = Rand(GLOBAL_RNG, X)
+Rand(::Type{X}=Float64) where {X} = Rand(GLOBAL_RNG, X)
+
+(R::Rand)(args...) = rand(R.rng, R.sp, args...)
+
+Base.start(R::Rand) = R
+
+function Base.next(::Union{Rand,Distribution}, R::Rand)
+    e = R()
+    e, R
+end
+
+Base.done(::Union{Rand,Distribution}, ::Rand) = false
+
+Base.IteratorSize(::Type{<:Rand}) = Base.IsInfinite()
+
+Base.IteratorEltype(::Type{<:Rand}) = Base.HasEltype()
+Base.eltype(::Type{<:Rand{R, <:Sampler{T}}}) where {R,T} = T
+
+# convenience iteration over distributions
+
+Base.start(d::Distribution) = Rand(GLOBAL_RNG, d)
+Base.IteratorSize(::Type{<:Distribution}) = Base.IsInfinite()

test/runtests.jl

@@ -88,3 +88,17 @@ end
     @test length(s) == 8
     @test Set(s) <= Set("asd")
 end
+
+@testset "Rand" for rng in ([], [MersenneTwister(0)], [RandomDevice()])
+    for XT = zip(([Int], [1:3], []), (Int, Int, Float64))
+        X, T = XT
+        r = Rand(rng..., X...)
+        @test collect(Iterators.take(r, 10)) isa Vector{T}
+        @test r() isa T
+        @test r(2, 3) isa Matrix{T}
+        @test r(.3, 2, 3) isa SparseMatrixCSC{T}
+    end
+    for d = (Uniform(1:10), Uniform(Int))
+        @test collect(Iterators.take(d, 10)) isa Vector{Int}
+    end
+end