use crate::array::Array;
use crate::error::{NumRs2Error, Result};
use num_traits::{Float, NumCast, ToPrimitive};
use scirs2_core::ndarray::distributions::uniform::SampleUniform;
use scirs2_core::random::prelude::*;
use scirs2_stats::distributions::{
lognormal::Lognormal, Bernoulli, Beta, Binomial, ChiSquare, Exponential, Gamma, Normal,
Poisson, Uniform, Weibull,
};
use std::fmt::Debug;
use std::fmt::Display;
use std::sync::{Arc, Mutex};
pub trait BitGenerator {
fn next_u64(&mut self) -> u64;
fn next_u32(&mut self) -> u32;
fn next_f64(&mut self) -> f64;
fn seed(&mut self, seed: u64);
}
pub struct StdBitGenerator {
rng: StdRng,
}
impl StdBitGenerator {
pub fn new(seed: u64) -> Self {
Self {
rng: StdRng::seed_from_u64(seed),
}
}
pub fn new_random() -> Self {
let mut rng = thread_rng();
let seed = rng.random::<u64>();
Self::new(seed)
}
}
impl BitGenerator for StdBitGenerator {
fn next_u64(&mut self) -> u64 {
self.rng.random()
}
fn next_u32(&mut self) -> u32 {
self.rng.random()
}
fn next_f64(&mut self) -> f64 {
self.rng.random()
}
fn seed(&mut self, seed: u64) {
self.rng = StdRng::seed_from_u64(seed);
}
}
pub struct PCG64BitGenerator {
state: u128,
inc: u128,
multiplier: u128,
}
impl PCG64BitGenerator {
pub fn new(seed: u64) -> Self {
let state = (seed as u128) << 64 | seed as u128;
let inc = ((seed.wrapping_add(1) as u128) << 64) | 1;
let multiplier = 0x2360ED051FC65DA44385DF649FCCF645;
let mut gen = Self {
state,
inc,
multiplier,
};
for _ in 0..10 {
gen.next_u64();
}
gen
}
pub fn new_random() -> Self {
let mut rng = thread_rng();
let seed = rng.random::<u64>();
Self::new(seed)
}
pub fn with_state_and_inc(state: u128, inc: u128) -> Self {
let multiplier = 0x2360ED051FC65DA44385DF649FCCF645;
Self {
state,
inc,
multiplier,
}
}
pub fn get_state(&self) -> u128 {
self.state
}
pub fn get_inc(&self) -> u128 {
self.inc
}
}
impl BitGenerator for PCG64BitGenerator {
fn next_u64(&mut self) -> u64 {
let old_state = self.state;
self.state = old_state
.wrapping_mul(self.multiplier)
.wrapping_add(self.inc);
let xorshifted = (((old_state >> 64) ^ old_state) >> 64) as u64;
let rot = (old_state >> 122) as u32;
xorshifted.rotate_right(rot)
}
fn next_u32(&mut self) -> u32 {
(self.next_u64() >> 32) as u32
}
fn next_f64(&mut self) -> f64 {
(self.next_u64() >> 11) as f64 * (1.0 / (1u64 << 53) as f64)
}
fn seed(&mut self, seed: u64) {
*self = PCG64BitGenerator::new(seed);
}
}
pub struct Generator<B: BitGenerator> {
bit_generator: Arc<Mutex<B>>,
}
impl<B: BitGenerator> Generator<B> {
pub fn new(bit_generator: B) -> Self {
Self {
bit_generator: Arc::new(Mutex::new(bit_generator)),
}
}
fn get_bit_generator(&self) -> Result<std::sync::MutexGuard<'_, B>> {
self.bit_generator.lock().map_err(|_| {
NumRs2Error::InvalidOperation("Failed to acquire bit generator lock".to_string())
})
}
pub fn random<T>(&self, shape: &[usize]) -> Result<Array<T>>
where
T: Clone + Float + NumCast,
{
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let dist = Uniform::new(0.0, 1.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create uniform distribution: {}", e))
})?;
for _ in 0..size {
let val_f64 = dist.rvs(1).expect("uniform sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert uniform sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn integers<
T: Clone + PartialOrd + SampleUniform + Into<i64> + TryFrom<i64> + ToPrimitive,
>(
&self,
low: T,
high: T,
shape: &[usize],
) -> Result<Array<T>>
where
<T as TryFrom<i64>>::Error: Debug,
{
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let low_f64 = low
.clone()
.into()
.to_f64()
.expect("integers: low bound should be convertible to f64");
let high_f64 = high
.clone()
.into()
.to_f64()
.expect("integers: high bound should be convertible to f64");
let uniform_dist = Uniform::new(low_f64, high_f64)
.expect("integers: uniform distribution should be valid for given bounds");
let val_f64 = uniform_dist.rvs(1).expect("uniform sampling failed")[0];
let val_i64 = val_f64.floor() as i64;
let val = T::try_from(val_i64).map_err(|_| {
NumRs2Error::InvalidOperation(
"Failed to convert integer sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn normal<T: Float + NumCast + Clone + Debug + Display>(
&self,
mean: T,
std: T,
shape: &[usize],
) -> Result<Array<T>> {
if std <= T::zero() {
return Err(NumRs2Error::InvalidOperation(format!(
"Standard deviation must be positive, got {}",
std
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let mean_f64 = mean.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert mean to f64".to_string())
})?;
let std_f64 = std.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert std to f64".to_string())
})?;
let dist = Normal::new(mean_f64, std_f64).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create normal distribution: {}", e))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert normal sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn standard_normal<T: Float + NumCast + Clone + Debug + Display>(
&self,
shape: &[usize],
) -> Result<Array<T>> {
self.normal(T::zero(), T::one(), shape)
}
pub fn lognormal<T: Float + NumCast + Clone + Debug + Display>(
&self,
mean: T,
sigma: T,
shape: &[usize],
) -> Result<Array<T>> {
if sigma <= T::zero() {
return Err(NumRs2Error::InvalidOperation(format!(
"Sigma must be positive, got {}",
sigma
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let mean_f64 = mean.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert mean to f64".to_string())
})?;
let sigma_f64 = sigma.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert sigma to f64".to_string())
})?;
let dist = Lognormal::new(mean_f64, sigma_f64, 0.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!(
"Failed to create log-normal distribution: {}",
e
))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert lognormal sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn beta<T: Float + NumCast + Clone + Debug + Display>(
&self,
a: T,
b: T,
shape: &[usize],
) -> Result<Array<T>> {
if a <= T::zero() || b <= T::zero() {
return Err(NumRs2Error::InvalidOperation(format!(
"Alpha and Beta parameters must be positive, got alpha={}, beta={}",
a, b
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let a_f64 = a.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert alpha parameter to f64".to_string())
})?;
let b_f64 = b.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert beta parameter to f64".to_string())
})?;
let dist = Beta::new(a_f64, b_f64, 0.0, 1.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create beta distribution: {}", e))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert beta sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn chisquare<T: Float + NumCast + Clone + Debug + Display>(
&self,
df: T,
shape: &[usize],
) -> Result<Array<T>> {
if df <= T::zero() {
return Err(NumRs2Error::InvalidOperation(format!(
"Degrees of freedom must be positive, got {}",
df
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let df_f64 = df.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert degrees of freedom to f64".to_string())
})?;
let dist = ChiSquare::new(df_f64, 0.0, 1.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!(
"Failed to create chi-square distribution: {}",
e
))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert chi-square sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn gamma<T: Float + NumCast + Clone + Debug + Display>(
&self,
shape_param: T,
scale: T,
size_shape: &[usize],
) -> Result<Array<T>> {
if shape_param <= T::zero() || scale <= T::zero() {
return Err(NumRs2Error::InvalidOperation(format!(
"Shape and scale parameters must be positive, got shape={}, scale={}",
shape_param, scale
)));
}
let arr_size: usize = size_shape.iter().product();
let mut vec = Vec::with_capacity(arr_size);
let shape_f64 = shape_param.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert shape to f64".to_string())
})?;
let scale_f64 = scale.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert scale to f64".to_string())
})?;
let corrected_scale = 1.0 / scale_f64;
let dist = Gamma::new(shape_f64, corrected_scale, 0.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create gamma distribution: {}", e))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..arr_size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert gamma sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(size_shape))
}
pub fn exponential<T: Float + NumCast + Clone + Debug + Display>(
&self,
scale: T,
shape: &[usize],
) -> Result<Array<T>> {
if scale <= T::zero() {
return Err(NumRs2Error::InvalidOperation(format!(
"Scale parameter must be positive, got {}",
scale
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let scale_f64 = scale.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert scale to f64".to_string())
})?;
let rate = 1.0 / scale_f64;
let dist = Exponential::new(rate, 0.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!(
"Failed to create exponential distribution: {}",
e
))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert exponential sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn weibull<T: Float + NumCast + Clone + Debug + Display>(
&self,
shape_param: T,
scale: T,
size_shape: &[usize],
) -> Result<Array<T>> {
if shape_param <= T::zero() || scale <= T::zero() {
return Err(NumRs2Error::InvalidOperation(format!(
"Shape and scale parameters must be positive, got shape={}, scale={}",
shape_param, scale
)));
}
let arr_size: usize = size_shape.iter().product();
let mut vec = Vec::with_capacity(arr_size);
let shape_f64 = shape_param.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert shape to f64".to_string())
})?;
let scale_f64 = scale.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert scale to f64".to_string())
})?;
let dist = Weibull::new(shape_f64, scale_f64, 0.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create Weibull distribution: {}", e))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..arr_size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert Weibull sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(size_shape))
}
pub fn uniform<T: Clone + PartialOrd + SampleUniform + ToPrimitive + NumCast>(
&self,
low: T,
high: T,
shape: &[usize],
) -> Result<Array<T>> {
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let low_f64 = low.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert low bound to f64".to_string())
})?;
let high_f64 = high.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert high bound to f64".to_string())
})?;
let uniform_dist = Uniform::new(low_f64, high_f64)
.expect("uniform: uniform distribution should be valid for given bounds");
let val_f64 = uniform_dist.rvs(1).expect("uniform sampling failed")[0];
let val = T::from(val_f64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert uniform sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn bernoulli<T: Float + NumCast + Clone + Debug + Display>(
&self,
p: T,
shape: &[usize],
) -> Result<Array<T>> {
if p < T::zero() || p > T::one() {
return Err(NumRs2Error::InvalidOperation(format!(
"Probability must be in [0, 1], got {}",
p
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let p_f64 = p.to_f64().ok_or_else(|| {
NumRs2Error::InvalidOperation("Failed to convert probability to f64".to_string())
})?;
let dist = Bernoulli::new(p_f64).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create Bernoulli distribution: {}", e))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_f64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val_bool = val_f64 > 0.5; let val = if val_bool { T::one() } else { T::zero() };
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn poisson<T: NumCast + Clone + Debug>(
&self,
lam: f64,
shape: &[usize],
) -> Result<Array<T>> {
if lam <= 0.0 {
return Err(NumRs2Error::InvalidOperation(format!(
"Lambda must be positive, got {}",
lam
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let dist = Poisson::new(lam, 0.0).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create Poisson distribution: {}", e))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_u64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_u64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert Poisson sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn binomial<T: NumCast + Clone + Debug>(
&self,
n: u64,
p: f64,
shape: &[usize],
) -> Result<Array<T>> {
if !(0.0..=1.0).contains(&p) {
return Err(NumRs2Error::InvalidOperation(format!(
"Probability must be in [0, 1], got {}",
p
)));
}
let size: usize = shape.iter().product();
let mut vec = Vec::with_capacity(size);
let dist = Binomial::new(n as usize, p).map_err(|e| {
NumRs2Error::InvalidOperation(format!("Failed to create Binomial distribution: {}", e))
})?;
let mut bit_gen = self.get_bit_generator()?;
for _ in 0..size {
let temp_rng = StdRng::seed_from_u64(bit_gen.next_u64());
let val_u64 = dist.rvs(1).expect("distribution sampling failed")[0];
let val = T::from(val_u64).ok_or_else(|| {
NumRs2Error::InvalidOperation(
"Failed to convert Binomial sample to target type".to_string(),
)
})?;
vec.push(val);
}
Ok(Array::from_vec(vec).reshape(shape))
}
pub fn integers_simple<T: Clone + PartialOrd + SampleUniform + num_traits::NumCast>(
&self,
low: T,
high: T,
shape: &[usize],
) -> Result<Array<T>> {
self.uniform(low, high, shape)
}
pub fn bit_generator(&self) -> Result<std::sync::MutexGuard<'_, B>> {
self.get_bit_generator()
}
}
pub fn default_rng() -> Generator<StdBitGenerator> {
Generator::new(StdBitGenerator::new_random())
}
pub fn seed_rng(seed: u64) -> Generator<StdBitGenerator> {
Generator::new(StdBitGenerator::new(seed))
}
pub fn pcg64_rng() -> Generator<PCG64BitGenerator> {
Generator::new(PCG64BitGenerator::new_random())
}
pub fn pcg64_seed_rng(seed: u64) -> Generator<PCG64BitGenerator> {
Generator::new(PCG64BitGenerator::new(seed))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_rng() {
let rng = default_rng();
let arr = rng
.random::<f64>(&[3, 3])
.expect("test: random should succeed");
assert_eq!(arr.shape(), vec![3, 3]);
}
#[test]
#[ignore = "Seeding behavior changed during SciRS2 migration - requires seeding implementation fix"]
fn test_seed_rng() {
let rng1 = seed_rng(42);
let arr1 = rng1
.random::<f64>(&[3, 3])
.expect("test: random should succeed");
let rng2 = seed_rng(42);
let arr2 = rng2
.random::<f64>(&[3, 3])
.expect("test: random should succeed");
assert_eq!(arr1.to_vec(), arr2.to_vec());
}
#[test]
fn test_generator_normal() {
let rng = default_rng();
let arr = rng
.normal(0.0, 1.0, &[10])
.expect("test: normal should succeed");
assert_eq!(arr.shape(), vec![10]);
}
#[test]
fn test_pcg64_generator() {
let rng = pcg64_rng();
let arr = rng
.random::<f64>(&[3, 3])
.expect("test: random should succeed");
assert_eq!(arr.shape(), vec![3, 3]);
}
#[test]
#[ignore = "Seeding behavior changed during SciRS2 migration - requires seeding implementation fix"]
fn test_pcg64_seed_produces_same_output() {
let rng1 = pcg64_seed_rng(42);
let arr1 = rng1
.random::<f64>(&[5])
.expect("test: random should succeed");
let rng2 = pcg64_seed_rng(42);
let arr2 = rng2
.random::<f64>(&[5])
.expect("test: random should succeed");
assert_eq!(arr1.to_vec(), arr2.to_vec());
}
#[test]
fn test_generator_distributions() {
let rng = default_rng();
let beta_arr = rng
.beta(2.0, 5.0, &[10])
.expect("test: beta should succeed");
assert_eq!(beta_arr.shape(), vec![10]);
let gamma_arr = rng
.gamma(2.0, 2.0, &[10])
.expect("test: gamma should succeed");
assert_eq!(gamma_arr.shape(), vec![10]);
let uniform_arr = rng
.uniform(0.0, 1.0, &[10])
.expect("test: uniform should succeed");
assert_eq!(uniform_arr.shape(), vec![10]);
let binomial_arr = rng
.binomial::<u32>(10, 0.5, &[10])
.expect("test: binomial should succeed");
assert_eq!(binomial_arr.shape(), vec![10]);
let poisson_arr = rng
.poisson::<u32>(5.0, &[10])
.expect("test: poisson should succeed");
assert_eq!(poisson_arr.shape(), vec![10]);
}
#[test]
fn test_pcg64_state() {
let mut rng = PCG64BitGenerator::new(42);
let initial_state = rng.get_state();
for _ in 0..10 {
rng.next_u64();
}
assert_ne!(initial_state, rng.get_state());
rng.seed(42);
let _state_after_reset = rng.get_state();
let rng2 = PCG64BitGenerator::new(42);
assert_eq!(rng.get_state(), rng2.get_state());
}
#[test]
fn test_bit_generator_methods() {
let mut std_rng = StdBitGenerator::new(42);
let mut pcg_rng = PCG64BitGenerator::new(42);
let std_u64 = std_rng.next_u64();
let pcg_u64 = pcg_rng.next_u64();
assert_ne!(std_u64, pcg_u64);
std_rng.seed(42);
assert_eq!(std_rng.next_u64(), std_u64);
pcg_rng.seed(42);
assert_eq!(pcg_rng.next_u64(), pcg_u64);
}
}