Reputation: 2501
I have two numpy 2-D arrays - A, B that have the same dimensions (n x n). I want to look at each cell on the both arrays and do some conditional arithmetic and return a new 2_D array (C).
For example, if A[i, j] = 0.0
and B[i, j]= 0.0
, then, I want to populate C[i, j] = 1.0
. If the cell values are different, then, I have to populate C[i, j] = f(A[i, j], B[i, j])
. There may be a few more if-conditions.
I can do brute-force looping, but, it is neither Pythonic nor an elegant way. I am reading about vectorization/from-functions etc,...but, unclear how to do it.
Any suggestions? appreciate it.
Thanks.
Upvotes: 1
Views: 49
Reputation: 8207
Here's a way you can do it by calculating f
for the entire matrix, then go back and assign the value 1.0 to the == indices from a and b:
a= np.random.randn(3, 3)
b= np.random.randn(3, 3)
c= np.zeros((3,3))
a[0,0]= 0 # for your example
b[0,0]= 0
In [532]: a
Out[532]:
array([[ 0. , 0.41134107, -0.58869891],
[ 0.09898662, -1.79136973, 1.28057148],
[ 0.14666895, -0.7022536 , 0.07162134]])
In [533]: b
Out[533]:
array([[ 0. , -0.19669429, -0.82305814],
[ 1.45526312, 1.7483546 , 0.41721631],
[ 0.50303962, -0.03359472, 0.09301669]])
Define your function:
def f(x,y):
return ((x+1.0) + (y+1.0))
apply it to the entire matrix:
c = f(a,b)
Now go back and assign the value of 1.0 where a==b
:
c[a==b] = 1.0
Results:
In [538]: c
Out[538]:
array([[ 1. , 2.21464677, 0.58824296],
[ 3.55424975, 1.95698487, 3.69778778],
[ 2.64970858, 1.26415167, 2.16463802]])
Upvotes: 2
Reputation: 2662
I'm not sure exactly what this is called, but I often refer to it as magic indexing. Given two matrices A and B, where they are entirely equal except for their last column and the upper most left value (index (0,0)). I first find get the indices where their values are equal use the ==
test. This returns a boolean matrix which can be used. Then I preinitialize the result matrix C.
#initialize two identical 2d arrays
A = np.array([np.arange(10) for i in range(5)])
B = np.array([np.arange(10) for i in range(5)])
A[0,0] = -1
A[:,-1] = range(10,15)
B[:,-1] = 1
#now everything but the last column is the same and upper left corner
equal_mat = A == B
print equal_mat #boolean matrix
def someFun(A, B):
toRet = A*B +1
return toRet
C = np.zeros(A.shape) #initialize result matrix
C[equal_mat] = 0.0 #set result to zero wherever A == B
opp = np.logical_not(equal_mat) # get the opposite indices
C[opp] = someFun(A[opp], B[opp]) # assign the output of a function
print C #zero everywhere except index (0,0) and the last column
Note that someFun must take in two arrays of shape (m,) and return a single array of shape (m,). Hope that helps.
Upvotes: 0