Can I create a new array by comparing arrays with numpy?
I have 3 arrays (A1, A2, A3).
How can I find all indexes where A1 == 2 and A2 > A3 and write there the value 5 in a new array?
I have this matlab code which is doing this:
index = find(A1==2 & A2>A3);
new_array(index) = 5;
I found putmask and logical_and but not sure if this are the right tools and how do use it in my case. Thanks!
The code below uses & to chain together conditionals. As such whenever A1 == 2 and A2 > A3 are both True, the index array will be True
import numpy as np
A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3
new_array = np.zeros(len(A1))
index = (A1 == 2) & (A2 > A3)
new_array[index] = 5
# array([ 0., 0., 5., 5., 0., 0.])
You could use np.logical_and of course. But this restricts you to only two conditionals, whilst you can effectively chain as many as you like when using &.
You can use the np.where function
import numpy as np
A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3
out = np.where((A1 == 2) & (A2 >= A3), 5, 0)
or more simple
out = ((A1 == 2) & (A2 >= A3))*5
