How to obtain and use boolean masks of ndarrays?

[rkshthrmsh]

I would like to implement the following function, where sample is a 2D array and offset is a 1D array.

  \begin{array}{ c l }
    (sample + offset)  & \quad \textrm{if } sample[i, j] \geq 1 \\
    0                            & \quad \textrm{otherwise}
  \end{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?

[nilgoyette]

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

[rkshthrmsh]

The solution with map_collect seems closest to what I was looking for. Thank you, @nilgoyette! :)