Writing a parallel iterator for a random number generator (ish) with Rayon

Hello all,

So I have this struct, which implements Iter and where indexes really do not mean anything.

use rand::prelude::*;

pub struct HorizontalDiskUniformSampler{
    n_samples: usize,
    i: usize,
    rng: ThreadRng,
    radius: f32,
}


impl HorizontalDiskUniformSampler {
    pub fn new(radius: f32, n_samples: usize)->Self{
        Self{
            i: 0,
            rng: rand::thread_rng(),
            radius,
            n_samples,
        }
    }
}

impl Iterator for HorizontalDiskUniformSampler{
    type Item = (f32, f32);
    fn next(&mut self)->Option<Self::Item>{
        if self.n_samples == self.i {
            return None
        }
        self.i += 1;
        
       let (r, theta): (f32, f32) = self.rng.gen();
    
       let r = self.radius * r.sqrt();
       let theta = 2. * std::f32::consts::PI * theta;
       let (theta_sin, theta_cos) = theta.sin_cos();
       let local_x = r * theta_sin;
       let local_y = r * theta_cos;
       Some((local_x , local_y))
    }
}


fn main(){
    
    HorizontalDiskUniformSampler::new(1., 5).for_each(|x|{
        println!("{:?}", x); 
    });
/* 
PRINTS SOMETHING LIKE
(-0.60948163, 0.26467818)
(0.8898232, -0.026498131)
(0.33927688, -0.72550493)
(-0.72247183, 0.22290573)
(-0.32823578, -0.77099776)
*/    
}

I am now trying to implement a parallel version of this with Rayon. I think it should be easy, as there are no serious data races (except for the ThreadRng)... but I have not been able to figure out Rayon's necessary drive_unindexed function. How does that work? I did not find any guide or tutorial.

Or, is there any other way of doing this properly?

Thanks a lot in advance!

So, this works... but I am still interested on a more elegant and memory-efficient method...?

let seq_iter = HorizontalDiskUniformSampler::new(1., 12);

let parallel_iter = {
       let aux : Vec<(f32,f32)>= seq_iter.map(|v|v).collect();
       aux.into_par_iter()
};

Thanks! But I think the code is correct, as r is found using the square root of a uniform random variable ( radius * self.rng.gen().sqrt()). Check THIS for my source.

If you don't need a named iterator type, you can use map_init, something like:

(0..n)
    .into_par_iter()
    .map_init(
        || HorizontalDiskUniformSampler::new(),
        |sampler, _i| sampler.gen(),
    )
4 Likes

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.