Generalise over arbitrary number of dimensions

[marcpabst]

ndarray is truely great, but one thing I really love about numpy is how easy it is to write code that generalises over unspecified dimensions. I found this to be much harder with ndarray. For example, imagine I want to apply the following function along the innermost axis:

fn do_something_1(a: [f64; 3]) -> f64 {
    a.iter().sum::<f64>()
}

fn do_something_2(a: [f64; 3], b: [f64; 4]) -> f64 {
    a.iter().sum::<f64>() + b.iter().sum::<f64>()
}

Let's say I have two arrays, arrayA with shape (100, 20, 3) and arrayB with shape (100, 20, 4). How can an apply functions to the inntermost axis, leaving all other dimensions in place?

I can easily apply do_something_1() with map_axis along the last axis of arrayA, resulting in an array of shape (100, 20). But how can I do the same for do_something_2, iterating over both arrays? I checked the documentation and thought Zip and lanes() might work, but everything I tried either ended up being way too complicated or didn't work at all.

Any ideas?

[bluss]

Thanks for the nice words.

I think I could maybe try to answer if there were some clarifications. As it is, I don't see what it means to apply do_something_2 to an array with shape (100, 20, 3) or (100, 20, 4). Maybe you could be even more specific, with concrete examples for two different arrays? To some extent, the heading of the question and the body seem to go in a bit different directions? Either way, I'm just trying to understand.

[marcpabst]

Thanks for getting back to me and sorry for being unclear! The goal here is to take two arrays with shape (X, Y, Z, N) and (X, Y, Z, M), then iterate over (X, Y, Z), applying a function fn compute(a: [f64; M], b:[f64; N]) -> f64 to produce an array of shape (X, Y, Z). Does that make more sense?

And idealy, it doesn't matter if its 3 additional dimensions, or 2 or 5.

[bluss]

You want kind of the inverse of an axis_iter, I see that.

I think .lanes() is that inverse and it's what you should use. In this case .rows() since you want the last axis?

Zip::from(a.rows()).and(b.rows()).map_collect( apply here )

The doc says

For example, in a 2 × 2 × 3 array, each row is 3 elements long and there are 2 × 2 = 4 rows in total.

We apply that to your examples so we have that for a, b:

Each row is N (or M) elem long and there are X x Y x Z rows in total.

[bluss]

Follow-up question, are the rows guaranteed to match up between a and b when using the rows interface like this? I think so, but it needs to be thought about and verified.

[bluss]

did it work?