How to boost/enhance the speed of the subtraction 2d array in Python?

Viewed 53

all!

I'm establishing a class with Python. In one of the function in the class, I need to subtract two large arrays, here is the code:

def gradient(x, h, data):
    u = (x - data) / h ** 2
    c = np.exp(- (h ** 2 / 2) * np.sum(u ** 2, axis=-1))
    c = c.reshape(-1, 1)
    g = -(u.transpose().dot(c)) / np.sum(c)

    return g

where x is a 1d numpy array of length n, data is an 2d numpy array with shape = (N, n) (in other word, a matrix of size N * n).

I test the speed of this function when N = 30,000, n = 10,000. The time cost is between 20s~100s (I also don't understand why the time cost could vary so much) And the 1st line -- subtraction of x and data is the main part of the time cost. (p.s. In numpy, if x is a 1d array of size n, data is a 2d array of size (N, n), then x - data will be automatically regard as the row-wise subtraction of data from x.)

from time import time

data = np.random.normal(size = 30000 * 10000).reshape(30000, 10000)
x = np.random.normal(size = 10000)
h = 0.2
start_time = time()
gradient(x, h, data)
print("Elapsed time is: ", time() - start_time)


Elapsed time is:  36.4828040599823

This is too low to serve the purpose of my class. Any method that could boost the speed of this 4-line function in the class? The ideal time for this function should be around 2s~10s, is that possible to reach this goal? Great thanks in advance for any suggestions/ideas

1 Answers

With smaller arrays (your's kills my memory):

In [1]: M, N = 3000,1000
   ...: data = np.random.normal(size = (M,N))
   ...: x = np.random.normal(size = N)
   ...: h = 0.2*10

I had to increase h, since otherwise I was getting divide by 0 warnings; c got too small.

time:

In [3]: timeit gradient(x,h,data)
43.6 ms ± 168 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

lets check the steps:

In [5]: u=(x-data)/h**2
In [6]: timeit u=(x-data)/h**2
21.8 ms ± 79.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

So that's about half of the total time; it is working with the data size. You can't do much about that speed - at least not with numpy itself.

Only u**2 works with the full data size; sum cuts it down substantially.

In [7]: c = np.exp(- (h ** 2 / 2) * np.sum(u ** 2, axis=-1))
In [8]: timeit c = np.exp(- (h ** 2 / 2) * np.sum(u ** 2, axis=-1))
15 ms ± 61.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [9]: c = c.reshape(-1, 1)

dot uses fast BLAS functions where possible, so has good speed, even when multiplying a (1000,3000) with (3000).

In [10]: g = -(u.transpose().dot(c)) / np.sum(c)
In [11]: timeit g = -(u.transpose().dot(c)) / np.sum(c)
3.19 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [12]: u.shape
Out[12]: (3000, 1000)
In [13]: c.shape
Out[13]: (3000, 1)
In [14]: g.shape
Out[14]: (1000, 1)

numba and numexpr may improve speeds here, but I don't use them enough to explore that for you.

Times in seconds suggest you are pushing your memory's limits.

Related