Repository: bluss/rust-ndarray
Pull request #246 (2016-11-22 10:13:03)

Add .axes() iterator and .max_stride_axis()
The axes iterator makes it easy to do any type of axis query.

Add max_stride_axis (for an axis with len > 1), which is what we need for when splitting views for rayon (see ndarray-experimental).

We don't add min_stride_axis yet, don't know if it should use len > 1 or not.


Patch:
diff --git a/src/dimension/axes.rs b/src/dimension/axes.rs
new file mode 100644
index 00000000..891c1d60
--- /dev/null
+++ b/src/dimension/axes.rs
@@ -0,0 +1,107 @@
+
+use {Dimension, Axis, Ix, Ixs};
+
+/// Create a new Axes iterator
+pub fn axes_of<'a, D>(d: &'a D, strides: &'a D) -> Axes<'a, D>
+    where D: Dimension,
+{
+    Axes {
+        dim: d,
+        strides: strides,
+        start: 0,
+        end: d.ndim(),
+    }
+}
+
+/// An iterator over the length and stride of each axis of an array.
+///
+/// See [`.axes()`](struct.ArrayBase.html#method.axes) for more information.
+#[derive(Debug)]
+pub struct Axes<'a, D: 'a> {
+    dim: &'a D,
+    strides: &'a D,
+    start: usize,
+    end: usize,
+}
+
+/// Description of the axis, its length and its stride.
+#[derive(Copy, Clone, Debug)]
+pub struct AxisDescription(pub Axis, pub Ix, pub Ixs);
+
+impl AxisDescription {
+    /// Return axis
+    #[inline(always)]
+    pub fn axis(self) -> Axis { self.0 }
+    /// Return length
+    #[inline(always)]
+    pub fn len(self) -> Ix { self.1 }
+    /// Return stride
+    #[inline(always)]
+    pub fn stride(self) -> Ixs { self.2 }
+}
+
+impl<'a, D> Copy for Axes<'a, D> { }
+impl<'a, D> Clone for Axes<'a, D> {
+    fn clone(&self) -> Self { *self }
+}
+
+impl<'a, D> Iterator for Axes<'a, D>
+    where D: Dimension,
+{
+    /// Description of the axis, its length and its stride.
+    type Item = AxisDescription;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.start < self.end {
+            let i = self.start.post_inc();
+            Some(AxisDescription(Axis(i), self.dim[i], self.strides[i] as Ixs))
+        } else {
+            None
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.end - self.start;
+        (len, Some(len))
+    }
+}
+
+impl<'a, D> DoubleEndedIterator for Axes<'a, D>
+    where D: Dimension,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if self.start < self.end {
+            let i = self.end.pre_dec();
+            Some(AxisDescription(Axis(i), self.dim[i], self.strides[i] as Ixs))
+        } else {
+            None
+        }
+    }
+}
+
+trait IncOps : Copy {
+    fn post_inc(&mut self) -> Self;
+    fn post_dec(&mut self) -> Self;
+    fn pre_dec(&mut self) -> Self;
+}
+
+impl IncOps for usize {
+    #[inline(always)]
+    fn post_inc(&mut self) -> Self {
+        let x = *self;
+        *self += 1;
+        x
+    }
+    #[inline(always)]
+    fn post_dec(&mut self) -> Self {
+        let x = *self;
+        *self -= 1;
+        x
+    }
+    #[inline(always)]
+    fn pre_dec(&mut self) -> Self {
+        *self -= 1;
+        *self
+    }
+}
+
diff --git a/src/dimension/dimension_trait.rs b/src/dimension/dimension_trait.rs
index 3904f81d..88b80665 100644
--- a/src/dimension/dimension_trait.rs
+++ b/src/dimension/dimension_trait.rs
@@ -17,12 +17,14 @@ use {Ix, Ixs, Ix0, Ix1, Ix2, Ix3, IxDyn, Dim, Si};
 use IntoDimension;
 use {ArrayView1, ArrayViewMut1};
 use {zipsl, zipsl_mut, ZipExt};
+use Axis;
 use super::{
     stride_offset,
     stride_offset_checked,
     DimPrivate,
 };
 use super::conversion::Convert;
+use super::axes_of;
 
 /// Array shape and index trait.
 ///
@@ -298,6 +300,36 @@ pub unsafe trait Dimension : Clone + Eq + Debug + Send + Sync + Default +
         indices.slice_mut().sort_by_key(|&i| strides[i]);
         indices
     }
+
+    /// Compute the minimum stride axis (absolute value), under the constraint
+    /// that the length of the axis is > 1;
+    #[doc(hidden)]
+    fn min_stride_axis(&self, strides: &Self) -> Axis {
+        let n = match self.ndim() {
+            0 => panic!("min_stride_axis: Array must have ndim > 0"),
+            1 => return Axis(0),
+            n => n,
+        };
+        axes_of(self, strides)
+            .rev()
+            .min_by_key(|ax| ax.stride().abs())
+            .map_or(Axis(n - 1), |ax| ax.axis())
+    }
+
+    /// Compute the maximum stride axis (absolute value), under the constraint
+    /// that the length of the axis is > 1;
+    #[doc(hidden)]
+    fn max_stride_axis(&self, strides: &Self) -> Axis {
+        match self.ndim() {
+            0 => panic!("max_stride_axis: Array must have ndim > 0"),
+            1 => return Axis(0),
+            _ => { }
+        }
+        axes_of(self, strides)
+            .filter(|ax| ax.len() > 1)
+            .max_by_key(|ax| ax.stride().abs())
+            .map_or(Axis(0), |ax| ax.axis())
+    }
 }
 
 // utility functions
@@ -379,6 +411,16 @@ unsafe impl Dimension for Dim<[Ix; 1]> {
         Ix1(0)
     }
 
+    #[inline(always)]
+    fn min_stride_axis(&self, _: &Self) -> Axis {
+        Axis(0)
+    }
+
+    #[inline(always)]
+    fn max_stride_axis(&self, _: &Self) -> Axis {
+        Axis(0)
+    }
+
     #[inline]
     fn first_index(&self) -> Option<Self> {
         if get!(self, 0) != 0 {
@@ -476,6 +518,17 @@ unsafe impl Dimension for Dim<[Ix; 2]> {
         if get!(self, 0) as Ixs <= get!(self, 1) as Ixs { Ix2(0, 1) } else { Ix2(1, 0) }
     }
 
+    #[inline]
+    fn min_stride_axis(&self, strides: &Self) -> Axis {
+        let s = get!(strides, 0) as Ixs;
+        let t = get!(strides, 1) as Ixs;
+        if s.abs() < t.abs() {
+            Axis(0)
+        } else {
+            Axis(1)
+        }
+    }
+    
     #[inline]
     fn is_contiguous(dim: &Self, strides: &Self) -> bool {
         let defaults = dim.default_strides();
diff --git a/src/dimension/mod.rs b/src/dimension/mod.rs
index 86339e7c..51eb4559 100644
--- a/src/dimension/mod.rs
+++ b/src/dimension/mod.rs
@@ -16,6 +16,7 @@ pub use self::conversion::IntoDimension;
 pub use self::dimension_trait::Dimension;
 pub use self::ndindex::NdIndex;
 pub use self::remove_axis::RemoveAxis;
+pub use self::axes::{axes_of, Axes, AxisDescription};
 
 #[macro_use] mod macros;
 mod axis;
@@ -24,6 +25,7 @@ pub mod dim;
 mod dimension_trait;
 mod ndindex;
 mod remove_axis;
+mod axes;
 
 /// Private constructor and accessors for Dim
 pub trait DimPrivate<I> {
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index 58487c72..2c11d692 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -21,6 +21,7 @@ use error::{self, ShapeError};
 use super::zipsl;
 use super::ZipExt;
 use dimension::IntoDimension;
+use dimension::{axes_of, Axes};
 
 use {
     NdIndex,
@@ -1002,6 +1003,25 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
         self.view().reversed_axes()
     }
 
+    /// Return an iterator over the length and stride of each axis.
+    pub fn axes(&self) -> Axes<D> {
+        axes_of(&self.dim, &self.strides)
+    }
+
+    /*
+    /// Return the axis with the least stride (by absolute value)
+    pub fn min_stride_axis(&self) -> Axis {
+        self.dim.min_stride_axis(&self.strides)
+    }
+    */
+
+    /// Return the axis with the greatest stride (by absolute value),
+    /// preferring axes with len > 1.
+    pub fn max_stride_axis(&self) -> Axis {
+        self.dim.max_stride_axis(&self.strides)
+    }
+
+
     fn pointer_is_inbounds(&self) -> bool {
         let slc = self.data._data_slice();
         if slc.is_empty() {
diff --git a/src/lib.rs b/src/lib.rs
index be960fe3..8dc58767 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -87,6 +87,8 @@ pub use dimension::{
     IntoDimension,
     RemoveAxis,
     Axis,
+    Axes,
+    AxisDescription,
 };
 pub use dimension::dim::*;
 


Test patch:
diff --git a/tests/dimension.rs b/tests/dimension.rs
index f57fbda2..37d36670 100644
--- a/tests/dimension.rs
+++ b/tests/dimension.rs
@@ -59,6 +59,66 @@ fn fastest_varying_order() {
     assert_eq!(Dim([2, 2, 3, 1, 2])._fastest_varying_stride_order(), [3, 0, 1, 4, 2]);
 }
 
+type ArrayF32<D> = Array<f32, D>;
+
+/*
+#[test]
+fn min_stride_axis() {
+    let a = ArrayF32::zeros(10);
+    assert_eq!(a.min_stride_axis(), Axis(0));
+
+    let a = ArrayF32::zeros((3, 3));
+    assert_eq!(a.min_stride_axis(), Axis(1));
+    assert_eq!(a.t().min_stride_axis(), Axis(0));
+
+    let a = ArrayF32::zeros(vec![3, 3]);
+    assert_eq!(a.min_stride_axis(), Axis(1));
+    assert_eq!(a.t().min_stride_axis(), Axis(0));
+
+    let min_axis = a.axes().min_by_key(|t| t.2.abs()).unwrap().axis();
+    assert_eq!(min_axis, Axis(1));
+
+    let mut b = ArrayF32::zeros(vec![2, 3, 4, 5]);
+    assert_eq!(b.min_stride_axis(), Axis(3));
+    for ax in 0..3 {
+        b.swap_axes(3, ax);
+        assert_eq!(b.min_stride_axis(), Axis(ax));
+        b.swap_axes(3, ax);
+    }
+
+    let a = ArrayF32::zeros((3, 3));
+    let v = a.broadcast((8, 3, 3)).unwrap();
+    assert_eq!(v.min_stride_axis(), Axis(0));
+}
+*/
+
+#[test]
+fn max_stride_axis() {
+    let a = ArrayF32::zeros(10);
+    assert_eq!(a.max_stride_axis(), Axis(0));
+
+    let a = ArrayF32::zeros((3, 3));
+    assert_eq!(a.max_stride_axis(), Axis(0));
+    assert_eq!(a.t().max_stride_axis(), Axis(1));
+
+    let a = ArrayF32::zeros(vec![1, 3]);
+    assert_eq!(a.max_stride_axis(), Axis(1));
+    let a = ArrayF32::zeros((1, 3));
+    assert_eq!(a.max_stride_axis(), Axis(1));
+
+    let a = ArrayF32::zeros(vec![3, 3]);
+    assert_eq!(a.max_stride_axis(), Axis(0));
+    assert_eq!(a.t().max_stride_axis(), Axis(1));
+
+    let mut b = ArrayF32::zeros(vec![2, 3, 4, 5]);
+    assert_eq!(b.max_stride_axis(), Axis(0));
+    for ax in 1..b.ndim() {
+        b.swap_axes(0, ax);
+        assert_eq!(b.max_stride_axis(), Axis(ax));
+        b.swap_axes(0, ax);
+    }
+}
+
 #[test]
 fn test_indexing() {
     let mut x = Dim([1, 2]);