concision_init/utils/
rand_utils.rs

1/*
2   Appellation: utils <mod>
3   Contrib: FL03 <jo3mccain@icloud.com>
4*/
5use crate::Initialize;
6use ndarray::{Array, ArrayBase, DataOwned, Dimension, IntoDimension, RawData, ShapeBuilder};
7use num::Num;
8use num::complex::{Complex, ComplexDistribution};
9use rand::{SeedableRng, rngs};
10use rand_distr::{
11    Distribution, StandardNormal,
12    uniform::{SampleUniform, Uniform},
13};
14
15/// Generate a random array of complex numbers with real and imaginary parts in the range [0, 1)
16pub fn randc<A, S, D>(shape: impl IntoDimension<Dim = D>) -> ArrayBase<S, D>
17where
18    A: Clone + Num,
19    D: Dimension,
20    S: RawData + DataOwned<Elem = Complex<A>>,
21    ComplexDistribution<A, A>: Distribution<S::Elem>,
22{
23    let distr = ComplexDistribution::<A, A>::new(A::one(), A::one());
24    ArrayBase::rand(shape, distr)
25}
26
27/// Given a shape, generate a random array using the StandardNormal distribution
28pub fn stdnorm<S, D, Sh>(shape: Sh) -> ArrayBase<S, D>
29where
30    D: Dimension,
31    S: DataOwned,
32    Sh: ShapeBuilder<Dim = D>,
33    StandardNormal: Distribution<S::Elem>,
34{
35    ArrayBase::rand(shape, StandardNormal)
36}
37
38pub fn stdnorm_from_seed<S, D, Sh>(shape: Sh, seed: u64) -> ArrayBase<S, D>
39where
40    D: Dimension,
41    S: DataOwned,
42    Sh: ShapeBuilder<Dim = D>,
43    StandardNormal: Distribution<S::Elem>,
44{
45    ArrayBase::rand_with(
46        shape,
47        StandardNormal,
48        &mut rngs::StdRng::seed_from_u64(seed),
49    )
50}
51/// Creates a random array from a uniform distribution using a given key
52pub fn uniform_from_seed<T, D>(
53    key: u64,
54    start: T,
55    stop: T,
56    shape: impl IntoDimension<Dim = D>,
57) -> crate::Result<Array<T, D>>
58where
59    D: Dimension,
60    T: SampleUniform,
61{
62    Uniform::new(start, stop)
63        .map(|distr| ArrayBase::rand_with(shape, &distr, &mut rngs::StdRng::seed_from_u64(key)))
64        .map_err(Into::into)
65}