iter map(v…) in rust - efficient recalculation of ndarray elements into two arrays at once?

Question

I’d like to iter over array elements and modify them. Additionally I would like (for efficiency purposes) compute at once two new results. The pseudocode in rust is below:

use ndarray::*; //definition of Array1

// let's assume examplary input/output
let x = Array1::random(100, Uniform::<f64>::new(0., 1.));
let y = Array1::random(100, Uniform::<f64>::new(0., 1.));

// transformation of the x array into two arrays (a,b) of the same size
let (a, b): (Array1<f64>, Array1<f64>) = x
    .iter()
    .map(|v| {
        // ...  computing
        let result1 = v; // just dummy example
        let result2 = v*2.0; // just dummy example
        return (result1, result2);
    })
    .collect(); // or unzip()

// after this I’d like to compute something on the result
let res1: f64 = a.dot(&y);
let res2: f64 = b.dot(&y);

It can be done when I map into vec! , but i would not like to make additional conversions/allocations if possible as I use later the result for further calculations (as dot product in the above code)
If I divide calculations in the closure into two parts and try to use for example mapv_inplace, I will have to compute a lot ot the same things again - I do not want to do it as its very expensive.

How to efficiently resolve such issue?

score 0 · Answer 1 · answered May 25 '23 at 12:04

You can use Zip for this. You just need to set up the output arrays a and b yourself.

let mut a = Array1::zeros(100);
let mut b = Array1::zeros(100);

// transformation of the x array into two arrays (a,b) of the same size
Zip::from(&x).and(&mut a).and(&mut b)
    .for_each(|&v, a, b| {
        // ...  computing
        *a = v; // just dummy example
        *b = v*2.0; // just dummy example
    });

(Try it in the playground.)

iter map(v…) in rust - efficient recalculation of ndarray elements into two arrays at once?

1 Answers1