ndarray_rand/
lib.rs

1// Copyright 2016-2019 bluss and ndarray developers.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! Constructors for randomized arrays: `rand` integration for `ndarray`.
10//!
11//! See [**`RandomExt`**](trait.RandomExt.html) for usage examples.
12//!
13//! ## Note
14//!
15//! `ndarray-rand` depends on [`rand` 0.8][rand].
16//!
17//! [`rand`][rand] and [`rand_distr`][rand_distr]
18//! are re-exported as sub-modules, [`ndarray_rand::rand`](rand/index.html)
19//! and [`ndarray_rand::rand_distr`](rand_distr/index.html) respectively.
20//! You can use these submodules for guaranteed version compatibility or
21//! convenience.
22//!
23//! [rand]: https://docs.rs/rand/0.8
24//! [rand_distr]: https://docs.rs/rand_distr/0.4
25//!
26//! If you want to use a random number generator or distribution from another crate
27//! with `ndarray-rand`, you need to make sure that the other crate also depends on the
28//! same version of `rand`. Otherwise, the compiler will return errors saying
29//! that the items are not compatible (e.g. that a type doesn't implement a
30//! necessary trait).
31
32use crate::rand::distributions::{Distribution, Uniform};
33use crate::rand::rngs::SmallRng;
34use crate::rand::seq::index;
35use crate::rand::{thread_rng, Rng, SeedableRng};
36
37use ndarray::{Array, Axis, RemoveAxis, ShapeBuilder};
38use ndarray::{ArrayBase, DataOwned, RawData, Data, Dimension};
39#[cfg(feature = "quickcheck")]
40use quickcheck::{Arbitrary, Gen};
41
42/// `rand`, re-exported for convenience and version-compatibility.
43pub mod rand {
44    pub use rand::*;
45}
46
47/// `rand-distr`, re-exported for convenience and version-compatibility.
48pub mod rand_distr {
49    pub use rand_distr::*;
50}
51
52/// Constructors for n-dimensional arrays with random elements.
53///
54/// This trait extends ndarray’s `ArrayBase` and can not be implemented
55/// for other types.
56///
57/// The default RNG is a fast automatically seeded rng (currently
58/// [`rand::rngs::SmallRng`], seeded from [`rand::thread_rng`]).
59///
60/// Note that `SmallRng` is cheap to initialize and fast, but it may generate
61/// low-quality random numbers, and reproducibility is not guaranteed. See its
62/// documentation for information. You can select a different RNG with
63/// [`.random_using()`](#tymethod.random_using).
64pub trait RandomExt<S, A, D>
65where
66    S: RawData<Elem = A>,
67    D: Dimension,
68{
69    /// Create an array with shape `dim` with elements drawn from
70    /// `distribution` using the default RNG.
71    ///
72    /// ***Panics*** if creation of the RNG fails or if the number of elements
73    /// overflows usize.
74    ///
75    /// ```
76    /// use ndarray::Array;
77    /// use ndarray_rand::RandomExt;
78    /// use ndarray_rand::rand_distr::Uniform;
79    ///
80    /// # fn main() {
81    /// let a = Array::random((2, 5), Uniform::new(0., 10.));
82    /// println!("{:8.4}", a);
83    /// // Example Output:
84    /// // [[  8.6900,   6.9824,   3.8922,   6.5861,   2.4890],
85    /// //  [  0.0914,   5.5186,   5.8135,   5.2361,   3.1879]]
86    /// # }
87    fn random<Sh, IdS>(shape: Sh, distribution: IdS) -> ArrayBase<S, D>
88    where
89        IdS: Distribution<S::Elem>,
90        S: DataOwned<Elem = A>,
91        Sh: ShapeBuilder<Dim = D>;
92
93    /// Create an array with shape `dim` with elements drawn from
94    /// `distribution`, using a specific Rng `rng`.
95    ///
96    /// ***Panics*** if the number of elements overflows usize.
97    ///
98    /// ```
99    /// use ndarray::Array;
100    /// use ndarray_rand::RandomExt;
101    /// use ndarray_rand::rand::SeedableRng;
102    /// use ndarray_rand::rand_distr::Uniform;
103    /// use rand_isaac::isaac64::Isaac64Rng;
104    ///
105    /// # fn main() {
106    /// // Get a seeded random number generator for reproducibility (Isaac64 algorithm)
107    /// let seed = 42;
108    /// let mut rng = Isaac64Rng::seed_from_u64(seed);
109    ///
110    /// // Generate a random array using `rng`
111    /// let a = Array::random_using((2, 5), Uniform::new(0., 10.), &mut rng);
112    /// println!("{:8.4}", a);
113    /// // Example Output:
114    /// // [[  8.6900,   6.9824,   3.8922,   6.5861,   2.4890],
115    /// //  [  0.0914,   5.5186,   5.8135,   5.2361,   3.1879]]
116    /// # }
117    fn random_using<Sh, IdS, R>(shape: Sh, distribution: IdS, rng: &mut R) -> ArrayBase<S, D>
118    where
119        IdS: Distribution<S::Elem>,
120        R: Rng + ?Sized,
121        S: DataOwned<Elem = A>,
122        Sh: ShapeBuilder<Dim = D>;
123
124    /// Sample `n_samples` lanes slicing along `axis` using the default RNG.
125    ///
126    /// If `strategy==SamplingStrategy::WithoutReplacement`, each lane can only be sampled once.
127    /// If `strategy==SamplingStrategy::WithReplacement`, each lane can be sampled multiple times.
128    ///
129    /// ***Panics*** when:
130    /// - creation of the RNG fails;
131    /// - `n_samples` is greater than the length of `axis` (if sampling without replacement);
132    /// - length of `axis` is 0.
133    ///
134    /// ```
135    /// use ndarray::{array, Axis};
136    /// use ndarray_rand::{RandomExt, SamplingStrategy};
137    ///
138    /// # fn main() {
139    /// let a = array![
140    ///     [1., 2., 3.],
141    ///     [4., 5., 6.],
142    ///     [7., 8., 9.],
143    ///     [10., 11., 12.],
144    /// ];
145    /// // Sample 2 rows, without replacement
146    /// let sample_rows = a.sample_axis(Axis(0), 2, SamplingStrategy::WithoutReplacement);
147    /// println!("{:?}", sample_rows);
148    /// // Example Output: (1st and 3rd rows)
149    /// // [
150    /// //  [1., 2., 3.],
151    /// //  [7., 8., 9.]
152    /// // ]
153    /// // Sample 2 columns, with replacement
154    /// let sample_columns = a.sample_axis(Axis(1), 1, SamplingStrategy::WithReplacement);
155    /// println!("{:?}", sample_columns);
156    /// // Example Output: (2nd column, sampled twice)
157    /// // [
158    /// //  [2., 2.],
159    /// //  [5., 5.],
160    /// //  [8., 8.],
161    /// //  [11., 11.]
162    /// // ]
163    /// # }
164    /// ```
165    fn sample_axis(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy) -> Array<A, D>
166    where
167        A: Copy,
168        S: Data<Elem = A>,
169        D: RemoveAxis;
170
171    /// Sample `n_samples` lanes slicing along `axis` using the specified RNG `rng`.
172    ///
173    /// If `strategy==SamplingStrategy::WithoutReplacement`, each lane can only be sampled once.
174    /// If `strategy==SamplingStrategy::WithReplacement`, each lane can be sampled multiple times.
175    ///
176    /// ***Panics*** when:
177    /// - creation of the RNG fails;
178    /// - `n_samples` is greater than the length of `axis` (if sampling without replacement);
179    /// - length of `axis` is 0.
180    ///
181    /// ```
182    /// use ndarray::{array, Axis};
183    /// use ndarray_rand::{RandomExt, SamplingStrategy};
184    /// use ndarray_rand::rand::SeedableRng;
185    /// use rand_isaac::isaac64::Isaac64Rng;
186    ///
187    /// # fn main() {
188    /// // Get a seeded random number generator for reproducibility (Isaac64 algorithm)
189    /// let seed = 42;
190    /// let mut rng = Isaac64Rng::seed_from_u64(seed);
191    ///
192    /// let a = array![
193    ///     [1., 2., 3.],
194    ///     [4., 5., 6.],
195    ///     [7., 8., 9.],
196    ///     [10., 11., 12.],
197    /// ];
198    /// // Sample 2 rows, without replacement
199    /// let sample_rows = a.sample_axis_using(Axis(0), 2, SamplingStrategy::WithoutReplacement, &mut rng);
200    /// println!("{:?}", sample_rows);
201    /// // Example Output: (1st and 3rd rows)
202    /// // [
203    /// //  [1., 2., 3.],
204    /// //  [7., 8., 9.]
205    /// // ]
206    ///
207    /// // Sample 2 columns, with replacement
208    /// let sample_columns = a.sample_axis_using(Axis(1), 1, SamplingStrategy::WithReplacement, &mut rng);
209    /// println!("{:?}", sample_columns);
210    /// // Example Output: (2nd column, sampled twice)
211    /// // [
212    /// //  [2., 2.],
213    /// //  [5., 5.],
214    /// //  [8., 8.],
215    /// //  [11., 11.]
216    /// // ]
217    /// # }
218    /// ```
219    fn sample_axis_using<R>(
220        &self,
221        axis: Axis,
222        n_samples: usize,
223        strategy: SamplingStrategy,
224        rng: &mut R,
225    ) -> Array<A, D>
226    where
227        R: Rng + ?Sized,
228        A: Copy,
229        S: Data<Elem = A>,
230        D: RemoveAxis;
231}
232
233impl<S, A, D> RandomExt<S, A, D> for ArrayBase<S, D>
234where
235    S: RawData<Elem = A>,
236    D: Dimension,
237{
238    fn random<Sh, IdS>(shape: Sh, dist: IdS) -> ArrayBase<S, D>
239    where
240        IdS: Distribution<S::Elem>,
241        S: DataOwned<Elem = A>,
242        Sh: ShapeBuilder<Dim = D>,
243    {
244        Self::random_using(shape, dist, &mut get_rng())
245    }
246
247    fn random_using<Sh, IdS, R>(shape: Sh, dist: IdS, rng: &mut R) -> ArrayBase<S, D>
248    where
249        IdS: Distribution<S::Elem>,
250        R: Rng + ?Sized,
251        S: DataOwned<Elem = A>,
252        Sh: ShapeBuilder<Dim = D>,
253    {
254        Self::from_shape_simple_fn(shape, move || dist.sample(rng))
255    }
256
257    fn sample_axis(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy) -> Array<A, D>
258    where
259        A: Copy,
260        S: Data<Elem = A>,
261        D: RemoveAxis,
262    {
263        self.sample_axis_using(axis, n_samples, strategy, &mut get_rng())
264    }
265
266    fn sample_axis_using<R>(
267        &self,
268        axis: Axis,
269        n_samples: usize,
270        strategy: SamplingStrategy,
271        rng: &mut R,
272    ) -> Array<A, D>
273    where
274        R: Rng + ?Sized,
275        A: Copy,
276        S: Data<Elem = A>,
277        D: RemoveAxis,
278    {
279        let indices: Vec<_> = match strategy {
280            SamplingStrategy::WithReplacement => {
281                let distribution = Uniform::from(0..self.len_of(axis));
282                (0..n_samples).map(|_| distribution.sample(rng)).collect()
283            }
284            SamplingStrategy::WithoutReplacement => {
285                index::sample(rng, self.len_of(axis), n_samples).into_vec()
286            }
287        };
288        self.select(axis, &indices)
289    }
290}
291
292/// Used as parameter in [`sample_axis`] and [`sample_axis_using`] to determine
293/// if lanes from the original array should only be sampled once (*without replacement*) or
294/// multiple times (*with replacement*).
295///
296/// [`sample_axis`]: trait.RandomExt.html#tymethod.sample_axis
297/// [`sample_axis_using`]: trait.RandomExt.html#tymethod.sample_axis_using
298#[derive(Debug, Clone)]
299pub enum SamplingStrategy {
300    WithReplacement,
301    WithoutReplacement,
302}
303
304// `Arbitrary` enables `quickcheck` to generate random `SamplingStrategy` values for testing.
305#[cfg(feature = "quickcheck")]
306impl Arbitrary for SamplingStrategy {
307    fn arbitrary<G: Gen>(g: &mut G) -> Self {
308        if bool::arbitrary(g) {
309            SamplingStrategy::WithReplacement
310        } else {
311            SamplingStrategy::WithoutReplacement
312        }
313    }
314}
315
316fn get_rng() -> SmallRng {
317    SmallRng::from_rng(thread_rng()).expect("create SmallRng from thread_rng failed")
318}
319
320/// A wrapper type that allows casting f64 distributions to f32
321///
322/// ```
323/// use ndarray::Array;
324/// use ndarray_rand::{RandomExt, F32};
325/// use ndarray_rand::rand_distr::Normal;
326///
327/// # fn main() {
328/// let distribution_f64 = Normal::new(0., 1.).expect("Failed to create normal distribution");
329/// let a = Array::random((2, 5), F32(distribution_f64));
330/// println!("{:8.4}", a);
331/// // Example Output:
332/// // [[ -0.6910,   1.1730,   1.0902,  -0.4092,  -1.7340],
333/// //  [ -0.6810,   0.1678,  -0.9487,   0.3150,   1.2981]]
334/// # }
335#[derive(Copy, Clone, Debug)]
336pub struct F32<S>(pub S);
337
338impl<S> Distribution<f32> for F32<S>
339where
340    S: Distribution<f64>,
341{
342    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
343        self.0.sample(rng) as f32
344    }
345}