Sequence sampling: `seq`, `WeightedChoice`

[dhardy]

This part of rand has needed work for a while, and hasn't had much love recently, excepting the excellent work on seq::sample* by @vitiral.

WeightedChoice is a distribution for sampling one item from a slice, with weights. The current API only accepts a slice (reference) and clones objects when sampling.

The seq module has functions to:

• sample a set number of items from an iterator of unknown length, returning some of the iterated items
• sample a set number of items from a slice, returning clones of the items
• sample a set number of items from a slice, returning references into the original slice

Issues to solve:

• Is this sufficiently generic/flexible? Do we cover all expected functionality?
• Is the code succinct?
• Is the documentation sufficient?
• Why are WeightedChoice and seq::* in completely separate modules, the first as a distribution the second as simple functions?
• Is there excess functionality which arguably doesn't belong in rand?

References:

• implementation of seq::*
• seq optimisation: investigate using sample_iter instead of sample_indicies_inplace rust-random/rand#203 [opt]: investigate using sample_iter instead of sample_indicies_inplace
• sample with random ordering rust-random/rand#169 [alg]: sampling with randomised ordering
• investigate generic sample_indicies rust-random/rand#202 [API]: investigate generic sample_indicies
• WeightedChoice is hard to use with borrow checker rust-random/rand#142 [API]: WeightedChoice is hard to use with borrow checker
• Consider adding an owned version of WeightedChoice rust-random/rand#90 [API]: Consider adding an owned version of WeightedChoice
• Add an owned variant of WeightedChoice. rust-random/rand#152 [API PR]: Add an owned variant of WeightedChoice
• Merging Rust Random Choice rust-random/rand#110 [API+alg]: RandomChoice impl
• Consider pruning gen_weighted_bool rust-random/rand#47 [API]: Consider pruning gen_weighted_bool
• Add Rng::pick method rust-random/rand#198 [API+alg PR]: Add Rng::pick / Choose
• Missing runtime in WeightedChoice rust-random/rand#116 [doc]: no doc for algorithm complexity
• Partial refactoring I did in this branch: API doc (makes WeightedChoice owning, adds Choose, some reorganisation)

I'm going to say nothing is fixed at this point. In theory it would be nice not to break the new seq code.

[pitdicker]

I have not looked all that closely at what the sample functions are supposed to do to be honest. Does this summarize it?

The new seq::sample* functions and Robert Floyds combination algorithm rust-random#144 are supposed to pick values from a range without returning any duplicates. And the problem is that all the previous results have to be kept in memory and checked against collisions.

Is there really not a better solution? It should be possible to adapt RNG's like PCG (but anything LCG or MCG-based) to work on a certain number of bits. What if we use an RNG with just enough bits to be able to generate all possible values, and reject the ones outside the range? Then all values are guaranteed to only occur once, and no vector with many comparisons or a more complex hashmap are necessary. And no dependency on alloc or std.

This has the disadvantage of not working with a user-supplied RNG. And maybe I am missing something big? cc @vitiral, @burdges

[pitdicker]

maybe I am missing something big?

😄 Realized it just after commenting. If we have to pick values from, say, 0..800 we need 10 bits. That would also mean the seed of the RNG is 10 bits. Then the RNG can never generate more than 2^10 different orderings, 1024. While there are 800! possible orderings, which is basically infinite in my book. So that is an argument to not use this method, but one that can utilize an RNG with a larger period. With an RNG with a period of 2^64 or 2^128 it is at least impossible to exhaustively list all possible orderings.

[dhardy]

That's some out-of-the-box thinking! I guess your custom sampling RNG can be seeded from the provided RNG, so your idea could work, if there is a way to reliably generate a high-quality random-hop RNG visiting all nodes of the given range exactly once in random-like order. That sounds hard to me!

Mostly though I made this issue in the hope that someone can design a nice API tying all the above things together.

[pitdicker]

Mostly though I made this issue in the hope that someone can design a nice API tying all the above things together.

But that is much more difficult! 😄.

I see WeightedChoice as a somewhat odd functionality that not many will use. Do you think there is a place in rand for functions like this? That implement some smart algorithm to solve a not entirely uncommon problem?

Just today I found out about this nice problem of picking a random point on a sphere: http://mathworld.wolfram.com/SpherePointPicking.html. This is the same problem as picking a direction in 3D, which comes up in games.

Would it makes sense to add a module collecting stuff like this?

They may use the Distribution trait if that is convenient, but I don't see that as a reason to live in the distributions module.

[dhardy]

I see WeightedChoice as a somewhat odd functionality that not many will use.

That was my thought too; removing it is definitely an option.

[burdges]

An PRNG can only generate unique values in a range with high probability if the range is much larger than the number of values requested.

I noticed sample_slice incurs a potentially undesirable allocation. At least sample_slice_ref should avoid allocation by converting all the indices to references directly via something like:

fn index_to_ref<'a,T>(slice: &'a [T], index: &'a mut [usize]) -> &'a [T] {
    let s = slice.as_ptr();
    let l = slice.len();
    for i in index.iter_mut() {
        assert!(*i < l);
        *i = unsafe { s.offset(*i) } as usize;
    }
    unsafe { mem::transmute(index) };
}
// I donno if this `mem::transmute` call could/should be replaced with some more idiomatic raw pointer call.
// There is no safe linear time way to implement the mutable variant
//   fn index_to_mut<'a,T>(slice: &'a mut [T], index: &'a mut [usize]) -> &'a mut [T] 
// but unsafely index_to_ref could be coerced into doing the mutable version for us.

Isn't there a use case for returning the second half of a Fisher-Yates shuffle applied to permute the first half? I think it gives you a faster combination in some situations: rust-random#169 (comment)

In principle, we could provide no_std variants that took in this user supplied buffer and call it for the -> Vex<T> versions. At least Robert Floyds combination algorithm works:

fn combination_slice(rng: &mut R, range: std::ops::Range<usize>, s: &mut [usize]) {
    let mut s = ::slicevec::SliceVec::new(s);  // https://github.com/jonas-schievink/slicevec
    let amount = s.capacity()
    let n = range.end;
    let m = std::cmp::min(amount,n-1);
    // Should use #![feature(inclusive_range_syntax)] here in future.
    for j in n-m+1 .. n+1 {
        let t = rng.gen_range(range.start,j);
        let t = if s.contains(&t) { j-1 } else { t };
        s.push( t );
    } 
}

fn combination_vec(rng: &mut R, range: std::ops::Range<usize>, amount: usize) -> Vec<usize> {
    let mut s = Vec::with_capacity((Zero::zero()..amount).len());
    unsafe { s.set_len(amount) };
    combination_slice(rng,range,s.as_mut_slice());
    s
}

fn combination_array<const AMOUNT: usize>(rng: &mut R, range: std::ops::Range<usize>) -> [usize; AMOUNT] {
    let mut s = [0usize; AMOUNT];
    combination_slice(rng,range,&mut s);
    s
}

// Deprecated form: 
fn combination_array<A: Array>(rng: &mut R, range: std::ops::Range<usize>) -> A {
    let mut s = ArrayVec::<A>::new();
    let amount = s.capacity()
    combination_slice(rng,range,s.as_mut_slice());
    s.into_inner()
}

It becomes tricky with T: Drop types of course, but SliceVec seemingly handles them, albeit with unnecessary drops. I suppose allocators might become the preferred mechanism here.

I suppose sample_indices_cache is kinda meant to integrate Robert Floyds combination algorithm's benefits while avoiding the amount^2 / 2 comparisons by using a HashMap? It's "running backwards" so that you can output form the main loop though? Assuming that's done correctly then sample_indices_cache could avoid allocating entirely by becoming an Iterator itself, no?

pub struct SampleIndicesCache<'a> {
    rng: &'a mut R,
    length: usize,  // or Range<usize>
    cache: HashMap<usize>,
    count: usize,
}

fn sample_indices_cache<'a>(rng: &mut R, length: usize) -> SampleIndicesCache<'a> {
    let cache = HashMap::with_capacity(amount);
    SampleIndicesCache { rng, length, cache, count: 0 }
}

impl<'a> Iterator for SampleIndicesCache<'a> {
    type Item = usize;
    fn next(&mut self) {
        let i = self.count;
        if i > length { return None; }
        self.count += 1;
        let j: usize = rng.gen_range(i, self.length);

        let tmp = self.cache.remove(&i).unwrap_or(i);       // equiv: let tmp = slice[i];
        if i == j { return Some(tmp); }  // Avoid reinserting or botching return

        // let x = self.cache.remove(&j).unwrap_or(j);        // equiv: slice[i] = slice[j];
        // self.cache.insert(j, tmp);                                     // equiv: slice[j] = tmp;
        // It's faster than the previous two lines if we use the Enter API
        Some( mem::replace(self.cache.entry(j).or_insert(j), tmp) )
        // note that in the inplace version, slice[i] is automatically "returned" value
    }
}

After writing this, there was seemingly a mistake in sample_indices_cache in which j comes up twice and i==j the second time, but an early return fixes it.

I'm still vaguely uncomfortable doing all this only for usize too. Anyone want these for u64 on a 32 bit platform?

[vitiral]

@burdges I'm not really sure about the rest of your issue, but generic sample_indices is covered in rust-random#202

[burdges]

We might as well provide the ref and mut iterators as well:

pub struct SampleRefsIter<'a,T: 'a> {
    slice: &'a [T],
    indices: SampleIndicesCache<'a>,
}

fn sample_refs_iter<'a,T>(rng: &'a mut R, slice: &'a [T]) -> SampleRefsIter<'a,T> {
    SampleRefsIter { slice, indices: sample_indices_cache(rng,slice.len()) }
}

impl<'a> Iterator for SampleRefsIter<'a,T> {
    type item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.indices.next().map(|x| {
            assert!(x < slice.len());
            unsafe { self.slice.as_ptr().offset(x) as &T } 
        })
    }
}

pub struct SampleMutsIter<'a,T: 'a>(SampleRefsIter<'a,T>);

fn sample_muts_iter<'a,T>(rng: &'a mut R, slice: &'a mut [T]) -> SampleMutsIter<'a,T> {
    SampleMutsIter(sample_refs_iter(rng,slice))
}

impl<'a> Iterator for SampleMutsIter<'a,T> {
    type item = &'a mut T;
    fn next(&mut self) -> Option<&mut T> {
        // In principle, this should be safe because the sample cannot return the same value
        // twice, but `mem::transmute` should be replaced with more idiomatic raw pointer calls
        // and worse this makes memory safety dependent upon the correctness of `HashMap`.
        // I proposed this approach to remain compatible with other mutable iterators, ala
        // https://doc.rust-lang.org/src/core/slice/mod.rs.html#1422-1426
        // but a bound like where Self: 'a might insulate us from `HashMap` but at the cost of 
        // some usability, specifically references cannot be kept between iterations 
        self.0.next().map(|x| unsafe { mem::transmute(x) }).
    }
}

[dhardy]

We have roughly this functionality:

• choose 1 elt from a slice/iter
• sample n elts from a slice/iter, by value/index/ref
• sample values from a weighted distribution

Some of this functionality is currently in Rng (see rust-random#293), some in other places.

Rough proposal: add a trait something like the following, with impls for both iterators and slices (optimised to each appropriately):

pub trait SampleSeq<T> {
    type Item;
    // as the current seq::sample:
    fn sample<R>(self, n: usize, rng: &mut R) -> Result<Vec<T>, Vec<T>>;
    // as above, but with key extraction per element:
    fn sample_weighted<R>(self, weight_fn: Fn<Self::Item -> f64>, n: usize, rng: &mut R) -> Result<Vec<T>, Vec<T>>;
    // sample but with n==1:
    fn choose<R>(self, rng: &mut R);
}

pub fn sample_indices<R>(rng: &mut R, length: usize, n: usize) -> Vec<usize>;