Merge pull request #495 from jturner314/merge-more-axes

Support merging axes in cases with lengths <= 1

Author: bluss

src/dimension/mod.rs

@@ -287,17 +287,27 @@ pub fn do_slice(
 pub fn merge_axes<D>(dim: &mut D, strides: &mut D, take: Axis, into: Axis) -> bool
     where D: Dimension,
 {
-    let il = dim.axis(into);
-    let is = strides.axis(into) as Ixs;
-    let tl = dim.axis(take);
-    let ts = strides.axis(take) as Ixs;
-    if il as Ixs * is != ts {
-        return false;
+    let into_len = dim.axis(into);
+    let into_stride = strides.axis(into) as isize;
+    let take_len = dim.axis(take);
+    let take_stride = strides.axis(take) as isize;
+    let merged_len = into_len * take_len;
+    if take_len <= 1 {
+        dim.set_axis(into, merged_len);
+        dim.set_axis(take, if merged_len == 0 { 0 } else { 1 });
+        true
+    } else if into_len <= 1 {
+        strides.set_axis(into, take_stride as usize);
+        dim.set_axis(into, merged_len);
+        dim.set_axis(take, if merged_len == 0 { 0 } else { 1 });
+        true
+    } else if take_stride == into_len as isize * into_stride {
+        dim.set_axis(into, merged_len);
+        dim.set_axis(take, 1);
+        true
+    } else {
+        false
     }
-    // merge them
-    dim.set_axis(into, il * tl);
-    dim.set_axis(take, 1);
-    true
 }
 
 

src/impl_methods.rs

@@ -1478,12 +1478,35 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
 
     /// If possible, merge in the axis `take` to `into`.
     ///
+    /// Returns `true` iff the axes are now merged.
+    ///
+    /// This method merges the axes if movement along the two original axes
+    /// (moving fastest along the `into` axis) can be equivalently represented
+    /// as movement along one (merged) axis. Merging the axes preserves this
+    /// order in the merged axis. If `take` and `into` are the same axis, then
+    /// the axis is "merged" if its length is ≤ 1.
+    ///
+    /// If the return value is `true`, then the following hold:
+    ///
+    /// * The new length of the `into` axis is the product of the original
+    ///   lengths of the two axes.
+    ///
+    /// * The new length of the `take` axis is 0 if the product of the original
+    ///   lengths of the two axes is 0, and 1 otherwise.
+    ///
+    /// If the return value is `false`, then merging is not possible, and the
+    /// original shape and strides have been preserved.
+    ///
+    /// Note that the ordering constraint means that if it's possible to merge
+    /// `take` into `into`, it's usually not possible to merge `into` into
+    /// `take`, and vice versa.
+    ///
     /// ```
     /// use ndarray::Array3;
     /// use ndarray::Axis;
     ///
     /// let mut a = Array3::<f64>::zeros((2, 3, 4));
-    /// a.merge_axes(Axis(1), Axis(2));
+    /// assert!(a.merge_axes(Axis(1), Axis(2)));
     /// assert_eq!(a.shape(), &[2, 1, 12]);
     /// ```
     ///

tests/array.rs

@@ -573,6 +573,97 @@ fn diag()
     assert_eq!(d.dim(), 1);
 }
 
+/// Check that the merged shape is correct.
+///
+/// Note that this does not check the strides in the "merged" case!
+#[test]
+fn merge_axes() {
+    macro_rules! assert_merged {
+        ($arr:expr, $slice:expr, $take:expr, $into:expr) => {
+            let mut v = $arr.slice($slice);
+            let merged_len = v.len_of(Axis($take)) * v.len_of(Axis($into));
+            assert!(v.merge_axes(Axis($take), Axis($into)));
+            assert_eq!(v.len_of(Axis($take)), if merged_len == 0 { 0 } else { 1 });
+            assert_eq!(v.len_of(Axis($into)), merged_len);
+        }
+    }
+    macro_rules! assert_not_merged {
+        ($arr:expr, $slice:expr, $take:expr, $into:expr) => {
+            let mut v = $arr.slice($slice);
+            let old_dim = v.raw_dim();
+            let old_strides = v.strides().to_owned();
+            assert!(!v.merge_axes(Axis($take), Axis($into)));
+            assert_eq!(v.raw_dim(), old_dim);
+            assert_eq!(v.strides(), &old_strides[..]);
+        }
+    }
+
+    let a = Array4::<u8>::zeros((3, 4, 5, 4));
+
+    assert_not_merged!(a, s![.., .., .., ..], 0, 0);
+    assert_merged!(a, s![.., .., .., ..], 0, 1);
+    assert_not_merged!(a, s![.., .., .., ..], 0, 2);
+    assert_not_merged!(a, s![.., .., .., ..], 0, 3);
+    assert_not_merged!(a, s![.., .., .., ..], 1, 0);
+    assert_not_merged!(a, s![.., .., .., ..], 1, 1);
+    assert_merged!(a, s![.., .., .., ..], 1, 2);
+    assert_not_merged!(a, s![.., .., .., ..], 1, 3);
+    assert_not_merged!(a, s![.., .., .., ..], 2, 1);
+    assert_not_merged!(a, s![.., .., .., ..], 2, 2);
+    assert_merged!(a, s![.., .., .., ..], 2, 3);
+    assert_not_merged!(a, s![.., .., .., ..], 3, 0);
+    assert_not_merged!(a, s![.., .., .., ..], 3, 1);
+    assert_not_merged!(a, s![.., .., .., ..], 3, 2);
+    assert_not_merged!(a, s![.., .., .., ..], 3, 3);
+
+    assert_merged!(a, s![.., .., .., ..;2], 0, 1);
+    assert_not_merged!(a, s![.., .., .., ..;2], 1, 0);
+    assert_merged!(a, s![.., .., .., ..;2], 1, 2);
+    assert_not_merged!(a, s![.., .., .., ..;2], 2, 1);
+    assert_merged!(a, s![.., .., .., ..;2], 2, 3);
+    assert_not_merged!(a, s![.., .., .., ..;2], 3, 2);
+
+    assert_merged!(a, s![.., .., .., ..3], 0, 1);
+    assert_not_merged!(a, s![.., .., .., ..3], 1, 0);
+    assert_merged!(a, s![.., .., .., ..3], 1, 2);
+    assert_not_merged!(a, s![.., .., .., ..3], 2, 1);
+    assert_not_merged!(a, s![.., .., .., ..3], 2, 3);
+
+    assert_merged!(a, s![.., .., ..;2, ..], 0, 1);
+    assert_not_merged!(a, s![.., .., ..;2, ..], 1, 0);
+    assert_not_merged!(a, s![.., .., ..;2, ..], 1, 2);
+    assert_not_merged!(a, s![.., .., ..;2, ..], 2, 3);
+
+    assert_merged!(a, s![.., ..;2, .., ..], 0, 1);
+    assert_not_merged!(a, s![.., ..;2, .., ..], 1, 0);
+    assert_not_merged!(a, s![.., ..;2, .., ..], 1, 2);
+    assert_merged!(a, s![.., ..;2, .., ..], 2, 3);
+    assert_not_merged!(a, s![.., ..;2, .., ..], 3, 2);
+
+    let a = Array4::<u8>::zeros((3, 1, 5, 1).f());
+    assert_merged!(a, s![.., .., ..;2, ..], 0, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 0, 3);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 0);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 2);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 3);
+    assert_merged!(a, s![.., .., ..;2, ..], 2, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 2, 3);
+    assert_merged!(a, s![.., .., ..;2, ..], 3, 0);
+    assert_merged!(a, s![.., .., ..;2, ..], 3, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 3, 2);
+    assert_merged!(a, s![.., .., ..;2, ..], 3, 3);
+
+    let a = Array4::<u8>::zeros((3, 0, 5, 1));
+    assert_merged!(a, s![.., .., ..;2, ..], 0, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 2, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 3, 1);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 0);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 2);
+    assert_merged!(a, s![.., .., ..;2, ..], 1, 3);
+}
+
 #[test]
 fn swapaxes()
 {