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