rustrust-ndarray

ndarray map_inplace with index


How do you map_inplace or par_map_inplace in an ndarray while using the index of the current element you're mapping over?

array.map_inplace(|val| {
    *val = (expression involving the index)
})

You can keep track of a global counter for the non-parallel version, but this won't work for the parallel version.


Solution

  • For a non-parallel in-place mutation with index, you can use .indexed_iter_mut() like so:

    array.indexed_iter_mut()
        .for_each(|(index, val)| {
            *val = (expression involving the index)
        });
    

    For a parallel version, you would need to use rayon's .par_bridge() since IndexedIterMut isn't inherently a parallel iterator (ndarray uses rayon internally for parallelization):

    use rayon::iter::{ParallelBridge, ParallelIterator};
    
    array.indexed_iter_mut()
        .par_bridge()
        .for_each(|(index, val)| {
            *val = (expression involving the index)
        });