Get all 1-dimensional slices of data from N-dimensional array

Question

My goal is to make a function that takes an n-dimensional &ndarray::ArrayD and do an operation on each 1-D slice in each dimension.

Here is an example for a 3-D array:

let values = array![
    [
        [0.0, 0.1, 0.2], // (x0, y0, z0), (x0, y0, z1), (x0, y0, z2)
        [0.3, 0.4, 0.5], // (x0, y1, z0), (x0, y1, z1), (x0, y1, z2)
        [0.6, 0.7, 0.8], // (x0, y2, z0), (x0, y0, z1), (x0, y2, z2)
    ],
    [
        [0.9, 1.0, 1.1], // (x1, y0, z0), (x1, y0, z1), (x1, y0, z2)
        [1.2, 1.3, 1.4], // (x1, y1, z0), (x1, y1, z1), (x1, y1, z2)
        [1.5, 1.6, 1.7], // (x1, y2, z0), (x1, y2, z1), (x1, y2, z2)
    ],
    [
        [1.8, 1.9, 2.0], // (x2, y0, z0), (x2, y0, z1), (x2, y0, z2)
        [2.1, 2.2, 2.3], // (x2, y1, z0), (x2, y1, z1), (x2, y1, z2)
        [2.4, 2.5, 2.6], // (x2, y2, z0), (x2, y2, z1), (x2, y2, z2)
    ],
].into_dyn();

So I want an iterator or similar over these items (preferably uncopied, and order doesn't matter, I just typed them in the easiest order to read):

// 1-D arrays in first direction
[0.0, 0.1, 0.2],
[0.3, 0.4, 0.5],
...
[2.1, 2.2, 2.3],
[2.4, 2.5, 2.6],
// 1-D arrays in second direction
[0.0, 0.3, 0.6],
[0.1, 0.4, 0.7],
...
[1.9, 2.2, 2.5]
[2.0, 2.3, 2.6],
// 1-D arrays in third direction
[0.0, 0.9, 1.8],
[0.1, 1.0, 1.9],
...
[0.7, 1.6, 2.5]
[0.8, 1.7, 2.6],

I tried indexing with a slice of length n - 1 (easily generated by a function that returns all possible n-1 dimensional indeces given the shape, unlike with slice where .. is inserted), but this panics; I think because the dimension of a slice must match the array dimensions. It seems like indexing taking any length from 1 to N would be useful, so maybe that's a topic for an ndarray feature request.

This is the actual operation I want to do on all these 1-D slices, to check that any non-NaN data is continuous within the slice.

// assuming arr is one of the 1-D arrays
assert!(
    arr.windows(2)
        .map(|w| w[0].is_nan() != w[1].is_nan())
        .filter(|&b| b)
        .count()
        <= 2
);
// e.g. the below would fail this assert!
// [f64::NAN, 0.0, 0.0, f64::NAN, 0.0, 0.0, f64::NAN]
// and this would be ok
// [f64::NAN, 0.0, 0.0, 0.0, 0.0, 0.0, f64::NAN]

Any thoughts on a nice way to solve this problem? Thanks!!

Answer 1

Sure, you are looking for the lanes function.

use ndarray::prelude::*;

fn main() {
    let values = array![
        [
            [0.0, 0.1, 0.2], // (x0, y0, z0), (x0, y0, z1), (x0, y0, z2)
            [0.3, 0.4, 0.5], // (x0, y1, z0), (x0, y1, z1), (x0, y1, z2)
            [0.6, 0.7, 0.8], // (x0, y2, z0), (x0, y0, z1), (x0, y2, z2)
        ],
        [
            [0.9, 1.0, 1.1], // (x1, y0, z0), (x1, y0, z1), (x1, y0, z2)
            [1.2, 1.3, 1.4], // (x1, y1, z0), (x1, y1, z1), (x1, y1, z2)
            [1.5, 1.6, 1.7], // (x1, y2, z0), (x1, y2, z1), (x1, y2, z2)
        ],
        [
            [1.8, 1.9, 2.0], // (x2, y0, z0), (x2, y0, z1), (x2, y0, z2)
            [2.1, 2.2, 2.3], // (x2, y1, z0), (x2, y1, z1), (x2, y1, z2)
            [2.4, 2.5, 2.6], // (x2, y2, z0), (x2, y2, z1), (x2, y2, z2)
        ],
    ]
    .into_dyn();

    // Axis 0 is the outermost, n-1 is the innermost
    // So we use rev() to iterate from inner to outer dimension
    for axis in (0..values.ndim()).rev() {
        for lane in values.lanes(Axis(axis)) {
            println!("{:?}", lane);
        }
    }
}

[0.0, 0.1, 0.2], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[0.3, 0.4, 0.5], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[0.6, 0.7, 0.8], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[0.9, 1.0, 1.1], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[1.2, 1.3, 1.4], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[1.5, 1.6, 1.7], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[1.8, 1.9, 2.0], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[2.1, 2.2, 2.3], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[2.4, 2.5, 2.6], shape=[3], strides=[1], layout=CFcf (0xf), const ndim=1
[0.0, 0.3, 0.6], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[0.1, 0.4, 0.7], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[0.2, 0.5, 0.8], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[0.9, 1.2, 1.5], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[1.0, 1.3, 1.6], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[1.1, 1.4, 1.7], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[1.8, 2.1, 2.4], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[1.9, 2.2, 2.5], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[2.0, 2.3, 2.6], shape=[3], strides=[3], layout=Custom (0x0), const ndim=1
[0.0, 0.9, 1.8], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.1, 1.0, 1.9], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.2, 1.1, 2.0], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.3, 1.2, 2.1], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.4, 1.3, 2.2], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.5, 1.4, 2.3], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.6, 1.5, 2.4], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.7, 1.6, 2.5], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1
[0.8, 1.7, 2.6], shape=[3], strides=[9], layout=Custom (0x0), const ndim=1