Idiomatic Rust index filtering (numpy style)-CodePudding

A common idiom that I use in python for initializing arrays is

arr = np.zeros(10)
x = np.arange(-5,5,10)
arr[np.abs(x)<2] = 1.

That is, conditionally changing elements of an array using a "view". Is there an idiomatic way of doing this (also with a vector or other collection) in Rust? Or is the only option to do this:

let mut arr = [0.; 10];
let x = linspace::<f64>[-5., 5., 10];
let mut i = 0;
for ele in x
{
    if (ele < 2.) arr[i] = 1.;
    i  = 1;
}

I've attempted to do something similar, which could be wrapped in a macro (haven't learned how to do that yet), but I'm not sure how to consume the iterator to actually execute it. Is this at least on the right track?

u.iter()
 .filter(|x: &&f64| x.abs()<1./3.)
 .map(|x: &f64| 1.);

CodePudding user response：

You are looking for something to the tune of

fn main() {
    let mut arr = [0.; 10];
    arr.iter_mut()
        .filter(|x: &&mut f64| x.abs() < 1. / 3.)
        .for_each(|x: &mut f64| *x = 1.);
}

We use .iter_mut() to create an iterator over mutable references; we need .iter_mut() because we are about the change/mutate the elements that the iterator is based on. The items the iterator yields will be &mut f64, mutable references to float64s.

We then call .filter() on that iterator, which creates a new iterator that only yields those elements that pass the given closure. The closure we give as an argument to .filter() must take immutable references to the items in the iterator, as the closure is only allowed to inspect the elements but not change them; notice that the signature for .filter() says that the closure needs to be a FnMut(&Self::Item) -> bool (notice the &); this where the "extra" reference (&&) comes from. So the type of the parameter x in the closure passed to filter() is an "immutable reference to a mutable reference to a f64" (a &&mut f64).

Finally, we call for_each() on that filtered iterator, which will consume the iterator it is called on and executes the given closure on each element. As that closure consumes the elements, it doesn't just get references (&&mut f64) but the actual &mut f64.

The (example) closure simply assigns a new value via *x = 1.. We need one dereference step (*) because don't want to assign to the reference (changing what the &mut f64 points at), but assign to the value the reference points to (changing the value behind the &mut f64).

CodePudding user response：

A direct translation of your example would be:

fn main() {
    let mut arr = [0_f64; 10]; //arr = np.zeros(10)
    let x = (-5..5); // x = np.arange(-5,5,10)
    for x in (-1..=1) {
        arr[(arr.len() as i32   x) as usize % arr.len()] = 1.0 // arr[np.abs(x)<2] = 1.
    }
    println!("{arr:?}"); // [1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
}