Unzip array of tuples

[niklasmueboe]

I have a function that I am mapping along an axis of an array to reduce it, but it doesn't return a single value but rather a tuple.

let tuple_values = arr.map_axis(Axis(0), reduce_to_tuple);

To get the values as separate arrays I am currently doing this
let val1 = tuple_values.mapv(|(x, _)| x);
let val2 = tuple_values.mapv(|(_, x)| x);

I was wondering whether there is a simpler way to unzip an array of tuples, e.g. along the lines of

let (val1, val2) = arr.map_axis(Axis(0), reduce_to_tuple).unzip();

ALos it isn't necessarily a 2-tuple but could be e.g. a 3-tuple, ...

[nilgoyette]

AFAIK, there's nothing in ndarray that you can call to do what you want. However, if you have a contiguous array, you can get the slice (as_slice, as_slice_memory_order) and use itertools.

[akern40]

This can be accomplished through the use of Zip:

use ndarray::{Array, Array1, ArrayView, Dimension, Zip};
use ndarray_rand::{rand_distr::Uniform, RandomExt};
use num_traits::Zero;

fn unzip<A, B, D>(arr: &ArrayView<(A, B), D>) -> (Array<A, D>, Array<B, D>)
where
    D: Dimension,
    A: Copy + Zero,
    B: Copy + Zero,
{
    let mut arr_a = Array::<A, D>::zeros(arr.raw_dim());
    let mut arr_b = Array::<B, D>::zeros(arr.raw_dim());
    Zip::from(arr)
        .and(&mut arr_a)
        .and(&mut arr_b)
        .for_each(|orig, a, b| {
            *a = orig.0;
            *b = orig.1;
        });

    (arr_a, arr_b)
}

fn main() {
    let arr = Array1::<i32>::random(10, Uniform::new(0, 20));
    let (div, rem) = unzip(&arr.map(|v| (v / 2, v % 2)).view());
    println!("Original: {arr}");
    println!("Divided: {div}");
    println!("Remainder: {rem}");
}

Unfortunately, doing this in the way itertools does it (i.e., the same function for multiple different arities of tuples) requires some more machinery and - as itertools does - probably a macro to implement them all.

On a separate note, I might welcome a sort of "itertools for ndarray" to make these things more coherent within the Rust ecosystem. But that's a huge task.