How to create Rust ndarray with custom explicit indices

help
1

I need help creating Rust ndarrays with custom begin and end indices. This question is based on my StackOverflow question. Unfortunately, I don't receive answers to it on there.

In short, I'm writing a stencil-based finite difference method. I would like to calculate boundary conditions. To do that I need two planes of field values right next to the boundary. For example, my electric field array for the entire 3D space is:

use ndarray::Array3;
...
const NX: usize = 100;
const NY: usize = 100;
const NZ: usize = 100;

let mut ex = Array3::<f32>::zeros((NX, NY, NZ));

This will create an array with indices running from 0 to 99. For the boundary calculation I need two 2D slices of this array (in xy-plane) at coordinates z=NZ-1 and z=NZ. Ideally, I would need Rust support for the following syntax:

// forgive my use of Fortran-inspired syntax
let mut ex_kmax = Array3::<f32>::zeros((NX, NY, NZ-1:NZ));
...
// forgive my use of Fortran-inspired syntax
ex_kmax[[:,:,NZ-1:NZ]] = ex[[:,:,NZ-1:NZ]];

This new array is not a reference. It needs to contain a copy of the original ex values. Is that possible in Rust? How can I achieve that in a native Rust way?

2

Have a look at the different slicing methods. For example, try this:

let mut ex_kmax = ex.slice(s![.., .., NZ-1..NZ]).clone()

Then, ex_kmax has shape (100, 100, 1).

slice returns a view of the data, but you can clone with clone if you need to.

2 Likes
3

I still don't get why you need custom begin index. Care to provide some concrete use case?

1 Like
4

I think this method invocation must be replaced with either ArrayBase::to_owned or ArrayBase::into_owned. If you clone the view you only create a clone of the view, not the underlying data.

use ndarray::{arr3, s, Array3};

fn main() {
    let ex = arr3(&[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]);

    let ex_kmax: Array3<u8> = ex.slice(s!(.., .., -2..)).into_owned();

    assert_eq!(ex_kmax, arr3(&[[[2, 3], [5, 6]], [[8, 9], [11, 12]]]));
}

Playground.

1 Like
5

The use case is in the example above. It is a 'thin' 3D array ex_kmax. It is composed of two 2D xy-planes of the original array ex. These would be ex[[:,:,NZ-1]] and ex[[:,:,NZ]]. Stacked together they would make this 'thin' array with the coordinates ex_kmax[[:,:,NZ-1:NZ]].

What is this new array and why do I need it? It is a copy, a recording of the original array values at the previous time step. A brute, inelegant way to save the values at the previous time step would be a copy of the entire ex array. But I don't need all the ex values. Saving all of them would be a waste of memory. I just need those two xy-planes of values. I hope this makes it clearer.

6

Just regular slicing like @jofas provided will do, right? The "custom begin index" part seems unnecessary.

2 Likes
7

Can you please tell me, what would be the indices in this ex_kmax array? Would they still start from zero? It would be very convenient for me if the new array indices were (.., .., 1..). This is the main aim of my question. If the new array indices still start from zero, I need to do a lot of manual labour and counting...

8

ex_kmax is zero-indexed. Can you elaborate what kind of manual labour is required? Isn't an offset variable enough and you index ex_kmax with the offset and ex with NZ minus the offset?

1 Like
9

Here is a specific example:

    for i in IMIN..=IMAX {
        for j in JMIN..=JMAX {

            // Ex @ kmin
            ex[[i,j,KMIN-1]] =
                calc_mur1_pt::<IMIN, IMAX, JMIN, JMAX, KMIN, KMAX>
                    (&ex_kmin_m[[i,j,1]], &ex[[i,j,KMIN]],
                     &ex_kmin_m[[i,j,0]],  c0, dt, dx);

            // Ex @ kmax
            ex[[i,j,KMAX+1]] =
                calc_mur1_pt::<IMIN, IMAX, JMIN, JMAX, KMIN, KMAX>
                    (&ex_kmax_m[[i,j,0]], &ex[[i,j,KMAX]],
                     &ex_kmax_m[[i,j,1]],  c0, dt, dx);
        }
    }

Here I'm using the zero based indices of the ex_kmin and ex_kmax arrays. If I had access to custom begin indices this example would turn to:

    for i in IMIN..=IMAX {
        for j in JMIN..=JMAX {

            // Ex @ kmin
            ex[[i,j,KMIN-1]] =
                calc_mur1_pt::<IMIN, IMAX, JMIN, JMAX, KMIN, KMAX>
                    (&ex_kmin_m[[i,j,KMIN]],   &ex[[i,j,KMIN]],
                     &ex_kmin_m[[i,j,KMIN-1]],  c0, dt, dx);

            // Ex @ kmax
            ex[[i,j,KMAX+1]] =
                calc_mur1_pt::<IMIN, IMAX, JMIN, JMAX, KMIN, KMAX>
                    (&ex_kmax_m[[i,j,KMAX]],   &ex[[i,j,KMAX]],
                     &ex_kmax_m[[i,j,KMAX+1]],  c0, dt, dx);
        }
    }

This is much more intuitive and can be simplified further.

10

I have a strong feeling, that my loops can be also optimised away. But this is a topic of a new question. I'm a Rust novice (;

11

I dont get why the second case is more intuitive :smiley:

I think your usecase can be simplified by using the s! and azip! macros of the crate but you have to embrace zero based indexing and exclusive ranges. What is your reason to stick to 1 based indexing? Are you translating fortran code?

1 Like
12

Because kmin-1 and kmin can be used in place of 0 and 1. This way I can generalise my code and only keep track of kmin-1, kmin, not both kmin-1, kmin and zeros and ones.

13

One based indexing is not very important here. Custom begin indices are. There are two key parameters in my arrays: the effective region from 1..NX and the boundary points 0 and NX+1. Boundary points are special cases. To update the array values at the boundaries I need the values one point before the boundary. In case of kmin these would be the values at locations kmin-1 and kmin itself. In case of kmax, these would be the values at spatial points kmax and kmax+1. I hope this makes it clearer.

No, I'm not translating a Fortran code. I just have Fortran experience.

14

This is how I allocate and initialise these 'boundary plane' arrays at the moment:

    let mut ex_kmin_m = Array3::<f32>::zeros((IMAX+2,JMAX+2,2));
    let mut ex_kmax_m = Array3::<f32>::zeros((IMAX+2,JMAX+2,2));

Here I'm using the default zero-based indexing. Ideally, I would like to do:

    let mut ex_kmin_m = Array3::<f32>::zeros((IMAX+2,JMAX+2,KMIN-1..KMIN));
    let mut ex_kmax_m = Array3::<f32>::zeros((IMAX+2,JMAX+2,KMAX..KMAX+1));
15

The intuitive thing for me would be to slice the array in the first step and work with 0..idx_len and the "offset" IDX_MIN for the rest of the calculation. Something like

let ex_view = ex.slice_mut(s![I_START..I_STOP, J_START..J_STOP, K_START..K_STOP]);
// alternatively use split 
let mut ex_view_iter = ex_view.axis_iter_mut(Axis(2));
let ex_view_0 = ex_view_iter.next().unwrap();
let ex_view_1 = ex_view_iter.next().unwrap();

// or par_azip if you want multi threading
azip!(
    (a0 in &mut ex_view_0,
     a1 in &ex_view_1,
     b0 in &ex_kmin_m_0,
     b1 in &ex_kmin_m_1)
*a0 = calc_mur1_pt::<...>(b1, a1, b0, ...));

The names and the structure are not very nice but I dont know your exact problem. The important part is that we split of the mutably accessed part, which allows you to use rayon/multi-threading in a very simple way. I guess you are useing fortran memory order? The parallel performance should benefit if the k axis is not continuous.

2 Likes
16

Thank you! I need to read the documentation of slice_mut(), axis_iter_mut() and azip!(). Now your example looks very cryptic for me.

What multidimensional array order does Rust use? Is it the row-major C-like order? Yes, ideally I would like to use multithreading for future calculation.

17

The default layout is row major but you can change it to column major.

1 Like
18

Just a reminder that ndarray is not equal to rust. Rust itself doesn't have special syntax support for multidimensional arrays. In pure rust a 2d array looks like this

const ROW = 3;
const COL = 4;

// "Row major"
let a2dr= [[0; ROW]; COL];

// "Column major"
let a2dc = [[0; COL]; ROW];

// Set a specific element
let (row, col) = (1, 2);
a2dr[col][row] = 7;
a2dc[row][col] = 7;


// Same thing in ndarray
use ndarray::Array2;
let mut array = Array2::zeros((ROW, COL)); // These don't have to be constants.
array[[row, col]] = 7;

As you can see row or column major in pure rust is just a matter of how you define the array. Also these are compile time constant and live on the stack. Both things you don't want in ndarray or in general for large arrays.

2 Likes