Iterate through two axis: flat_map?

[UnderTheCarpet]

Given a [n0, n1, n2, n3] array, I'd like an iterator going through the first two axis. Is this possible? I'd like to use this to parallely go through all components of a matrix-field, where n0 times n1 is the size of the matrix and n2 times n3 the number of grid points.

So kind of what

for mut ax0 in a.outer_iter_mut() {
    for mut ax1 in ax0.outer_iter_mut() {
        //...
    }
}

gives me, but as a flat iterator.

So something along the lines of

a.outer_iter_mut().flat_map(|ax| ax.outer_iter_mut()).par_iter_mut().etc()

which does not work. I'm a bit lost here.

[UnderTheCarpet]

I've tried

a
    .axis_iter_mut(Axis(0))
    .into_par_iter()
    .flat_map(|ax| ax.axis_iter_mut(Axis(0)).into_par_iter())
    .for_each(|mut v| f(&mut v)));

using ndarray-parallel, but this only gives me

|
|    .flat_map(|ax| ax.axis_iter_mut(Axis(0)).into_par_iter())
|                   -- borrow occurs here                   ^ `ax` dropped here while still borrowed
|    .for_each(|mut v| pf(&mut v));
|                                - borrowed value needs to live until here

[UnderTheCarpet]

Solved it by reshaping. So first reshape the array to flatten the first two axes. Make the calculation. Reshape back to old shape.

[bluss]

The flat map issue is unfortunate. It's partly due to the fact how borrowing works and how views / iterators in ndarray can't do some of the neat tricks that references do. A custom view is not a reference.

The difference, very briefly stated is that if we have a reference like view: &'a [T] we can make an iterator Iter<'a> with the same lifetime parameter. If we have an ndarray view, ArrayView<'a>, you can't make such an iterator.(*) You have to give the iterator method a &'b ArrayView<'a> for the self parameter and you get an iterator using borrow 'b.

(*) can't is not entirely accurate. With .into_iter() you can make such an iterator. But the rest of the iterator methods don't have the corresponding .into_* equivalents that so to speak convey the existing borrow instead of making a new.

[bluss]

That said, we don't want to use flat_map anyway, at least not for inner loops, because that iteration will not optimize well.

[UnderTheCarpet]

Reshaping seems to work fine anyway. Reshaping is basically free, right? Since it shouldn't touch the memory of the data?

Closing this.