odem_rs_util/

random.rs

//! Contains facilities to ease the creation, distribution and use of random
//! number generators.
//!
//! The module offers [RngStream] as a utility structure to deal with
//! pseudo-randomness. It provides a two-level hierarchy for generating
//! random numbers, ensuring reproducibility in both sequential and parallel
//! simulations.
//!
//! The two levels are:
//! 1.  *Inter-Simulation Streams*: This level is for creating independent
//!     random number "universes" for separate, top-level simulation runs (e.g.,
//!     in a parameter sweep). The main entry points are [`RngStream::iter()`]
//!     for sequential runs and [`RngStream::par_iter()`] for parallel runs.
//!     The latter function requires `feature = "rayon"`.
//!
//! 2.  *Intra-Simulation RNGs*: This level is for drawing individual,
//!     independent PRNGs from a single `RngStream`. These are used for
//!     different tasks *within* one simulation (e.g., one RNG for arrival
//!     times, another for service times). This is handled by methods on an
//!     `RngStream` instance, like [`rng()`][RngStream::rng] or by iterating
//!     over it.

use core::cell::RefCell;

#[cfg(feature = "rand_pcg")]
use rand_pcg::*;

#[cfg(feature = "rand_xoshiro")]
use rand_xoshiro::*;

#[cfg(feature = "rayon")]
#[cfg_attr(docsrs, doc(cfg(feature = "rayon")))]
mod rayon_support;

#[cfg(feature = "rayon")]
pub use rayon_support::RngStreamParIter;

/// Type-alias for the default random-number generator that is used in absense
/// of another user-selected one.
///
/// This defaults to the 64-Bit PCG for `feature="rand_pcg"` with 256-Bit
/// state-space. If instead of `rand_pcg` feature `rand_xoshiro` is active,
/// Xoshiro256++ is used as a default PRNG.
#[cfg(all(feature = "rand_pcg", target_pointer_width = "64"))]
#[cfg_attr(docsrs, doc(cfg(feature = "rand_pcg")))]
pub type DefaultRng = Pcg64Dxsm;

/// Type-alias for the default random-number generator that is used in absense
/// of another user-selected one.
///
/// This defaults to the 64-Bit PCG for `feature="rand_pcg"` with 256-Bit
/// state-space. If instead of `rand_pcg` feature `rand_xoshiro` is active,
/// Xoshiro256++ is used as a default PRNG.
#[cfg(all(feature = "rand_pcg", target_pointer_width = "32"))]
pub type DefaultRng = Pcg32;

/// Type-alias for the default random-number generator that is used in absense
/// of another user-selected one.
///
/// This defaults to the 64-Bit PCG for `feature="rand_pcg"` with 256-Bit
/// state-space. If instead of `rand_pcg` feature `rand_xoshiro` is active,
/// Xoshiro256++ is used as a default PRNG.
#[cfg(all(
	not(feature = "rand_pcg"),
	feature = "rand_xoshiro",
	target_pointer_width = "64"
))]
#[cfg_attr(
	docsrs,
	doc(cfg(all(not(feature = "rand_pcg"), feature = "rand_xoshiro")))
)]
pub type DefaultRng = Xoshiro256PlusPlus;

/// Type-alias for the default random-number generator that is used in absense
/// of another user-selected one.
///
/// This defaults to the 64-Bit PCG for `feature="rand_pcg"` with 256-Bit
/// state-space. If instead of `rand_pcg` feature `rand_xoshiro` is active,
/// Xoshiro256++ is used as a default PRNG.
#[cfg(all(
	not(feature = "rand_pcg"),
	feature = "rand_xoshiro",
	target_pointer_width = "32"
))]
pub type DefaultRng = Xoshiro128PlusPlus;

/// A trait that abstracts over the state management of a splittable PRNG.
///
/// This trait provides the core logic for the two-level hierarchy of randomness
/// used by [`RngStream`].
pub trait RngGen {
	/// The internal state necessary to generate random-number generators.
	type State;

	/// Initializes the state using a default value.
	fn default_state() -> Self::State;

	/// Returns an independent PRNG for use within a single simulation.
	///
	/// This method advances the internal state by a small, fixed amount (a
	/// "jump") and returns a new PRNG instance. Each call produces an
	/// independent generator for different components within a simulation run.
	fn next_rng(state: &mut Self::State) -> Self;

	/// Splits the stream of streams, advancing the current state and returning
	/// a new state for a sub-stream.
	///
	/// This is used for creating independent simulation runs. It partitions the
	/// state space, returning the current state for a new `RngStream` and
	/// advancing the parent state by a large amount corresponding to `count`
	/// full streams.
	fn split_at(state: &mut Self::State, count: usize) -> Self::State;
}

/// Represents a stream of independent random-number generators.
///
/// An `RngStream` is the main entry point for randomness in a simulation.
///
/// Use static methods like [`RngStream::iter()`] or [`RngStream::par_iter()`]
/// to create sequences of independent streams for multiple simulation runs.
///
/// Use instance methods like [`rng()`][RngStream::rng] or iterate over an
/// instance to get individual PRNGs for tasks within a single simulation.
///
/// # Usage
///
/// ### Sequential Simulation with Multiple RNGs
/// Inside a single simulation, you can create multiple independent PRNGs.
///
/// ```
/// # use odem_rs_util::random::{RngStream, DefaultRng};
/// # use rand::RngExt;
/// let rng_stream = RngStream::<DefaultRng>::default();
///
/// // Get an RNG for event generation.
/// let mut event_rng = rng_stream.rng();
/// // Get another, independent RNG for agent decisions.
/// let mut agent_rng = rng_stream.rng();
///
/// println!("Event: {}", event_rng.random::<u64>());
/// println!("Agent: {}", agent_rng.random::<u64>());
/// ```
///
/// ### Multiple Sequential Simulation Runs
/// For sequential runs, use [`RngStream::iter()`] combined with `take()` to
/// specify the number of runs.
///
/// ```
/// # use odem_rs_util::random::{RngStream, DefaultRng};
/// // Create two independent simulation streams sequentially.
/// for rng_stream in RngStream::<DefaultRng>::iter().take(2) {
///     // ... run a full simulation with `rng_stream` ...
/// }
/// ```
///
/// ### Multiple Parallel Simulation Runs
/// For parallel runs, use [`RngStream::par_iter()`] combined with `take()` to
/// specify the number of runs. This is the idiomatic approach when using
/// `rayon` and requires feature `"rayon"` to be enabled.
///
/// ```no_run
/// # #[cfg(rayon)]
/// # mod rayon_specific {
/// # use odem_rs_util::{random::{RngStream, DefaultRng}, random_variable::RandomVariable};
/// # use rayon::prelude::*;
/// # fn sim_main(rng_stream: RngStream) -> f64 { 0.0 }
/// # fn main() {
/// let stats: RandomVariable<f64> = RngStream::par_iter()
///     .take(100) // <- run 100 experiments in parallel
///     .map(|rng_stream| sim_main(rng_stream))
///     .fold(RandomVariable::default, |var, duration| {
///         var.tabulate(duration);
///         var
///     })
///     .reduce(RandomVariable::default, RandomVariable::join);
/// # }}
/// ```
pub struct RngStream<R: RngGen = DefaultRng>(RefCell<R::State>);

impl<R: RngGen> RngStream<R> {
	/// Returns an infinite iterator over independent `RngStream`s.
	///
	/// This is the entry point for creating multiple sequential simulation
	/// runs. Use `take(n)` to specify the number of streams you need.
	pub fn iter() -> RngStreamIter<R> {
		RngStreamIter(RngStream::default())
	}

	/// Returns an infinite parallel iterator over independent `RngStream`s.
	///
	/// This is the preferred method for running multiple simulations in
	/// parallel. It requires the `rayon` feature to be enabled.
	#[cfg(feature = "rayon")]
	#[cfg_attr(docsrs, doc(cfg(feature = "rayon")))]
	pub fn par_iter() -> RngStreamParIter<R> {
		RngStreamParIter::new(RngStream::default())
	}

	/// Creates a new RngStream from the inner state.
	pub const fn new(inner: R::State) -> Self {
		RngStream(RefCell::new(inner))
	}

	/// Creates a new `RngStream` from a seed value.
	///
	/// The `seed` is a positional index into the sequence of possible streams.
	/// This is useful for deterministically recreating a specific simulation
	/// run.
	pub fn from_seed(seed: u64) -> Self {
		let mut default_state = R::default_state();
		R::split_at(&mut default_state, seed as usize);
		RngStream(RefCell::new(default_state))
	}

	/// Creates and returns a new, independent random number generator from this
	/// stream.
	///
	/// This is the primary method for getting an RNG for a task within a
	/// simulation.
	pub fn rng(&self) -> R {
		R::next_rng(&mut self.0.borrow_mut())
	}

	/// Creates an infinite iterator yielding independent random number
	/// generators from this stream.
	pub fn rng_iter(&self) -> RngIter<'_, R> {
		RngIter(self)
	}
}

// --- Trait Implementations (Default, Clone, IntoIterator) ---

impl<R: RngGen> Default for RngStream<R> {
	fn default() -> Self {
		RngStream::new(R::default_state())
	}
}

impl<R: RngGen> Clone for RngStream<R>
where
	R::State: Clone,
{
	fn clone(&self) -> Self {
		RngStream(self.0.clone())
	}
}

impl<R: RngGen> IntoIterator for RngStream<R> {
	type Item = R;
	type IntoIter = RngIntoIter<R>;

	fn into_iter(self) -> Self::IntoIter {
		RngIntoIter(self)
	}
}

impl<'r, R: RngGen> IntoIterator for &'r RngStream<R> {
	type Item = R;
	type IntoIter = RngIter<'r, R>;

	fn into_iter(self) -> Self::IntoIter {
		RngIter(self)
	}
}

/* ************************************************************ Rng Iterators */

/// An infinite iterator over individual PRNGs drawn from a shared
/// [`RngStream`].
///
/// # Usage
///
/// Use it within a simulation to generate a sequence of RNGs:
/// ```
/// # use odem_rs_util::random::RngStream;
/// # fn sim_main(rng_stream: RngStream) {
/// for rng in rng_stream.rng_iter().take(20) {
///     // Use the random-number generators,
///     // e.g., for initializing independent jobs.
/// }
/// # }
/// ```
pub struct RngIter<'r, R: RngGen>(&'r RngStream<R>);

impl<R: RngGen> Iterator for RngIter<'_, R> {
	type Item = R;

	fn next(&mut self) -> Option<Self::Item> {
		Some(self.0.rng())
	}
}

impl<R: RngGen> Clone for RngIter<'_, R> {
	fn clone(&self) -> Self {
		*self
	}
}

impl<R: RngGen> Copy for RngIter<'_, R> {}

/// An infinite owning iterator over individual PRNGs drawn from an
/// [`RngStream`].
///
/// # Usage
///
/// Use it within a simulation to generate a sequence of RNGs:
/// ```
/// # use odem_rs_util::random::RngStream;
/// # fn sim_main(rng_stream: RngStream) {
/// for rng in rng_stream.into_iter().take(20) {
///     // Use the random-number generators,
///     // e.g., for initializing independent jobs.
/// }
/// # }
/// ```
pub struct RngIntoIter<R: RngGen>(RngStream<R>);

impl<R: RngGen> Iterator for RngIntoIter<R> {
	type Item = R;

	fn next(&mut self) -> Option<Self::Item> {
		Some(self.0.rng())
	}
}

impl<R: RngGen> Clone for RngIntoIter<R>
where
	R::State: Clone,
{
	fn clone(&self) -> Self {
		RngIntoIter(self.0.clone())
	}
}

/// An infinite iterator over independent [`RngStream`] generators.
///
/// This is used to set up multiple, independent simulations sequentially.
///
/// # Usage
///
/// Use it during the set-up of a simulation run to generate a sequence of
/// independent `RngStream`s to pass into the simulation to repeat an experiment
/// with different random events.
///
/// ```
/// # use odem_rs_util::random::{DefaultRng, RngStream};
/// # fn run_sim(s: RngStream) -> usize { 0 }
/// let sum: usize = RngStream::iter()
///     .take(100)
///     .map(|rng_stream| run_sim(rng_stream))
///     .sum();
/// ```
pub struct RngStreamIter<R: RngGen>(RngStream<R>);

impl<R: RngGen> Iterator for RngStreamIter<R> {
	type Item = RngStream<R>;

	fn next(&mut self) -> Option<Self::Item> {
		Some(RngStream::new(R::split_at(self.0.0.get_mut(), 1)))
	}

	fn nth(&mut self, n: usize) -> Option<Self::Item> {
		let _left_stream = R::split_at(self.0.0.get_mut(), n);
		self.next()
	}
}

impl<R: RngGen> Clone for RngStreamIter<R>
where
	R::State: Clone,
{
	fn clone(&self) -> Self {
		RngStreamIter(self.0.clone())
	}
}

/* ************************************************* Built-In Implementations */

// Xoshiro-family of PRNGs
#[cfg(feature = "rand_xoshiro")]
macro_rules! impl_stream_xoshiro {
	($($UID:ident),* $(,)?) => {$(
		#[cfg_attr(docsrs, doc(cfg(feature = "rand_xoshiro")))]
		impl RngGen for $UID {
			type State = $UID;

			fn default_state() -> $UID {
				use rand_xoshiro::rand_core::SeedableRng;
				$UID::seed_from_u64(0)
			}

			fn next_rng(state: &mut $UID) -> $UID {
				let res = state.clone();
				state.jump();
				res
			}

			fn split_at(state: &mut $UID, steps: usize) -> $UID {
				let res = state.clone();

				for _ in 0..steps {
					state.long_jump();
				}

				res
			}
		}
	)*};
}

#[cfg(feature = "rand_xoshiro")]
impl_stream_xoshiro!(
	Xoroshiro128Plus,
	Xoroshiro128PlusPlus,
	Xoroshiro128StarStar,
	// Xoshiro128Plus, // doesn't support long_jump
	Xoshiro128PlusPlus,
	Xoshiro128StarStar,
	Xoshiro256Plus,
	Xoshiro256PlusPlus,
	Xoshiro256StarStar,
	Xoshiro512Plus,
	Xoshiro512PlusPlus,
	Xoshiro512StarStar,
);

/// A constant related to the golden ratio, used for incrementing 128-bit PCG
/// sequence numbers.
///
/// This value is `floor((φ - 1) * 2^128)` made odd, where φ is the golden
/// ratio. Using an increment based on an irrational number ensures that
/// successive sequence numbers are well-distributed, avoiding simple
/// arithmetic patterns. This contributes to the statistical independence
/// of the generated streams. The value must be odd to ensure that it is
/// coprime to 2^128, which guarantees that repeated additions will visit
/// every possible sequence number before repeating (a full period).
#[cfg(feature = "rand_pcg")]
const PHI_MINUS_ONE_128: u128 = 210306068529402873165736369884012333109;

/// A constant related to the golden ratio, used for incrementing 64-bit PCG
/// sequence numbers.
///
/// This is the 64-bit equivalent of `PHI_MINUS_ONE_128`. See that constant's
/// documentation for a detailed explanation of the principle.
#[cfg(feature = "rand_pcg")]
const PHI_MINUS_ONE_64: u64 = 11400714819323198485;

// PCG-family of PRNGs
#[cfg(feature = "rand_pcg")]
#[cfg_attr(docsrs, doc(cfg(feature = "rand_pcg")))]
impl RngGen for Pcg32 {
	type State = (u64, u64);

	fn default_state() -> Self::State {
		(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7)
	}

	fn next_rng((state, seq): &mut Self::State) -> Self {
		let mut res = Pcg32::new(*state, *seq);

		// Advance the state significantly. For PCG variants with weaker output
		// functions, creating streams with very close sequence numbers can lead
		// to initial correlations. Advancing the state mitigates this.
		res.advance(1u64 << 32);

		// increase the sequence number
		*seq = seq.wrapping_add(PHI_MINUS_ONE_64);

		res
	}

	fn split_at(state: &mut Self::State, steps: usize) -> Self::State {
		let res = *state;
		state.1 = state
			.1
			.wrapping_add((steps as u64).wrapping_mul(PHI_MINUS_ONE_64));
		res
	}
}

#[cfg(feature = "rand_pcg")]
#[cfg_attr(docsrs, doc(cfg(feature = "rand_pcg")))]
impl RngGen for Pcg64 {
	type State = (u128, u128);

	fn default_state() -> Self::State {
		(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96)
	}

	fn next_rng((state, seq): &mut Self::State) -> Self {
		let mut res = Self::new(*state, *seq);

		// Advance the state significantly. For PCG variants with weaker output
		// functions, creating streams with very close sequence numbers can lead
		// to initial correlations. Advancing the state mitigates this.
		res.advance(1u128 << 64);

		// increase the sequence number
		*seq = seq.wrapping_add(PHI_MINUS_ONE_128);

		res
	}

	fn split_at(state: &mut Self::State, steps: usize) -> Self::State {
		let res = *state;
		state.1 = state
			.1
			.wrapping_add((steps as u128).wrapping_mul(PHI_MINUS_ONE_128));
		res
	}
}

#[cfg(feature = "rand_pcg")]
#[cfg_attr(docsrs, doc(cfg(feature = "rand_pcg")))]
impl RngGen for Pcg64Mcg {
	type State = u128;

	fn default_state() -> Self::State {
		0xcafef00dd15ea5e5
	}

	fn next_rng(state: &mut Self::State) -> Self {
		let mut res = Self::new(*state);

		// Advance the state significantly. For PCG variants with weaker output
		// functions, creating streams with very close sequence numbers can lead
		// to initial correlations. Advancing the state mitigates this.
		res.advance(1u128 << 64);

		// increase the sequence number
		*state = state.wrapping_add(PHI_MINUS_ONE_128);

		res
	}

	fn split_at(state: &mut Self::State, steps: usize) -> Self::State {
		let res = *state;
		*state = state.wrapping_add((steps as u128).wrapping_mul(PHI_MINUS_ONE_128));
		res
	}
}

#[cfg(feature = "rand_pcg")]
#[cfg_attr(docsrs, doc(cfg(feature = "rand_pcg")))]
impl RngGen for Pcg64Dxsm {
	type State = (u128, u128);

	fn default_state() -> Self::State {
		(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96)
	}

	fn next_rng((state, seq): &mut Self::State) -> Self {
		let res = Self::new(*state, *seq);

		// This specific PRNG ensures practical statistical independence through
		// increased complexity of its output function and so does not benefit
		// from the mitigation strategy employed for the other variants.

		// increase the sequence number
		*seq = seq.wrapping_add(PHI_MINUS_ONE_128);

		res
	}

	fn split_at(state: &mut Self::State, steps: usize) -> Self::State {
		let res = *state;
		state.1 = state
			.1
			.wrapping_add((steps as u128).wrapping_mul(PHI_MINUS_ONE_128));
		res
	}
}