CODEWITHSUNDEEP
pythonpython-2.7performancecythonnumba
hockey5692
hockey5692

Reputation: 25

Suggestions on how to speed up this python function?

Any suggestions on how to speed up this function? 

def smooth_surface(z,c):
    hph_arr_list = []
    for x in xrange(c,len(z)-(c+1)):
        new_arr = np.hstack(z[x-c:x+c])
        hph_arr_list.append(np.percentile(new_arr[((new_arr >= np.percentile(new_arr,15)) & (new_arr <= np.percentile(new_arr,85)))],99))
    return np.array(map(float,hph_arr_list))

The variable z has a length of ~15 million and c is value of window size + and -. The function is basically a sliding window that calculates a percentile value per iteration. Any help would be appreciated! z is an array of arrays (hence the np.hstack). Maybe any idea if numba would help with this. If so, how to implement?

Upvotes: 1

Views: 322

Answers (1)

J&#233;r&#244;me Richard
J&#233;r&#244;me Richard

Reputation: 50668

The slow part of the computation appear to be the line np.percentile(new_arr[((new_arr >= np.percentile(new_arr,15)) & (new_arr <= np.percentile(new_arr,85)))],99). This is due to the unexpectedly slow np.percentile on small arrays as well as the creation of several intermediate arrays.

Since new_arr is actually quite small, it is much faster to just sort it and make the interpolation yourself. Moreover, numba can also help to speed the computation up.

@njit #Use @njit instead of @jit to increase speed
def filter(arr):
    arr = arr.copy() # This line can be removed to modify arr in-place
    arr.sort()
    lo = int(math.ceil(len(arr)*0.15))
    hi = int(len(arr)*0.85)
    interp = 0.99 * (hi - 1 - lo)
    interp = interp - int(interp)
    assert lo <= hi-2
    return arr[hi-2]* (1.0 - interp) + arr[hi-1] * interp

This code is 160 times faster with arrays of size 20 on my machine and should produce the same result.

Finally, you can speed up smooth_surface too by using automatic parallelization in numba (see here for more information). Here is an untested prototype:

@jit(parallel=True)
def smooth_surface(z,c):
    hph_arr = np.zeros(len(z)-(c+1)-c)
    for x in prange(c,len(z)-(c+1)):
        hph_arr[x-c] = filter(np.hstack(z[x-c:x+c]))
    return hph_arr

Upvotes: 1

Related Questions