Numpy: broadcasting + boolean indexing
-
29-06-2021 - |
Question
I have the following numpy arrays
A: shape (n1, n2) array of float
B: shape (n2,) array of float
M: shape (n1, n2) array of bool
How do I turn the following pseduo-code inte efficient real code? The arrays may be huge, possibly > 100 million elements.
A[M] = ("B broadcast to shape (n1,n2)")[M]
Solution
Broadcasting is simple and memory efficient:
A, B, M = np.broadcast_arrays(A, B, M)
However using this B
in your code A[M] = B[M]
would not be memory efficient because B[M]
has as many real elements as M
has True
values.
Instead use:
np.putmask(A, M, B)
Since B
is repeated automatically with the putmask
function, you should not even have to broadcast it. Though I guess it cannot hurt to do that.
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow