CODEWITHSUNDEEP
pythonnumpy
Jojo
Jojo

Reputation: 1117

Vectorized implementation of conditional numpy multiplication

I haven't used numpy much and am struggling to figure out the vectorized implementation of the following:

Change and pct_change are numpy arrays which are 5*2000 in size.

I wanted to calculate each element in the array Change_after_calculation like:

###Carry out this calculation for each number 
###in Change_after_calculation where Change has a correspondingly -ve number
Change_after_calculation = (numpy.ones((5,2000)) + pct_change) * Change 

###If Change has a correspondingly +ve number do the below
Change_after_calculation = Change

I could do the above with multiple for loops like this, but if there is a vectorized way to do it, I'd prefer to do that:

for p in 2000:
    for n in 5:
       if Change[n,p] < 0:
            Change_after_calculation[n,p] = (1+pct_change[n,p]) *  Change[n,p]
       else:
            Change_after_calculation[n,p] = Change[n,p]

Upvotes: 2

Views: 316

Answers (3)

Dmitry Petrov
Dmitry Petrov

Reputation: 120

Your code is not very reproducible (it doesn't run, so I constructed my own snippet based on what you wrote).

If I understood correctly, you just need to create a mask based on your condition (change[n,p] < 0) and do an assignment for the arrays based on this mask.

# defining arrays for testing

P, N = 2000, 5
change = np.random.randint(-1, high=2, size=(N,P))
change_after_calculation1 = np.zeros((N,P))
change_after_calculation2 = np.zeros((N,P))
pct_change = np.random.random((N,P))


# your code
for p in range(P):
    for n in range(N):
        if change[n,p] < 0:
            change_after_calculation1[n,p] = (1+pct_change[n,p]) *  change[n,p]
        else:
            change_after_calculation1[n,p] = change[n,p]
            
# vectorized version
mask = change < 0
change_after_calculation2[mask] = (1 + pct_change[mask]) * change[mask]
change_after_calculation2[~mask] = change[~mask]

# test that arrays match
print(np.allclose(change_after_calculation1, change_after_calculation2))

Upvotes: 0

Akshay Sehgal
Akshay Sehgal

Reputation: 19332

Try this oneliner without any need for defining another array from the start to hold the result -

(Change>=0)*Change + (Change<0)*(pct_change+1)*Change

On a toy example -

Change = np.random.randint(-2,4,size=(2,5))
pct_change = np.random.randint(1,3,size=(2,5))
Change_after_calculation = np.zeros_like(pct_change)

for p in range(5):
    for n in range(2):
        if Change[n,p] < 0:
            Change_after_calculation[n,p] = (1+pct_change[n,p]) *  Change[n,p]
        else:
            Change_after_calculation[n,p] = Change[n,p]

#OUTPUT - 
array([[ 0, -1,  0,  3,  3],
       [ 2, -1,  3,  2, -2]])


#using one-liner
Change_after_calculation = (Change>=0)*Change + (Change<0)*(pct_change+1)*Change

#OUTPUT - 
array([[ 0, -2,  0,  3,  3],
       [ 2, -3,  3,  2, -4]])

Upvotes: 3

David Hoffman
David Hoffman

Reputation: 2343

You can use NumPy's fancy indexing for this

# calculate a mask array for fancy indexing
less_than = Change < 0
# Do conditional math
Change_after_calculation[less_than] = (1 + pct_change[less_than]) *  Change[less_than]
Change_after_calculation[~less_than] = Change[~less_than]

If Change_after_calculation isn't initialized beforehand you can initialize it as

Change_after_calculation = np.empty_like(Change)

Upvotes: 5

Related Questions