Numpy 数组和 for 循环：如何改进它

Question

I am new with Python and I have a question about a for loop to speed up.

Let "u" be an numpy array of dimension (N,K) and let "kernel_vect" be a numpy array of dimension (K,) both of float64 numbers. I would like to speed up the following code (by eliminating the for loop for example)

Kernel_appo = np.zeros((N**2,))
    for k in range(K):
        uk = u[:,k]
        Mat_appo = np.outer(uk,uk)
        Kernel_appo = Kernel_appo  + kernel_vect[k] * routines.vec(Mat_appo)

Any idea? Thanks!

Answer 1

Here a faster implementation without a loop over k:

# Version 2
Kernel_appo = np.zeros((N**2,))
for n1 in range(N):
    for n2 in range(N):
        Kernel_appo[n1*N+n2] = (u[n1,:] * u[n2,:] * kernel_vect).sum()

We can make it faster by using the symmetry of the u product:

# Version 3
Kernel_appo = np.zeros((N,N))
for n1 in range(N):
    for n2 in range(n1,N):
        Kernel_appo[n1,n2] = (u[n1,:] * u[n2,:] * kernel_vect).sum()
Kernel_appo = np.triu(Kernel_appo, 1) + np.tril(Kernel_appo.transpose(), 0) # make the matrix symmetric
Kernel_appo = np.ravel(Kernel_appo, order='C')

Here is a version that remove one of the loop:

# Version 4
Kernel_appo = np.zeros((N,N))
for n1 in range(N):
    Kernel_appo[n1,n1:N] = ((u[n1,:] * kernel_vect) * u[n1:N,:]).sum(axis=1)
Kernel_appo = np.triu(Kernel_appo, 1) + np.tril(Kernel_appo.transpose(), 0)
Kernel_appo = np.ravel(Kernel_appo, order='C')

We still have a loop over N. However, this seems reasonable to keep it since N is small. Removing it will certainly force numpy to create huge matrices in memory which should cause a performance drop (and can even crash if N and K a very big).

Note that Version 4 will probably not be so fast if K is much bigger (since temporary numpy matrices could not fit in the cache of the CPU).

UPDATE : I just discover that this is possible to use the awesome np.einsum in this case:

# Version 5
Kernel_appo = np.ravel(np.einsum('ji,ki,i->jk', u, u, kernel_vect, optimize=True), order='C')

Be prepared, because this simpler implementation is also much faster (because numpy is able to vectorize the code and run it in parallel).

Here are performance results with N=50 and K=5000 on my machine:

Initial code: 58.15 ms
Version 2:    19.94 ms
Version 3:    10.11 ms
Version 4:     5.08 ms
Version 5:     0.57 ms

The final implementation is now about 100 times faster than the initial one!