map_axis with FnMut(ArrayViewMut1<'a, A>)

[LukeMathWalker]

Currently map_axis accepts only a mapping taking as argument a 1-dimensional ArrayView (https://docs.rs/ndarray/0.11.2/ndarray/struct.ArrayBase.html#method.map_axis).
What if I wanted to use a function that takes an ArrayViewMut1 as argument?

Is there another way to achieve the effect of map_axis using other iterators/Zip/etc.?
I have tried to reason around that but I have not yet been able to come up with a solution...

The use-case I have in mind is a reducer that takes a lane in the array as argument and applies a function that modifies the one-dimensional array in place to return a single value for each lane.

[rcarson3]

Normally, I would use the azip! macro to do something like what you want. I actually use this quite often to zip together different ArrayView types and change one of them or even just a single ArrayViewMut type with nothing else zipped with it. Typically I use the .axis_iter() or .axis_iter_mut() methods in the azip! macro.

However, I can't seem to get it to work with the .lanes() method. I really haven't had a need to use that method, so I'm more unfamiliar with it. However looking at the documents, it looks like it should work, since it implements the NdProducer trait. If anyone more familiar with this method has gotten it to work I'd be curious to see how to use this with azip! or even the Zip iterator.

[LukeMathWalker]

You gave me an idea! I'll work on it tomorrow and then I'll report on the results.
Even though, having had a look at map_axis implementation, my question becomes: what could go wrong replacing ArrayView::new_ with ArrayViewMut::new_ in

    pub fn map_axis<'a, B, F>(&'a self, axis: Axis, mut mapping: F)
        -> Array<B, D::Smaller>
        where D: RemoveAxis,
              F: FnMut(ArrayView1<'a, A>) -> B,
              A: 'a,
    {
        let view_len = self.len_of(axis);
        let view_stride = self.strides.axis(axis);
        // use the 0th subview as a map to each 1d array view extended from
        // the 0th element.
        self.subview(axis, 0).map(|first_elt| {
            unsafe {
                mapping(ArrayView::new_(first_elt, Ix1(view_len), Ix1(view_stride)))
            }
        })
    }

with all the necessary adjustments in the signature to produce a map_axis_mut?

[jturner314]

Something like this should do what you want:

#[macro_use]
extern crate ndarray;

use ndarray::prelude::*;

/// function that modifies the one-dimensional array in place to return a
/// single value for each lane
fn do_something(mut lane: ArrayViewMut1<i32>) -> i32 {
    lane += 1;
    lane[0] * 2
}

fn main() {
    let mut arr = array![[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]];
    let mut out = Array::zeros(arr.view().remove_axis(Axis(1)).raw_dim());
    println!("arr =\n{}", arr);
    println!("out =\n{}", out);
    println!("==================");

    azip!(mut lane (arr.lanes_mut(Axis(1))), mut out in {
        println!("lane = {}", lane);
        *out = do_something(lane);
    });
    println!("==================");
    println!("arr =\n{}", arr);
    println!("out =\n{}", out);
}

Output:

arr =
[[[1, 2, 3],
  [4, 5, 6]],
 [[7, 8, 9],
  [10, 11, 12]]]
out =
[[0, 0, 0],
 [0, 0, 0]]
==================
lane = [1, 4]
lane = [2, 5]
lane = [3, 6]
lane = [7, 10]
lane = [8, 11]
lane = [9, 12]
==================
arr =
[[[2, 3, 4],
  [5, 6, 7]],
 [[8, 9, 10],
  [11, 12, 13]]]
out =
[[4, 6, 8],
 [16, 18, 20]]

We could add a map_axis_mut to make this easier. I'm curious, though -- what's the use-case for this? It seems unusual to want to modify the lanes and return a value in the same operation.

[LukeMathWalker]

Thanks a lot for the solution @jturner314!
The use-case I am working on is k-th percentile computation along an axis: I want to be able to compute the k-th percentile without having to allocate extra memory.
Hence I want to be able to reorder lanes elements in-place and return the k-th percentile along each axis.

In most cases you'd want to keep an untouched copy of the original data, but it's nicer to be able to choose - if you only care about the output you don't want to pay the memory cost of allocating a copy of the original data.

[jturner314]

You're welcome. Thanks for the explanation.

I'm not sure about @bluss, but I wouldn't object to adding a .map_axis_mut() method if you want to submit a PR.

[LukeMathWalker]

I have been trying to implement .map_axis_mut adapting the code in .map_axis but all the necessary adjustments in the signature have proven to be more complicated than I expected.
After a couple of tries my best stab was the following:

    /// Reduce the values along an axis into just one value, producing a new
    /// array with one less dimension.
    ///
    /// Elements are visited in arbitrary order.
    ///
    /// Return the result as an `Array`.
    ///
    /// **Panics** if `axis` is out of bounds.
    pub fn map_axis_mut<'a, B, F>(&'a mut self, axis: Axis, mut mapping: F)
        -> Array<B, D::Smaller>
        where D: RemoveAxis,
              F: FnMut(ArrayViewMut1<'a, A>) -> B,
              A: 'a + Clone,
              S: DataMut,
    {
        let view_len = self.len_of(axis);
        let view_stride = self.strides.axis(axis);
        // use the 0th subview as a map to each 1d array view extended from
        // the 0th element.
        self.subview_mut(axis, 0).mapv(|mut first_elt: A| {
            unsafe {
                mapping(ArrayViewMut::new_(&mut first_elt, Ix1(view_len), Ix1(view_stride)))
            }
        })
    }

that triggers the following error from the compiler:

error[E0387]: cannot borrow data mutably in a captured outer variable in an `Fn` closure
    --> src/impl_methods.rs:1844:17
     |
1844 |                 mapping(ArrayViewMut::new_(&mut first_elt, Ix1(view_len), Ix1(view_stride)))
     |                 ^^^^^^^
     |
help: consider changing this closure to take self by mutable reference
    --> src/impl_methods.rs:1842:40
     |
1842 |           self.subview_mut(axis, 0).mapv(|mut first_elt: A| {
     |  ________________________________________^
1843 | |             unsafe {
1844 | |                 mapping(ArrayViewMut::new_(&mut first_elt, Ix1(view_len), Ix1(view_stride)))
1845 | |             }
1846 | |         })
     | |_________^

I have thus tried to follow its suggestion:

    pub fn map_axis_mut<'a, B, F>(&'a mut self, axis: Axis, mut mapping: F)
        -> Array<B, D::Smaller>
        where D: RemoveAxis,
              F: FnMut(ArrayViewMut1<'a, A>) -> B,
              A: 'a + Clone,
              S: DataMut,
    {
        let view_len = self.len_of(axis);
        let view_stride = self.strides.axis(axis);
        // use the 0th subview as a map to each 1d array view extended from
        // the 0th element.
        self.subview_mut(axis, 0).map(|mut first_elt: &mut A| {
            unsafe {
                mapping(ArrayViewMut::new_(first_elt, Ix1(view_len), Ix1(view_stride)))
            }
        })
    }

but map (or mapv) signature comes into play:

error[E0631]: type mismatch in closure arguments
    --> src/impl_methods.rs:1842:35
     |
1842 |         self.subview_mut(axis, 0).map(|mut first_elt: &mut A| {
     |                                   ^^^ ----------------------- found signature of `for<'r> fn(&'r mut A) -> _`
     |                                   |
     |                                   expected signature of `fn(&A) -> _`

Any idea on how to get out of this pinch?

If you want I can polish your zip example and use that to draft a PR out of that, but I am curious to see how it could be done using this more obscure path @jturner314

[jturner314]

In the first example, the compiler is complaining about mutably borrowing mapping in a non-mutable (Fn) closure. .mapv takes an argument f: F, where F: Fn(A) -> B, so the closure cannot mutably borrow mapping; .mapv would need to take an argument F: FnMut(A) -> B for that to work. Note that a mutable borrow of mapping is necessary because the constraint on mapping: F is F: FnMut....

It's possible to make the compiler happy in this case by changing the F: FnMut(ArrayViewMut1<'a, A>) -> B constraint to F: Fn(ArrayViewMut1<'a, A>) -> B. However, there is a second, perhaps less-obvious issue with this example that the compiler can't catch because it's in the unsafe block.

        self.subview_mut(axis, 0).mapv(|mut first_elt: A| {

.subview_mut(axis, 0) returns the first subview as an ArrayViewMut. .mapv clones an element and passes it as an argument on each call to the closure, so on each call of the closure, first_elt is actually a clone of an element.

                mapping(ArrayViewMut::new_(&mut first_elt, Ix1(view_len), Ix1(view_stride)))

Here, ArrayViewMut::new_ is creating a new mutable view with the given shape and stride, where the pointer to the first element is a pointer to first_elt. This is bad because first_elt is a cloned value on the stack, not an element in the array. If we were to run this code, we'd create an view of data in the stack, not the data in the array.

In the second example, yeah the issue is that .map gives a &A, while we need a &mut A.

What I would suggest is to first add a method .map_mut() (equivalent to .map() but pass &'a mut A to the closure instead of &'a A). Then, implement .map_axis_mut() using .map_mut() essentially the same way .map_axis() is implemented using .map().

[LukeMathWalker]

#460 got merged, so I can close this.