FEAT: Add dimension::squeeze to remove dimensions with len == 1

rust-ndarray · bluss · Apr 11, 2021 · Apr 11, 2021 · Mar 15, 2024 · Apr 11, 2021
commit d223bc61c33fe2c0bb694abceee37238456edd77
diff --git a/src/dimension/mod.rs b/src/dimension/mod.rs
@@ -785,6 +785,43 @@ where D: Dimension
     }
 }
 
+/// Remove axes with length one, except never removing the last axis.
+///
+/// This only has effect on dynamic dimensions.
+pub(crate) fn squeeze<D>(dim: &mut D, strides: &mut D)
+where D: Dimension
+{
+    if let Some(_) = D::NDIM {
+        return;
+    }
+    debug_assert_eq!(dim.ndim(), strides.ndim());
+
+    // Count axes with dim == 1; we keep axes with d == 0 or d > 1
+    let mut ndim_new = 0;
+    for &d in dim.slice() {
+        if d != 1 {
+            ndim_new += 1;
+        }
+    }
+    ndim_new = Ord::max(1, ndim_new);
+    let mut new_dim = D::zeros(ndim_new);
+    let mut new_strides = D::zeros(ndim_new);
+    let mut i = 0;
+    for (&d, &s) in izip!(dim.slice(), strides.slice()) {
+        if d != 1 {
+            new_dim[i] = d;
+            new_strides[i] = s;
+            i += 1;
+        }
+    }
+    if i == 0 {
+        new_dim[i] = 1;
+        new_strides[i] = 1;
+    }
+    *dim = new_dim;
+    *strides = new_strides;
+}
+
 #[cfg(test)]
 mod test
 {
@@ -797,6 +834,7 @@ mod test
         slice_min_max,
         slices_intersect,
         solve_linear_diophantine_eq,
+        squeeze,
         IntoDimension,
     };
     use crate::error::{from_kind, ErrorKind};
@@ -1146,4 +1184,35 @@ mod test
             s![.., 3..;6, NewAxis]
         ));
     }
+
+    #[test]
+    #[cfg(feature = "std")]
+    fn test_squeeze()
+    {
+        let dyndim = Dim::<&[usize]>;
+
+        let mut d = dyndim(&[1, 2, 1, 1, 3, 1]);
+        let mut s = dyndim(&[!0, !0, !0, 9, 10, !0]);
+        let dans = dyndim(&[2, 3]);
+        let sans = dyndim(&[!0, 10]);
+        squeeze(&mut d, &mut s);
+        assert_eq!(d, dans);
+        assert_eq!(s, sans);
+
+        let mut d = dyndim(&[1, 1]);
+        let mut s = dyndim(&[3, 4]);
+        let dans = dyndim(&[1]);
+        let sans = dyndim(&[1]);
+        squeeze(&mut d, &mut s);
+        assert_eq!(d, dans);
+        assert_eq!(s, sans);
+
+        let mut d = dyndim(&[0, 1, 3, 4]);
+        let mut s = dyndim(&[2, 3, 4, 5]);
+        let dans = dyndim(&[0, 3, 4]);
+        let sans = dyndim(&[2, 4, 5]);
+        squeeze(&mut d, &mut s);
+        assert_eq!(d, dans);
+        assert_eq!(s, sans);
+    }
 }