user2770149 user2770149 - 2 months ago 13
Python Question

broadcast / extend boolean array along axis in python

I am trying to broadcast a boolean 1D array along a vector to isolate certain operations.

For example, if my boolean array is:

b = array([True False False True])


np.shape(b) = (4,)

My matrix I would like to do row wise comparison with is of this form:

A = array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0]])


np.shape(A) = (4, 3)

The end result I would like would be:

C = b * A (row-wise comparison)

C = array([[1.0, 2.0, 3.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [10.0, 11.0, 12.0]])

I realize that this can be done by slicing, but I need the speed boost from using logical arrays. Is there a way to "extend" the boolean vector so that it works over the rows?

The idea is to be able to calculate a function all in one line so that

D = b * A * (some operation done over all vectors) + ~b * A * (some different operation done over all vectors)

Thank you!


Extend bto 2D keeping the second axis as a singleton dimension (dimension with length = 1) and multiply with A. Thus, under the hood b would be broadcasted to the same shape as of A and element-wise multiplications would be performed.

Thus, the vectorized solution would be -