Rewritning if condition to speed up in python

Question

I have following piece of code with if statement within a function. When I run it would take long time and it is a way to rewrite if condition or a way to speed up this sample code?

import numpy as np

def func(S, R, H):
    ST =  S * R
    if ST <= - H:
      result = - H
    elif ST >= - H and ST < 0:
      result = ST
    else:
      result = min(ST, H)
    return result

y=[]
t1= time()
for x in np.arange(0, 10000, 0.001): 
    y.append(func(3, 5, x))
t2 = time()
print("time with numpy arange:", t2-t1)

time taken to run the code:

 10 s

This is reproduced sample of real code, and in the real code ST becomes both negative and positive value, we may keep conditions but changing if statement to something else may help perfom task faster!

Since you are always passing 3 and 5 as first two arguments to `func`, why not make `ST` = 15? — Shiva, Jan 08 '21 at 10:54
Will you be running this function like this in the real code? Have you identified this particular function as a bottle-neck somewhere? — Moberg, Jan 08 '21 at 11:01
This function is reproduced sample code of real code. The algorithm must be the same, in real code ST become noth negative and postive value. — Pavel.D, Jan 08 '21 at 12:18

Patrick Artner · Accepted Answer · 2021-01-08T12:36:54.063

If you want your functions parameters still available you would need to use boolean indexing in a creative way and replace your function with that:

from time import time
import numpy as np

ran = np.arange(-10, 10, 1)
s = 2
r = 3

st =  s * r

def func(S, R, H):
    ST =  S * R
    if ST <= - H:
      result = - H
    elif ST >= - H and ST < 0:
      result = ST
    else:
      result = min(ST, H)
    return result

# calculate with function
a = []
t1 = time()
for x in ran:
    a.append(func(s, r, x))
t2 = time()
print("time with function:", t2 - t1)
a = np.array(a)

# calculate with numpy
y = np.copy(ran)
neg_y = np.copy(y) * -1

# creative boolean indexing
t1 = time()
y[st <= neg_y] = neg_y[st <= neg_y]
if st < 0:
  y[st >= neg_y] = st
else:
  alike = np.full(ran.shape, st)[st >= neg_y]
  y[st > neg_y] = np.where(y[st > neg_y] > st, st, y[st > neg_y])
t2 = time()

print(a)
print(y)
print("time with numpy indexing:", t2 - t1)

Will give you (timings omitted):

# s=2, r=3
[10  9  8  7  6  5  4  3  2  1  0 -1 -2 -3 -4 -5 -6 -6 -6 -6] # function
[10  9  8  7  6  5  4  3  2  1  0 -1 -2 -3 -4 -5 -6 -6 -6 -6] # numpy

# s=-2, s=3
[10  9  8  7  6 -5 -4 -3 -2 -1  0  1  2  3  4  5  6  6  6  6] # function
[10  9  8  7  6 -5 -4 -3 -2 -1  0  1  2  3  4  5  6  6  6  6] # numpy

You might need to tweak it a bit more.

Using a

ran = np.arange(-1000, 1000, 0.001)

I get timings (s=3,r=5) of:

time with function: 5.606577634811401
[1000.     999.999  999.998 ...   15.      15.      15.   ]
[1000.     999.999  999.998 ...   15.      15.      15.   ]
time with numpy indexing: 0.06600046157836914

Thanks to your solution, indeed with creative boolean indexing difficult to understand but it helps and gives back same result with quicker performance. As you mention it above I need to tweak it to the code. — Pavel.D, Jan 08 '21 at 13:54

Rewritning if condition to speed up in python

1 Answers1