use rand::distributions::Distribution;
use rand::rngs::SmallRng;
use rand::{thread_rng, Rng, SeedableRng};
use ndarray::ShapeBuilder;
use ndarray::{ArrayBase, DataOwned, Dimension};
pub trait RandomExt<S, D>
where
S: DataOwned,
D: Dimension,
{
fn random<Sh, IdS>(shape: Sh, distribution: IdS) -> ArrayBase<S, D>
where
IdS: Distribution<S::Elem>,
Sh: ShapeBuilder<Dim = D>;
fn random_using<Sh, IdS, R>(shape: Sh, distribution: IdS, rng: &mut R) -> ArrayBase<S, D>
where
IdS: Distribution<S::Elem>,
R: Rng + ?Sized,
Sh: ShapeBuilder<Dim = D>;
}
impl<S, D> RandomExt<S, D> for ArrayBase<S, D>
where
S: DataOwned,
D: Dimension,
{
fn random<Sh, IdS>(shape: Sh, dist: IdS) -> ArrayBase<S, D>
where
IdS: Distribution<S::Elem>,
Sh: ShapeBuilder<Dim = D>,
{
let mut rng =
SmallRng::from_rng(thread_rng()).expect("create SmallRng from thread_rng failed");
Self::random_using(shape, dist, &mut rng)
}
fn random_using<Sh, IdS, R>(shape: Sh, dist: IdS, rng: &mut R) -> ArrayBase<S, D>
where
IdS: Distribution<S::Elem>,
R: Rng + ?Sized,
Sh: ShapeBuilder<Dim = D>,
{
Self::from_shape_fn(shape, |_| dist.sample(rng))
}
}
#[derive(Copy, Clone, Debug)]
pub struct F32<S>(pub S);
impl<S> Distribution<f32> for F32<S>
where
S: Distribution<f64>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
self.0.sample(rng) as f32
}
}
#[cfg(test)]
mod tests {
use ndarray::Array;
use rand::distributions::Uniform;
use crate::ndarray_rand::RandomExt;
#[test]
fn test_dim() {
let (mm, nn) = (5, 5);
for m in 0..mm {
for n in 0..nn {
let a = Array::random((m, n), Uniform::new(0., 2.));
assert_eq!(a.shape(), &[m, n]);
assert!(a.iter().all(|x| *x < 2.));
assert!(a.iter().all(|x| *x >= 0.));
}
}
}
}