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`]** 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)
19//! and [`ndarray_rand::rand_distr`](rand_distr) 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, Data, DataOwned, Dimension, RawData};
39#[cfg(feature = "quickcheck")]
40use quickcheck::{Arbitrary, Gen};
41
42/// `rand`, re-exported for convenience and version-compatibility.
43pub mod rand
44{
45    pub use rand::*;
46}
47
48/// `rand-distr`, re-exported for convenience and version-compatibility.
49pub mod rand_distr
50{
51    pub use rand_distr::*;
52}
53
54/// Constructors for n-dimensional arrays with random elements.
55///
56/// This trait extends ndarray’s `ArrayBase` and can not be implemented
57/// for other types.
58///
59/// The default RNG is a fast automatically seeded rng (currently
60/// [`rand::rngs::SmallRng`], seeded from [`rand::thread_rng`]).
61///
62/// Note that `SmallRng` is cheap to initialize and fast, but it may generate
63/// low-quality random numbers, and reproducibility is not guaranteed. See its
64/// documentation for information. You can select a different RNG with
65/// [`.random_using()`](Self::random_using).
66pub trait RandomExt<S, A, D>
67where
68    S: RawData<Elem = A>,
69    D: Dimension,
70{
71    /// Create an array with shape `dim` with elements drawn from
72    /// `distribution` using the default RNG.
73    ///
74    /// ***Panics*** if creation of the RNG fails or if the number of elements
75    /// overflows usize.
76    ///
77    /// ```
78    /// use ndarray::Array;
79    /// use ndarray_rand::RandomExt;
80    /// use ndarray_rand::rand_distr::Uniform;
81    ///
82    /// # fn main() {
83    /// let a = Array::random((2, 5), Uniform::new(0., 10.));
84    /// println!("{:8.4}", a);
85    /// // Example Output:
86    /// // [[  8.6900,   6.9824,   3.8922,   6.5861,   2.4890],
87    /// //  [  0.0914,   5.5186,   5.8135,   5.2361,   3.1879]]
88    /// # }
89    fn random<Sh, IdS>(shape: Sh, distribution: IdS) -> ArrayBase<S, D>
90    where
91        IdS: Distribution<S::Elem>,
92        S: DataOwned<Elem = A>,
93        Sh: ShapeBuilder<Dim = D>;
94
95    /// Create an array with shape `dim` with elements drawn from
96    /// `distribution`, using a specific Rng `rng`.
97    ///
98    /// ***Panics*** if the number of elements overflows usize.
99    ///
100    /// ```
101    /// use ndarray::Array;
102    /// use ndarray_rand::RandomExt;
103    /// use ndarray_rand::rand::SeedableRng;
104    /// use ndarray_rand::rand_distr::Uniform;
105    /// use rand_isaac::isaac64::Isaac64Rng;
106    ///
107    /// # fn main() {
108    /// // Get a seeded random number generator for reproducibility (Isaac64 algorithm)
109    /// let seed = 42;
110    /// let mut rng = Isaac64Rng::seed_from_u64(seed);
111    ///
112    /// // Generate a random array using `rng`
113    /// let a = Array::random_using((2, 5), Uniform::new(0., 10.), &mut rng);
114    /// println!("{:8.4}", a);
115    /// // Example Output:
116    /// // [[  8.6900,   6.9824,   3.8922,   6.5861,   2.4890],
117    /// //  [  0.0914,   5.5186,   5.8135,   5.2361,   3.1879]]
118    /// # }
119    fn random_using<Sh, IdS, R>(shape: Sh, distribution: IdS, rng: &mut R) -> ArrayBase<S, D>
120    where
121        IdS: Distribution<S::Elem>,
122        R: Rng + ?Sized,
123        S: DataOwned<Elem = A>,
124        Sh: ShapeBuilder<Dim = D>;
125
126    /// Sample `n_samples` lanes slicing along `axis` using the default RNG.
127    ///
128    /// If `strategy==SamplingStrategy::WithoutReplacement`, each lane can only be sampled once.
129    /// If `strategy==SamplingStrategy::WithReplacement`, each lane can be sampled multiple times.
130    ///
131    /// ***Panics*** when:
132    /// - creation of the RNG fails;
133    /// - `n_samples` is greater than the length of `axis` (if sampling without replacement);
134    /// - length of `axis` is 0.
135    ///
136    /// ```
137    /// use ndarray::{array, Axis};
138    /// use ndarray_rand::{RandomExt, SamplingStrategy};
139    ///
140    /// # fn main() {
141    /// let a = array![
142    ///     [1., 2., 3.],
143    ///     [4., 5., 6.],
144    ///     [7., 8., 9.],
145    ///     [10., 11., 12.],
146    /// ];
147    /// // Sample 2 rows, without replacement
148    /// let sample_rows = a.sample_axis(Axis(0), 2, SamplingStrategy::WithoutReplacement);
149    /// println!("{:?}", sample_rows);
150    /// // Example Output: (1st and 3rd rows)
151    /// // [
152    /// //  [1., 2., 3.],
153    /// //  [7., 8., 9.]
154    /// // ]
155    /// // Sample 2 columns, with replacement
156    /// let sample_columns = a.sample_axis(Axis(1), 1, SamplingStrategy::WithReplacement);
157    /// println!("{:?}", sample_columns);
158    /// // Example Output: (2nd column, sampled twice)
159    /// // [
160    /// //  [2., 2.],
161    /// //  [5., 5.],
162    /// //  [8., 8.],
163    /// //  [11., 11.]
164    /// // ]
165    /// # }
166    /// ```
167    fn sample_axis(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy) -> Array<A, D>
168    where
169        A: Copy,
170        S: Data<Elem = A>,
171        D: RemoveAxis;
172
173    /// Sample `n_samples` lanes slicing along `axis` using the specified RNG `rng`.
174    ///
175    /// If `strategy==SamplingStrategy::WithoutReplacement`, each lane can only be sampled once.
176    /// If `strategy==SamplingStrategy::WithReplacement`, each lane can be sampled multiple times.
177    ///
178    /// ***Panics*** when:
179    /// - creation of the RNG fails;
180    /// - `n_samples` is greater than the length of `axis` (if sampling without replacement);
181    /// - length of `axis` is 0.
182    ///
183    /// ```
184    /// use ndarray::{array, Axis};
185    /// use ndarray_rand::{RandomExt, SamplingStrategy};
186    /// use ndarray_rand::rand::SeedableRng;
187    /// use rand_isaac::isaac64::Isaac64Rng;
188    ///
189    /// # fn main() {
190    /// // Get a seeded random number generator for reproducibility (Isaac64 algorithm)
191    /// let seed = 42;
192    /// let mut rng = Isaac64Rng::seed_from_u64(seed);
193    ///
194    /// let a = array![
195    ///     [1., 2., 3.],
196    ///     [4., 5., 6.],
197    ///     [7., 8., 9.],
198    ///     [10., 11., 12.],
199    /// ];
200    /// // Sample 2 rows, without replacement
201    /// let sample_rows = a.sample_axis_using(Axis(0), 2, SamplingStrategy::WithoutReplacement, &mut rng);
202    /// println!("{:?}", sample_rows);
203    /// // Example Output: (1st and 3rd rows)
204    /// // [
205    /// //  [1., 2., 3.],
206    /// //  [7., 8., 9.]
207    /// // ]
208    ///
209    /// // Sample 2 columns, with replacement
210    /// let sample_columns = a.sample_axis_using(Axis(1), 1, SamplingStrategy::WithReplacement, &mut rng);
211    /// println!("{:?}", sample_columns);
212    /// // Example Output: (2nd column, sampled twice)
213    /// // [
214    /// //  [2., 2.],
215    /// //  [5., 5.],
216    /// //  [8., 8.],
217    /// //  [11., 11.]
218    /// // ]
219    /// # }
220    /// ```
221    fn sample_axis_using<R>(
222        &self, axis: Axis, n_samples: usize, strategy: SamplingStrategy, rng: &mut R,
223    ) -> Array<A, D>
224    where
225        R: Rng + ?Sized,
226        A: Copy,
227        S: Data<Elem = A>,
228        D: RemoveAxis;
229}
230
231impl<S, A, D> RandomExt<S, A, D> for ArrayBase<S, D>
232where
233    S: RawData<Elem = A>,
234    D: Dimension,
235{
236    fn random<Sh, IdS>(shape: Sh, dist: IdS) -> ArrayBase<S, D>
237    where
238        IdS: Distribution<S::Elem>,
239        S: DataOwned<Elem = A>,
240        Sh: ShapeBuilder<Dim = D>,
241    {
242        Self::random_using(shape, dist, &mut get_rng())
243    }
244
245    fn random_using<Sh, IdS, R>(shape: Sh, dist: IdS, rng: &mut R) -> ArrayBase<S, D>
246    where
247        IdS: Distribution<S::Elem>,
248        R: Rng + ?Sized,
249        S: DataOwned<Elem = A>,
250        Sh: ShapeBuilder<Dim = D>,
251    {
252        Self::from_shape_simple_fn(shape, move || dist.sample(rng))
253    }
254
255    fn sample_axis(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy) -> Array<A, D>
256    where
257        A: Copy,
258        S: Data<Elem = A>,
259        D: RemoveAxis,
260    {
261        self.sample_axis_using(axis, n_samples, strategy, &mut get_rng())
262    }
263
264    fn sample_axis_using<R>(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy, rng: &mut R) -> Array<A, D>
265    where
266        R: Rng + ?Sized,
267        A: Copy,
268        S: Data<Elem = A>,
269        D: RemoveAxis,
270    {
271        let indices: Vec<_> = match strategy {
272            SamplingStrategy::WithReplacement => {
273                let distribution = Uniform::from(0..self.len_of(axis));
274                (0..n_samples).map(|_| distribution.sample(rng)).collect()
275            }
276            SamplingStrategy::WithoutReplacement => index::sample(rng, self.len_of(axis), n_samples).into_vec(),
277        };
278        self.select(axis, &indices)
279    }
280}
281
282/// Used as parameter in [`sample_axis`] and [`sample_axis_using`] to determine
283/// if lanes from the original array should only be sampled once (*without replacement*) or
284/// multiple times (*with replacement*).
285///
286/// [`sample_axis`]: RandomExt::sample_axis
287/// [`sample_axis_using`]: RandomExt::sample_axis_using
288#[derive(Debug, Clone)]
289pub enum SamplingStrategy
290{
291    WithReplacement,
292    WithoutReplacement,
293}
294
295// `Arbitrary` enables `quickcheck` to generate random `SamplingStrategy` values for testing.
296#[cfg(feature = "quickcheck")]
297impl Arbitrary for SamplingStrategy
298{
299    fn arbitrary(g: &mut Gen) -> Self
300    {
301        if bool::arbitrary(g) {
302            SamplingStrategy::WithReplacement
303        } else {
304            SamplingStrategy::WithoutReplacement
305        }
306    }
307}
308
309fn get_rng() -> SmallRng
310{
311    SmallRng::from_rng(thread_rng()).expect("create SmallRng from thread_rng failed")
312}