Reputation: 115
I'm new to NumPy, and I've encountered a problem with running some conditional statements on numpy arrays. Let's say I have 3 numpy arrays that look like this:
a:
[[0, 4, 4, 2],
[1, 3, 0, 2],
[3, 2, 4, 4]]
b:
[[6, 9, 8, 6],
[7, 7, 9, 6],
[8, 6, 5, 7]]
and, c:
[[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]]
I have a conditional statement for a and b in which I would like to use the value of b (if the conditions of a and b are met) to calculate the value of c:
c[(a > 3) & (b > 8)]+=b*2
I get an error saying:
Traceback (most recent call last):
File "<interactive input>", line 1, in <module>
ValueError: non-broadcastable output operand with shape (1,) doesn't match the broadcast shape (3,4)
Any idea how I can accomplish this?
I would like the output of c to look as follows:
[[0, 18, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]]
Upvotes: 11
Views: 39063
Reputation: 477794
The problem is that you mask the receiving part, but do not mask the sender part. As a result:
c[(a > 3) & (b > 8)]+=b*2
# ^ 1x1 matrix ^3x4 matrix
The dimensions are not the same. Given you want to perform element-wise addition (based on your example), you can simply add the slicing to the right part as well:
c[(a > 3) & (b > 8)]+=b[(a > 3) & (b > 8)]*2
or make it more efficient:
mask = (a > 3) & (b > 8)
c[mask] += b[mask]*2
Upvotes: 10
Reputation: 22964
A slight change in the numpy expression would get the desired results:
c += ((a > 3) & (b > 8)) * b*2
Here First I create a mask matrix with boolean values, from ((a > 3) & (b > 8))
, then multiply the matrix with b*2
which in turn generates a 3x4
matrix which can be easily added to c
Upvotes: 3
Reputation: 215127
You can use numpy.where
:
np.where((a > 3) & (b > 8), c + b*2, c)
#array([[ 0, 18, 0, 0],
# [ 0, 0, 0, 0],
# [ 0, 0, 0, 0]])
Or arithmetically:
c + b*2 * ((a > 3) & (b > 8))
#array([[ 0, 18, 0, 0],
# [ 0, 0, 0, 0],
# [ 0, 0, 0, 0]])
Upvotes: 22