JuliaLang / julia

The Julia Programming Language
https://julialang.org/
MIT License
45.61k stars 5.48k forks source link

FR: Random sampler for `Iterators.ProductIterator` #43266

Open fonsp opened 2 years ago

fonsp commented 2 years ago

Feature

I love Iterators.product, and it would be great if I could call rand on it directly:

julia> Iterators.product([1,2],[3,4,5]) |> rand
(1, 5)

Currently, this does not work, because no sampler is defined:

Stacktrace ``` julia> Iterators.product([1,2],[3,4,5]) |> rand ERROR: ArgumentError: Sampler for this object is not defined Stacktrace: [1] Random.Sampler(#unused#::Type{Random.MersenneTwister}, sp::Random.SamplerTrivial{Base.Iterators.ProductIterator{Tuple{Vector{Int64}, Vector{Int64}}}, Tuple{Int64, Int64}}, #unused#::Val{1}) @ Random /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.6/Random/src/Random.jl:145 [2] Random.Sampler(rng::Random.MersenneTwister, x::Random.SamplerTrivial{Base.Iterators.ProductIterator{Tuple{Vector{Int64}, Vector{Int64}}}, Tuple{Int64, Int64}}, r::Val{1}) @ Random /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.6/Random/src/Random.jl:139 [3] rand(rng::Random.MersenneTwister, X::Random.SamplerTrivial{Base.Iterators.ProductIterator{Tuple{Vector{Int64}, Vector{Int64}}}, Tuple{Int64, Int64}}) (repeats 2 times) @ Random /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.6/Random/src/Random.jl:253 [4] rand(X::Base.Iterators.ProductIterator{Tuple{Vector{Int64}, Vector{Int64}}}) @ Random /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.6/Random/src/Random.jl:258 ```

Implementation

A simple implementation would be this:

function Random.rand(
    rng::Random.AbstractRNG,
    iterator::Random.SamplerTrivial{Base.Iterators.ProductIterator{T}},
) where {T}
    r(x) = rand(rng, x)
    r.(iterator[].iterators)
end

Would this PR be welcome? Where in the codebase should the implementation go? Should it be a method of Random.Sampler?

rfourquet commented 2 years ago

This would be a fine implementation, but indeed also a Sampler method would be ideal, which would store a sampler for each component of the product (I'm happy to help there)

fonsp commented 2 years ago

Thanks! Can you provide links to the codebase where I should start looking? Of course, if this would be trivial for you to add, feel free to do so! :)

rfourquet commented 2 years ago

These things are generally implemented in "Random/src/generation.jl", but I don't think there is already a sampler implementation there which stores multiple subsamplers. If you were really motivated, you could look at the RandomExtensions package which has a bunch of examples, where you can do rand(make(Tuple, 1:2, 3:5)) to produce the same as what is in the OP.... Ok you nerd-sniped me ;-) Basically it would look like

Sampler(::Type{RNG}, z::Base.Iterators.ProductIterator, n::Repetition) where {RNG<:AbstractRNG} =
    SamplerTag{typeof(z)}(map(x -> Sampler(RNG, x, n), z.iterators))

rand(r::AbstractRNG, z::SamplerTag{<:Base.Iterators.ProductIterator}) = map(x -> rand(r, x), z.data)