rustrust-ndarray

Macro to slice ndarray


I have a struct with one Array and other properties associated to it. I want to be able to slice the array in order to write values at specific indices. However, I don't want the user to have to know the field name where the array is stored, so just calling a slice method on their object should work.

More specifically:

struct Foo {
    data: Array2<u32>,
    v1: i32,
    v2: i32
}

fn main() {
    // Just for showing purpose, it's not like this in the main program
    bar = Foo {
        data: Array2::zeros((4, 3));
        v1: // Some value,
        v2: // Another value
    }
}

Then, I want something like slice!(bar, [0, 1]) == bar.data[[0, 1]]

I think a macro is what I need, but I have no idea if it's the best way.


Solution

  • Most likely you just want to implement and forward indexing for your type:

    use ndarray::{Array2, NdIndex, Dim};
    use std::ops::{Index, IndexMut};
    
    impl<Idx: NdIndex<Dim<[usize; 2]>>> Index<Idx> for Foo {
        type Output = <Array2<u32> as Index<Idx>>::Output;
        fn index(&self, i: Idx) -> &Self::Output {
            &self.data[i]
        }
    }
    
    impl<Idx: NdIndex<Dim<[usize; 2]>>> IndexMut<Idx> for Foo {
        fn index_mut (&mut self, i: Idx) -> &mut Self::Output {
            &mut self.data[i]
        }
    }
    

    After that you can just use indexing on your type directly:

    bar[[0, 1]] == bar.data[[0, 1]]