user2770149 - 1 year ago 108
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])
``````

with

``````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]])
``````

with

``````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 `b`to `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.
``````A*b[:,None]