# Efficient numpy value assignment via boolean mask

I have a boolean mask value assigning problem the requires efficient boolean mask operation.

It’s a multi-dimension mask and i’m using `einsum` to achieve the result, but the operation is not very efficient, and i’m wondering, if i can get some help with it
Here is my current solution: (both `mask, truth_value, false_value` are dummy data with dtype and shape matches to my problem.

``````mask = np.random.randn(1000, 50)> 0.5
truth_value = np.random.randn(50, 10)
false_value = np.random.randn(10)
``````

Is there any faster way to get `objective` given `mask, truth_value, false_value` ?

While i was waiting, figured out a faster way

``````objective = np.where(mask[...,np.newaxis], np.broadcast_to(truth_value, (1000, 50, 10)), np.broadcast_to(false_value,  (1000, 50, 10)))
``````

But is there any faster alternative ?

### >Solution :

You can use the Numba JIT to do that more efficiently.

``````import numpy as np
import numba as nb

@nb.njit('float64[:,:,::1](bool_[:,::1], float64[:,::1], float64[::1])')
l = false_value.shape[0]
assert truth_value.shape == (m, l)
result = np.empty((n, m, l), dtype=np.float64)
for i in range(n):
for j in range(m):
result[i, j, :] = truth_value[j, :]
else:
result[i, j, :] = false_value[:]
return result

mask = np.random.randn(1000, 50) > 0.5
truth_value = np.random.randn(50, 10)
false_value = np.random.randn(10)
``````

The computation of `objective` is 4.8 times faster on my machine.

If this is not fast enough, you can try to parallelize the code using the parameter `parallel=True` and using `nb.prange` instead of `range` in the i-based loop. This may not be faster due to the overhead of creating new threads. On my machine (with 6 cores), the parallel version is 7.4 times faster (the creation of threads is pretty expensive compared to the execution time).

Another possible optimization is to write directly the result in a buffer allocated ahead of time (this is only better if you call this function multiple times with the same array size).

Here are the overall timings on my machine:

``````np.einsum:         4.32 ms
np.where:          1.72 ms
numba sequential:  0.89 ms
numba parallel:    0.58 ms
``````