speed up looping over three axis 'in numpy

Question

Мне нужно ускорить цикл, который делает что-то вроде кода ниже, я смутно знаю, что эти вещи иногда можно сделать с помощью numpy.einsum, но не совсем уверен, как его использовать , Или если у кого-то есть другие предложения?

import numpy as np 


x = np.random.normal(size=(206,11,11)) 
y = np.random.normal(size=(206,11,11)) 
complx = x + 1j*y 
complx 

a,b,c = complx.shape 

for n in xrange(a): 
    #do somthing 
    z = np.zeros(b) 
    for i in xrange(b):  
     z[i] = (complx[n,:,:].real[i][i]*complx[n,:,:].real[i][i] +\ 
       complx[n,:,:].imag[i][i]*complx[n,:,:].imag[i][i])(**-0.25) 

Answer 1

В случае, если вы хотите ускорить внутреннюю for петли вы можете сделать что-то вроде этого

import numpy as np 

x = np.random.normal(size=(206,11,11)) 
y = np.random.normal(size=(206,11,11)) 
complx = x + 1j*y 

# takes only the diagonal part of all the 11x11 matrices 
complx_diag = np.diagonal(complx,0,1,2) 

# do the calc 
zn = np.abs(complx_diag)**(-0.5) 

for n in xrange(a): 
    z = zn[n] 
    # do your stuff 

В случае вашего stuff не слишком сложно это может быть векторизация, а также (весьма вероятно). Чем больше вы рассчитываете вне цикла for, тем быстрее ваш код.

Answer 2

Если я не ошибаюсь, это более или менее то, что вы хотите. Утверждения печати только для того, чтобы убедить себя, что расчет правилен.

def optimize_01(): 
    x = np.random.normal(size=(6, 11, 11)) 
    y = np.random.normal(size=(6, 11, 11)) 
    complx = x + 1j * y 

    a, b, _ = complx.shape 

    for n in range(a): 
     # do somthing 
     A = complx[n, :, :] 
     d = np.diagonal(A) 
     z = np.power(np.abs(d * d), -0.25) 

     print (d[0]) 
     print (z[0]) 
     print ((d[0].real * d[0].real + d[0].imag * d[0].imag) ** -0.25) 

EDIT: Если я сравню эту реализацию с вашей реализацией, я получаю следующее.

import timeit 

def optimize_02(): 
    x = np.random.normal(size=(206, 11, 11)) 
    y = np.random.normal(size=(206, 11, 11)) 
    complx = x + 1j * y 

    a, b, _ = complx.shape 

    for n in range(a): 
     # do somthing 
     A = complx[n, :, :] 
     d = np.diagonal(A) 
     z = np.power(np.abs(d * d), -0.25) 

def optimize_03(): 
    x = np.random.normal(size=(206, 11, 11)) 
    y = np.random.normal(size=(206, 11, 11)) 
    complx = x + 1j * y 

    a, b, _ = complx.shape 

    for n in range(a): 
     # do somthing 
     z = np.zeros(b) 
     for i in range(b):  
      z[i] = (complx[n, :, :].real[i][i] * complx[n, :, :].real[i][i] + \ 
        complx[n, :, :].imag[i][i] * complx[n, :, :].imag[i][i]) ** (-0.25) 

if __name__ == '__main__': 
    print (timeit.timeit(optimize_02, number=10)) 
    print (timeit.timeit(optimize_03, number=10)) 

Результат:

0.03474012700007734 
0.09025639800074714 

С 6 массивов 1100 элементов, вместо 206 массивов 11 элементов, результат:

5.762741210999593 
5.771216576999905 

Похоже, мое решение не так быстро в конце концов.