Extend ndarray-rand to be able to randomly sample from ArrayRef.

Prior to `ndarray` 0.17, the `RandomExt` trait exposed by `ndarray-rand` contained methods for both creating new arrays randomly whole-cloth (`random_using`) and sampling from existing arrays (`sample_axis_using`). With the introduction of reference types in `ndarray` 0.17, users should be able to sample from `ArrayRef` instances as well.

We choose to expose an additional extension trait, `RandomRefExt`, that provides this functionality. We keep the methods on the old trait for backwards compatibility, but collapse the implementation and documentation to the new trait to maintain a single source of truth.

Author: akern40

ndarray-rand/src/lib.rs

@@ -29,12 +29,14 @@
 //! that the items are not compatible (e.g. that a type doesn't implement a
 //! necessary trait).
 
+#![warn(missing_docs)]
+
 use crate::rand::distr::{Distribution, Uniform};
 use crate::rand::rngs::SmallRng;
 use crate::rand::seq::index;
 use crate::rand::{rng, Rng, SeedableRng};
 
-use ndarray::{Array, Axis, RemoveAxis, ShapeBuilder};
+use ndarray::{Array, ArrayRef, Axis, RemoveAxis, ShapeBuilder};
 use ndarray::{ArrayBase, Data, DataOwned, Dimension, RawData};
 #[cfg(feature = "quickcheck")]
 use quickcheck::{Arbitrary, Gen};
@@ -124,6 +126,43 @@ where
         S: DataOwned<Elem = A>,
         Sh: ShapeBuilder<Dim = D>;
 
+    /// Sample `n_samples` lanes slicing along `axis` using the default RNG.
+    ///
+    /// See [`RandomRefExt::sample_axis`] for additional information.
+    fn sample_axis(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy) -> Array<A, D>
+    where
+        A: Copy,
+        S: Data<Elem = A>,
+        D: RemoveAxis;
+
+    /// Sample `n_samples` lanes slicing along `axis` using the specified RNG `rng`.
+    ///
+    /// See [`RandomRefExt::sample_axis_using`] for additional information.
+    fn sample_axis_using<R>(
+        &self, axis: Axis, n_samples: usize, strategy: SamplingStrategy, rng: &mut R,
+    ) -> Array<A, D>
+    where
+        R: Rng + ?Sized,
+        A: Copy,
+        S: Data<Elem = A>,
+        D: RemoveAxis;
+}
+
+/// Constructors for sampling from [`ArrayRef`] with random elements.
+///
+/// This trait extends ndarray’s `ArrayRef` and can not be implemented
+/// for other types.
+///
+/// The default RNG is a fast automatically seeded rng (currently
+/// [`rand::rngs::SmallRng`], seeded from [`rand::thread_rng`]).
+///
+/// Note that `SmallRng` is cheap to initialize and fast, but it may generate
+/// low-quality random numbers, and reproducibility is not guaranteed. See its
+/// documentation for information. You can select a different RNG with
+/// [`.random_using()`](Self::random_using).
+pub trait RandomRefExt<A, D>
+where D: Dimension
+{
     /// Sample `n_samples` lanes slicing along `axis` using the default RNG.
     ///
     /// If `strategy==SamplingStrategy::WithoutReplacement`, each lane can only be sampled once.
@@ -168,7 +207,6 @@ where
     fn sample_axis(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy) -> Array<A, D>
     where
         A: Copy,
-        S: Data<Elem = A>,
         D: RemoveAxis;
 
     /// Sample `n_samples` lanes slicing along `axis` using the specified RNG `rng`.
@@ -225,7 +263,6 @@ where
     where
         R: Rng + ?Sized,
         A: Copy,
-        S: Data<Elem = A>,
         D: RemoveAxis;
 }
 
@@ -259,7 +296,7 @@ where
         S: Data<Elem = A>,
         D: RemoveAxis,
     {
-        self.sample_axis_using(axis, n_samples, strategy, &mut get_rng())
+        (**self).sample_axis(axis, n_samples, strategy)
     }
 
     fn sample_axis_using<R>(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy, rng: &mut R) -> Array<A, D>
@@ -268,6 +305,27 @@ where
         A: Copy,
         S: Data<Elem = A>,
         D: RemoveAxis,
+    {
+        (&**self).sample_axis_using(axis, n_samples, strategy, rng)
+    }
+}
+
+impl<A, D> RandomRefExt<A, D> for ArrayRef<A, D>
+where D: Dimension
+{
+    fn sample_axis(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy) -> Array<A, D>
+    where
+        A: Copy,
+        D: RemoveAxis,
+    {
+        self.sample_axis_using(axis, n_samples, strategy, &mut get_rng())
+    }
+
+    fn sample_axis_using<R>(&self, axis: Axis, n_samples: usize, strategy: SamplingStrategy, rng: &mut R) -> Array<A, D>
+    where
+        R: Rng + ?Sized,
+        A: Copy,
+        D: RemoveAxis,
     {
         let indices: Vec<_> = match strategy {
             SamplingStrategy::WithReplacement => {
@@ -284,9 +342,10 @@ where
 /// if lanes from the original array should only be sampled once (*without replacement*) or
 /// multiple times (*with replacement*).
 ///
-/// [`sample_axis`]: RandomExt::sample_axis
-/// [`sample_axis_using`]: RandomExt::sample_axis_using
+/// [`sample_axis`]: RandomRefExt::sample_axis
+/// [`sample_axis_using`]: RandomRefExt::sample_axis_using
 #[derive(Debug, Clone)]
+#[allow(missing_docs)]
 pub enum SamplingStrategy
 {
     WithReplacement,