user1717931
user1717931

Reputation: 2501

numerical operation on two numpy ndarrays

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

Answers (2)

Kevin
Kevin

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

Garrett R
Garrett R

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

Related Questions