How to obtain and use boolean masks of ndarrays?
rkshthrmsh opened this issue · comments
I would like to implement the following function, where sample
is a 2D array and offset
is a 1D array.
In numpy, it is possible to obtain a boolean mask and multiply it element-wise since Python's True and False can be interpreted numerically. Using this, the above function can be implemented as:
arr = np.multiply(np.add(sample, offset), sample >= 1) # here sample >= 1 creates the boolean mask
What would be the idiomatic way of doing this in Rust with ndarray?
Sorry, this won't be a one-liner :) My first version was
let mut arr = sample.clone();
for mut row in arr.rows_mut() {
Zip::from(row).and(offset).for_each(|r, &o| {
*r = if *r >= 1.0 { *r + o } else { 0.0 }
})
}
arr
but then I remembered and_broadcast
let mut arr = sample.clone();
Zip::from(&mut arr).and_broadcast(&offset).for_each(|a, &o| {
*a = if *a >= 1 { *a + o } else { 0 }
});
arr
but then I remembered map_collect
Zip::from(&sample).and_broadcast(&offset).map_collect(|&s, &o| {
if s >= 1 { s + o } else { 0 }
})
I think this is not too bad because:
- it allocates only what it should
- it's not too long. It's a one-liner, at least in my heart ;)
- it's shorter than what I thought it would be
The solution with map_collect seems closest to what I was looking for. Thank you, @nilgoyette! :)