csrk 1.1.4

//! PCG64 Random number generator

// Standard library modules
use rand::rngs::OsRng;
use rand::TryRngCore;
// Third Party
use rand_pcg::Pcg64;

/// A struct containing random number generators belonging to other structs
pub struct PCG64Stream {
    /// 128-bit seed used to initialize the generator
    seed: u128,
    /// Stream/Entropy selector (PCG sequence parameter)
    entropy: u128,
    /// The actual RNG (advanced as samples are drawn)
    rng: Pcg64,
    /// Draws consumed by RNG
    draws_consumed: u128,
    /// Hold on to spare standard normal value
    has_spare_std_norm: bool,
    /// Extra standard normal value since they are produced in pairs
    spare_std_norm: f64,
}
impl PCG64Stream {
    /// Constructor for PCG64Stream struct
    pub fn new(
        seed: u128,
        entropy: u128,
    ) -> Self {
    Self::build(1, Some(seed), Some(entropy)).remove(0)
    }
    /// Generate a single random number stream with an arbitrary
    /// seed and entropy determined by the OS
    #[allow(non_snake_case)]
    pub fn from_OsRng() -> Self {
    Self::build(1, None, None).remove(0)
    }
    /// Build one or many PCG64Stream structs
    pub fn build(
        nrealizations: usize,
        seed: Option<u128>,
        entropy: Option<u128>,
    ) -> Vec<Self> {
        // -- generate a base seed  from the OS if none provided --
        let base_seed = match seed {
            Some(s) => s,
            None => {
                let mut buf = [0u8; 16];
                let mut rng = OsRng;
                rng.try_fill_bytes(&mut buf).expect("OsRng failure");
                u128::from_le_bytes(buf)
            }
        };

        // Base entropy can be zero if not given
        let base_entropy = entropy.unwrap_or(0);

        // Initialize vector of PCG64Stream structs
        let mut out = Vec::with_capacity(nrealizations);

        // Loop realizations
        for i in 0..nrealizations {
            // Define entropy for this iteration
            let ent = base_entropy.wrapping_add(i as u128);
            // Get new RNG
            let rng = Pcg64::new(base_seed, ent);
            // push the realization
            out.push(Self{
                seed: base_seed,
                entropy: ent,
                draws_consumed: 0,
                rng,
                has_spare_std_norm: false,
                spare_std_norm: 0.0,
            });
        }
        out
    }
    /// Access the seed used to initialize this realization
    pub fn seed(&self) -> u128 {
        self.seed
    }
    /// Access the entropy/stream selector
    pub fn entropy(&self) -> u128 {
        self.entropy
    }
    /// Access the draws consumed
    pub fn draws_consumed(&self) -> u128 {
        self.draws_consumed
    }
    /// Reset this realization back to its initial RNG state
    pub fn reset(&mut self) {
        self.rng = Pcg64::new(self.seed(), self.entropy());
        self.draws_consumed = 0;
        self.has_spare_std_norm = false;
    }
    /// Random u64
    pub fn next_u64(&mut self) -> u64 {
        self.draws_consumed += 1;
        self.rng.try_next_u64().unwrap()
    }
    /// Fill u64
    pub fn fill_u64(&mut self, out: &mut [u64]) {
        // Count the size of out
        for point in out.iter_mut() {
            *point = self.next_u64();
        }
    }
    /// Create vector of u64s
    pub fn vec_u64(&mut self, n: usize) -> Vec<u64> {
        (0..n).map(|_| self.next_u64()).collect()
    }
    /// Random double in [0,1)
    pub fn next_f64(&mut self) -> f64 {
        // 53 random bits -> (0,1)
        let x = self.next_u64() >> 11;
        // divide by 2^53
        (x as f64) * (1.0 / 9007199254740992.0)
    }
    /// Fill f64
    pub fn fill_f64(&mut self, out: &mut [f64]) {
        // Count the size of out
        for point in out.iter_mut() {
            *point = self.next_f64();
        }
    }
    /// Create vector of f64s
    pub fn vec_f64(&mut self, n: usize) -> Vec<f64> {
        (0..n).map(|_| self.next_f64()).collect()
    }
    /// Random double in (0,1]
    pub fn next_f64_nonzero(&mut self) -> f64 {
        // If you really need it to not be zero
        1. - self.next_f64()
    }
    /// Fill f64
    pub fn fill_f64_nonzero(&mut self, out: &mut [f64]) {
        // Count the size of out
        for point in out.iter_mut() {
            *point = self.next_f64_nonzero();
        }
    }
    /// Create vector of f64s
    pub fn vec_f64_nonzero(&mut self, n: usize) -> Vec<f64> {
        (0..n).map(|_| self.next_f64_nonzero()).collect()
    }
    /// Pair of Random Standard Normal
    pub fn next_standard_normal_pair(&mut self) -> (f64, f64) {
        // Generate a pair of random numbers
        let u1 = self.next_f64_nonzero();
        let u2 = self.next_f64_nonzero();
        // Do some Box-Muller math
        let r = (-2.0 * u1.ln()).sqrt();
        let theta = 2.0 * std::f64::consts::PI * u2;
        // Calculate random normals
        let z0 = r * theta.cos();
        let z1 = r * theta.sin();
        (z0, z1)
    }
    /// Single Standard Normal
    pub fn next_standard_normal(&mut self) -> f64 {
        // Check if we have an extra floating around
        if self.has_spare_std_norm {
            self.has_spare_std_norm = false;
            return self.spare_std_norm;
        }
        // Generate a pair
        let (x0, x1) = self.next_standard_normal_pair();
        // Update spare
        self.spare_std_norm = x1;
        self.has_spare_std_norm = true;
        // Return x0
        x0
    }
    /// Fill standard normal
    pub fn fill_standard_normal(&mut self, out: &mut [f64]) {
        // Count the size of out
        let mut count = out.len();
        while count > 0 {
            // Assign one value
            if count == 1 {
                out[count - 1] = self.next_standard_normal();
                count -= 1;
            } else {
                (out[count - 1], out[count - 2]) = 
                    self.next_standard_normal_pair();
                count -= 2;
            }
        }
    }
    /// Create vector of standard normals
    pub fn vec_standard_normal(&mut self, n: usize) -> Vec<f64> {
        (0..n).map(|_| self.next_standard_normal()).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Instant;

    #[test]
    fn same_seed_entropy_reproducible_and_resettable() {
        let seed = 123456789u128;
        let entropy = 42u128;

        let mut r1 = PCG64Stream::new(seed,entropy);
        let mut r2 = PCG64Stream::new(seed,entropy);

        let a = r1.vec_u64(5);
        let b = r2.vec_u64(5);

        // Same seed + entropy => identical streams
        assert_eq!(a, b, "streams with same seed+entropy differ");

        // Advance r1, then reset and check again
        let _ = r1.vec_u64(10);
        r1.reset();
        let c = r1.vec_u64(5);

        assert_eq!(a, c, "reset did not restore original stream");
    }

    #[test]
    fn no_seed_produces_different_streams_but_self_consistent() {
        let mut rs = PCG64Stream::build(
            2,
            None,
            None,
        );

        let mut r1 = rs.remove(0);
        let mut r2 = rs.remove(0);

        let a = r1.vec_u64(5);
        let b = r2.vec_u64(5);

        // With OS entropy, these should almost surely differ
        assert_ne!(a, b, "independently seeded streams matched (extremely unlikely)");
        for x in &a {
            assert!(!b.contains(x), "overlap between streems 0 and 1");
        }

        // But each realization should be reproducible under reset
        r1.reset();
        r2.reset();

        let a2 = r1.vec_u64(5);
        let b2 = r2.vec_u64(5);

        assert_eq!(a, a2, "r1 not self-consistent after reset");
        assert_eq!(b, b2, "r2 not self-consistent after reset");
    }

    #[test]
    fn incremented_entropy_produces_independent_streams() {
        let seed = 9999999u128;

        let mut rs = PCG64Stream::build(
            3,
            Some(seed),
            Some(0),
        );

        let a = rs[0].vec_u64(20);
        let b = rs[1].vec_u64(20);
        let c = rs[2].vec_u64(20);

        // Extremely strong check: no overlaps
        for x in &a {
            assert!(!b.contains(x), "overlap between stream 0 and 1");
            assert!(!c.contains(x), "overlap between stream 0 and 2");
        }
        for x in &b {
            assert!(!c.contains(x), "overlap between stream 1 and 2");
        }

        // And definitely not identical
        assert_ne!(a, b);
        assert_ne!(a, c);
        assert_ne!(b, c);
    }
    
    #[test]
    fn next_f64_is_uniform() {
        // Get a random number generator
        let mut rs = PCG64Stream::from_OsRng();
        // Generate 10,000 random doubles
        let t0 = Instant::now();
        let uniform: Vec<f64> = (0..10000)
            .map(|_| rs.next_f64())
            .collect();
        let t_10k = t0.elapsed();
        println!("Generated 10,000 doubles in {:?}!",t_10k);
        // Check that none are equal to 1.0
        assert!(!uniform.contains(&1.0));
        // Find min, max, and mean
        let mut min = 1.0;
        let mut max = -1.0;
        let mut mean = 0.;
        for f in uniform.iter() {
            mean += *f;
            if *f > max {max = *f};
            if *f < min {min = *f};
        }
        mean /= uniform.len() as f64;
        // Check that values are reasonable
        assert!(min < 0.1);
        assert!(max > 0.9);
        assert!(mean > 0.4);
        assert!(mean < 0.6);
    }
    #[test]
    fn next_f64_nonzero_is_uniform() {
        // Get a random number generator
        let mut rs = PCG64Stream::from_OsRng();
        // Generate 10,000 random doubles
        let uniform: Vec<f64> = (0..10000)
            .map(|_| rs.next_f64())
            .collect();
        // Reset and generate 10,000 nonzero doubles
        rs.reset();
        let uniform_nz: Vec<f64> = (0..10000)
            .map(|_| rs.next_f64_nonzero())
            .collect();
        // Check exactness
        for i in 0..uniform.len() {
            assert!(uniform_nz[i] == 1. - uniform[i]);
        }
    }
    #[test]
    fn gaussian_reproducible_and_resettable() {
        // Get a pair of random number generators
        let seed = 123654;
        let entropy = 420;
        let mut r1 = PCG64Stream::new(seed, entropy);
        let mut r2 = PCG64Stream::new(seed, entropy);
        // Initialize two vectors of length 10
        let mut x1 = [0.;10];
        let mut x2 = [0.;10];
        // Fill them with standard normals
        r1.fill_standard_normal(&mut x1);
        r2.fill_standard_normal(&mut x2);
        // Check that they are equal
        assert_eq!(x1,x2);
        // Reset an rng
        r1.reset();
        let mut x3 = [0.;10];
        r1.fill_standard_normal(&mut x3);
        // Check that reset is equal
        assert_eq!(x1, x3);
    }
}