Split multiple arrays to multiple mutable chunks

[mrsanor]

Hey, I haven't figured out if I can split arrays into multiple mutable chunks somehow.

Example of what I would like to achieve:

// All the data in same index is related to each other.
// E.g., data1[0], data2[0] and data3.row(0) should all be handled at the same time.
let mut data1 = Array1::zeros(100000);
let mut data2 = Array1::zeros(100000);
let mut data3 = Array2::zeros((100000, 2));

// Split all the data to X chunks (X is known only at runtime).
// All the chunks are guaranteed not to be overlapping, so no need to worry about race conditions.
thread::scope(|scope| {
  for (start_index, end_index) in get_chunks() {
    let mut data1 = data1.slice_mut(s![start_index..end_index]);
    let mut data2 = data2.slice_mut(s![start_index..end_index]);
    let mut data3 = data3.slice_mut(s![start_index..end_index, ..]);
    scope.spawn(move || {
      // Do something with the data
    });
  }
});

multi_slice_mut won't work since the amount of chunks is known only at runtime.
axis_chunks_iter_mut won't work because zipped data can't be iterated with parallel iterator. Also you can't zip together data with different dimensionalities.
slice_mut won't work because you of course can't borrow as mutable many times. You can use into_owned() but that clones the data and there's lots of data to be processed to this should be avoided.

[jturner314]

Will you please clarify why axis_chunks_iter_mut doesn't meet your needs?

axis_chunks_iter_mut won't work because zipped data can't be iterated with parallel iterator.

Will you please clarify what you mean by this? When manually creating threads, array views produced by AxisChunksIterMut can be moved into different threads. Personally, though, I'd suggest using rayon for convenient parallelization, unless it doesn't meet your requirements for some reason. ndarray has built-in integration with rayon if you enable the rayon feature. For example, AxisChunksIterMut implements IntoParallelIterator. Here's some example code using rayon:

use ndarray::prelude::*;
use rayon::prelude::*;

fn main() {
    let mut data1 = Array1::<i32>::zeros(100000);
    let mut data2 = Array1::<i32>::zeros(100000);
    let mut data3 = Array2::<i32>::zeros((100000, 2));
    let chunk_size = 1000;
    data1
        .axis_chunks_iter_mut(Axis(0), chunk_size)
        .into_par_iter()
        .zip_eq(data2.axis_chunks_iter_mut(Axis(0), chunk_size))
        .zip_eq(data3.axis_chunks_iter_mut(Axis(0), chunk_size))
        .for_each(
            |((d1, d2), mut d3): (
                (ArrayViewMut1<'_, i32>, ArrayViewMut1<'_, i32>),
                ArrayViewMut2<'_, i32>,
            )| {
                d3.column_mut(0).assign(&d1);
                d3.column_mut(1).assign(&d2);
            },
        );
}

Also you can't zip together data with different dimensionalities.

Will you please clarify what you mean by this? You can zip together iterators regardless of the types of the iterators' items. If you're referring to iteration over multiple axes, you can use broadcasting with ndarray::Zip::and_broadcast.

[mrsanor]

Will you please clarify why axis_chunks_iter_mut doesn't meet your needs?

I didn't realize you need to first import rayon for parallel iterator zipping to work, my bad.

You can zip together iterators regardless of the types of the iterators' items.

Ah yeah sorry, I actually meant to write that you can't zip together immutable and mutable data, rather than different dimensionalities.

But anyway, I got this to work with axis_chunks_iter and axis_chunks_iter_mut (data1 was supposed to be immutable in my example) and manually consuming the iterators. I guess I was a bit sleep deprived when I posted the topic, the answer was right there in front of me 😞

Thank you for the quick help and awesome job with the library!

Here's the final solution in case anyone else in the future needs it:

let data1 = Array1::zeros(100000);
let mut data2 = Array1::zeros(100000);
let mut data3 = Array2::zeros((100000, 2));

let n_chunks = 20;
let chunk_size = 100000 / n_chunks;

let mut data1_iter = data1.axis_chunks_iter(Axis(0), chunk_size);
let mut data2_iter = data2.axis_chunks_iter_mut(Axis(0), chunk_size);
let mut data3_iter = data3.axis_chunks_iter_mut(Axis(0), chunk_size);

thread::scope(|scope| {
  for _ in 0..n_chunks {
    let d1 = data1_iter.next().unwrap();
    let mut d2 = data1_iter.next().unwrap();
    let mut d3 = data1_iter.next().unwrap();
    scope.spawn(move || {
      // Do something with the data
    });
  }
}

[jturner314]

Ah yeah sorry, I actually meant to write that you can't zip together immutable and mutable data, rather than different dimensionalities.

Fwiw, that's not true, either. Maybe if you have some example code which didn't compile as expected, I could help resolve the errors? Anyway, here's an example using zipping and manual threading. (Note that the rayon crate feature is not necessary for this example; the rayon feature is necessary only if you want to use rayon's traits, etc.)

use ndarray::prelude::*;
use std::thread;

fn main() {
    let data1 = Array1::<i32>::zeros(100000);
    let mut data2 = Array1::<i32>::zeros(100000);
    let mut data3 = Array2::<i32>::zeros((100000, 2));

    let n_chunks = 20;
    let chunk_size = 100000 / n_chunks;

    thread::scope(|scope| {
        for ((d1, mut d2), mut d3) in data1
            .axis_chunks_iter(Axis(0), chunk_size)
            .zip(data2.axis_chunks_iter_mut(Axis(0), chunk_size))
            .zip(data3.axis_chunks_iter_mut(Axis(0), chunk_size))
        {
            scope.spawn(move || {
                // Do something with the data, e.g.:
                d3.column_mut(0).assign(&d1);
                d3.column_mut(1).assign(&d2);
            });
        }
    });
}

(Playground)

[mrsanor]

Don't know what I messed up with zipping, when I thought it didn't work with mutable and immutable data. But yeah I see it actually does work, thanks for the help and sorry for all the confusion! 😄