Indexing `ndarray::ArrayBase` with Unknown Dimensions in Rust

[rivoltella]

I have a function, index_to_indices , that translates flattened indices to array indices based on shape. Since I need to account for row/column sums during indexing, using the flattened index is not enough.

fn index_to_indices(index: usize, shape: &[usize]) -> &[usize] 
{
    let mut indices = vec![0; shape.len()];
    let mut index = index;
    for (i, &s) in shape.iter().enumerate().rev() {
        indices[i] = index % s;
        index /= s;
    }
    &indices
}

fn other_func<R, D>(arr: &ArrayBase<R, D>) 
where R: RawData, D: Dimension
{
   // this function needs an immutable borrow of arr
   todo!()
}

fn calling_fn<R,D>(array: &mut ArrayBase<R, D>)
where R: RawData, D: Dimension 
{
   for index in 0..array.len() {
     let coordinates = index_to_indices(index, &array.shape());
     // borrows arr immutably:
     let value = other_func(arr); 
     // borrows arr mutably:
     array.get_mut(coordinates) += value; //doesn't work because coordinates is not of the proper type
  }
}

Question: How can I leverage index_to_indices 's output (array of indices) to effectively access elements in multi-dimensional arrays and modify them, while the effective dimension of the array is generic?

The issues arise from the fact that iterating over the array with indexed_iter_mut doesn't work because the function I need to pass needs to borrow the array itself, so I'm trying to implement something similar but with a for loop, where I can control where the borrows happen.

[rjc2013]

I think this can only work if you use a "dynamic dimension" array. Dynamic dimension arrays can be indexed using a &[usize]. Regular arrays cannot. According to this page:

NdIndex in ndarray - Rust

the NdIndex trait is only implemented for fixed-size arrays and tuples with the correct number of elements, with the one exception of NdIndex being implemented for &[usize] in the case of dynamic-dimension arrays.

I got it to build the way (I think) you want it to work, using dynamic-dimension arrays:

Rust Playground

[rivoltella]

I finally got it!

The main idea is using D::zeros, where D: Dimension in place of a vector or array of usize.

The idea is:

1. Ix1, Ix2, etc. are concrete types implementing the trait Dimension
2. for indexing an array, we need a type implementing NdIndex
3. NdIndex is implemented for D, where D: Dimension

We just need to use an object implementing the trait Dimension, as it will satisfy the NdIndex bound.

Since Dimension is implemented for tuples or arrays of usize, we can create an empty element and access/modify it with the same syntax I used for &[usize]. Thus, we can use generic programming around the dimension.

Here is the full code with a main function.

For easy reference, here is the main part of the code:

fn index_to_indices<D>(index: usize, shape: &[usize]) -> D
where
    D: Dimension,
{
    let ndim = shape.len();
    // NDIM is None for dynamic dimension, so any shape is ok in that case
    if D::NDIM.unwrap_or(ndim) != ndim {
        panic!("Dimension mismatch");
    }
    let mut indices = D::zeros(ndim);
    let mut index = index;
    for (i, &s) in shape.iter().enumerate().rev() {
        indices[i] = index % s;
        index /= s;
    }
    indices
}

fn calling_fn<R, D>(array: &mut ArrayBase<R, D>)
where
    R: DataMut + Data,
    R::Elem: AddAssign<R::Elem> + Zero,
    D: Dimension + RemoveAxis,
{
    for index in 0..array.len() {
        let coordinates = index_to_indices::<D>(index, &array.shape());
        // borrows arr immutably:
        let value = other_func(coordinates, array);
        // borrows arr mutably:
        *array.get_mut(coordinates).unwrap() += value;
    }
}

[rivoltella]

Actually, just found that ndarray::indices_of returns an iterator that yields Pattern objects.

So the code would become something like:

use ndarray::IntoDimension;

fn calling_fn<R, D>(array: &mut ArrayBase<R, D>)
where
    R: DataMut + Data,
    R::Elem: AddAssign<R::Elem> + Zero + Copy,
    D: Dimension + RemoveAxis + Copy,
{
    for pattern in ndarray::indices_of(array) {
        let coordinates = pattern.into_dimension();
        // borrows arr immutably:
        let value = other_func(&coordinates, array);
        // borrows arr mutably:
        *array.get_mut(coordinates).unwrap() += value;
    }
}

Playground link