WIP on .expand() for expanding array element into dimension

bluss · bluss · commit 961e25a92830 · 2021-11-08T21:01:40.000+01:00
diff --git a/src/impl_expand.rs b/src/impl_expand.rs
@@ -0,0 +1,64 @@
+// Copyright 2021 bluss and ndarray developers.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use crate::imp_prelude::*;
+
+use crate::data_traits::RawDataSubst;
+
+use num_complex::Complex;
+
+
+pub trait MultiElement {
+    type Elem;
+    const LEN: usize;
+}
+
+impl<A, const N: usize> MultiElement for [A; N] {
+    type Elem = A;
+    const LEN: usize = N;
+}
+
+impl<A> MultiElement for Complex<A> {
+    type Elem = A;
+    const LEN: usize = 2;
+}
+
+impl<'a, A, D, const N: usize> ArrayView<'a, [A; N], D>
+where
+    D: Dimension,
+{
+    ///
+    /// Note: expanding a zero-element array leads to a new axis of length zero,
+    /// i.e. the array becomes empty.
+    ///
+    /// **Panics** if the product of non-zero axis lengths overflows `isize`.
+    pub fn expand(self, new_axis: Axis) -> ArrayView<'a, A, D::Larger> {
+        let mut strides = self.strides.insert_axis(new_axis);
+        let mut dim = self.dim.insert_axis(new_axis);
+        let len = N as isize;
+        for ax in 0..strides.ndim() {
+            if Axis(ax) == new_axis {
+                continue;
+            }
+            if dim[ax] > 1 {
+                strides[ax] = ((strides[ax] as isize) * len) as usize;
+            }
+        }
+        dim[new_axis.index()] = N;
+        // TODO nicer assertion
+        crate::dimension::size_of_shape_checked(&dim).unwrap();
+
+        // safe because
+        // size still fits in isize;
+        // new strides are adapted to new element type, inside the same allocation.
+        unsafe {
+            ArrayBase::from_data_ptr(self.data.data_subst(), self.ptr.cast())
+                .with_strides_dim(strides, dim)
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,4 +1,9 @@
-// Copyright 2014-2020 bluss and ndarray developers.
+// Copyright 2014-2021 ndarray developers.
+// Main authors:
+//
+// Ulrik Sverdrup "bluss"
+// Jim Turner "jturner314"
+// and many others
 //
 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
@@ -1515,6 +1520,7 @@ mod impl_internal_constructors;
 mod impl_constructors;
 
 mod impl_methods;
+mod impl_expand;
 mod impl_owned_array;
 mod impl_special_element_types;
 
diff --git a/tests/expand.rs b/tests/expand.rs
@@ -0,0 +1,55 @@
+
+use ndarray::prelude::*;
+
+#[test]
+fn test_expand_from_zero() {
+    let a = Array::from_elem((), [[[1, 2], [3, 4], [5, 6]],
+                                  [[11, 12], [13, 14], [15, 16]]]);
+    let av = a.view();
+    println!("{:?}", av);
+    let av = av.expand(Axis(0));
+    println!("{:?}", av);
+    let av = av.expand(Axis(1));
+    println!("{:?}", av);
+    let av = av.expand(Axis(2));
+    println!("{:?}", av);
+    assert!(av.is_standard_layout());
+    assert_eq!(av, av.to_owned());
+
+    let av = a.view();
+    println!("{:?}", av);
+    let av = av.expand(Axis(0));
+    println!("{:?}", av);
+    let av = av.expand(Axis(0));
+    println!("{:?}", av);
+    let av = av.expand(Axis(0));
+    println!("{:?}", av);
+    assert!(av.t().is_standard_layout());
+    assert_eq!(av, av.to_owned());
+}
+
+#[test]
+fn test_expand_zero() {
+    let a = Array::from_elem((3, 4), [0.; 0]);
+
+    for ax in 0..=2 {
+        let mut new_shape = [3, 4, 4];
+        new_shape[1] = if ax == 0 { 3 } else { 4 };
+        new_shape[ax] = 0;
+        let av = a.view();
+        let av = av.expand(Axis(ax));
+        assert_eq!(av.shape(), &new_shape);
+    }
+}
+
+#[test]
+fn test_expand1() {
+    let a = Array::from_elem((3, 3), [1, 2, 3]);
+    println!("{:?}", a);
+    let b = a.view().expand(Axis(2));
+    println!("{:?}", b);
+    let b = a.view().expand(Axis(1));
+    println!("{:?}", b);
+    let b = a.view().expand(Axis(0));
+    println!("{:?}", b);
+}
