Modify an existing array and assign to another one of same shape.

[daniellga]

Is the following code the most efficient way of modifying an existing ndarray and assigning to another pre-allocated one of the same shape but different type?

    let arr = Array2::from_shape_vec((2,2), vec![1_f32,2.,3.,4.]).unwrap();
    let dim = arr.dim();
    let mut new_arr = Array2::<Complex<f32>>::zeros(dim);

    new_arr.zip_mut_with(&arr, |val_new, val| *val_new = Complex {re: *val, im: 0.});

I was also checking zip_mut_with source code. Why does it check self.dim.ndim() == rhs.dim.ndim()? Isn't the self.shape() == rhs.shape() check enough?

    pub fn zip_mut_with<B, S2, E, F>(&mut self, rhs: &ArrayBase<S2, E>, f: F)
    where
        S: DataMut,
        S2: Data<Elem = B>,
        E: Dimension,
        F: FnMut(&mut A, &B),
    {
        if rhs.dim.ndim() == 0 {
            // Skip broadcast from 0-dim array
            self.zip_mut_with_elem(rhs.get_0d(), f);
        } else if self.dim.ndim() == rhs.dim.ndim() && self.shape() == rhs.shape() {
            self.zip_mut_with_same_shape(rhs, f);
        } else {
            let rhs_broadcast = rhs.broadcast_unwrap(self.raw_dim());
            self.zip_mut_with_by_rows(&rhs_broadcast, f);
        }
    }

Thanks!

[nilgoyette]

I would probably use

let new_arr = arr.mapv(|a| Complex::new(a, 0.0));

because it's short and clear. Is it the fastest method? I don't know. You will have to test. Maybe

let new_arr = Zip::from(&arr).apply_collect(|&a| Complex::new(a, 0.0));

is faster (but uglier). Create a benchmark and test it.

As for your second question. No. Arrays of shapes &[3, 2] (ndim = 2) and &[2] (ndim = 1) could be used with zip_mut_with.

[daniellga]

In my use case I preallocate and cache the allocation, that's why I can't use your suggestions.

But in addition to the differente ndim, wouldn't shape of both Arrays also be different in this example you gave? Isn't it redundant?

[nilgoyette]

pre-allocated

Oh, right. I forgot about that. In that case, what you use is fine! It's probably as fast as

Zip::from(&arr).map_assign_into(new_arr, |&a| Complex::new(a, 0.0));

bluss answered your other question below.

[bluss]

Why does it check self.dim.ndim() == rhs.dim.ndim()?

For all Dimension impls (except IxDyn), ndim is a compile time known quantity. Thus this is like a compile time known conditional, which can be useful to have there for optimization.