переписать al oop in numpy для более быстрого выполнения

Question

Я пишу функцию, которая принимает массив numpy a длины 200 и матрицу M размера 200 x 200 и выполняет следующую операцию:

for i in range(len(a)):
    x = a[i]
    for j in range(len(a)):
        y = a[j]
        z = M[i][j]
        d[i][j] = 2 * z/(y+x)
return d

Как можно Я векторизовал этот кусок кода, чтобы увеличить время выполнения?

Answer 1

Все функции ufuncs

Numpy имеют метод outer для выполнения операций "крест-накрест" над двумя массивами. Таким образом, чтобы избежать большинства промежуточных вычислений и векторизации, насколько это возможно:

def f(M, a):
    return 2 * M / np.add.outer(a, a)

---

Ответ для старой версии вопроса (слева, потому что он все еще полезен) :

Для таких вещей я всегда старался работать поэтапно и пытаться найти правильное выражение einsum.

# the definition given in the original question,
# before the z / (y + x) update
def f0():
    d = np.empty((3,3))
    for i in range(len(a)):
        x = a[i]
        for j in range(len(a)):
            y = a[j]
            z = M[i][j]
            d[i][j] = 2 * x/(y+z)
    return d

# rewrite things inlined
def f1():
    d = np.empty((3,3))
    for i in range(len(a)):
        for j in range(len(a)):
            d[i, j] = 2 * a[i]/(a[j] + M[i, j])
    return d

# factor out broadcasting
def f2():
    d = np.empty((3,3))
    for i in range(len(a)):
        m = a + M[i, :]
        for j in range(len(a)):
            d[i,j] = 2 * a[i]/m[j]
    return d

# more broadcasting
def f3():
    d = np.empty((3,3))
    m = a + M
    for i in range(len(a)):
        for j in range(len(a)):
            d[i,j] = 2 * a[i]/m[i,j]
    return d

# now turn loops into einsums
def f4():
    d = np.empty((3,3))
    m = 1/(a + M)
    d[:,:] = 2 * np.einsum('i,ij->ij', a, m)
    return d

# collect everything
def f5():
    return np.einsum('i,ij->ij', a, 2 / (a + M))

Answer 2

Помимо numpy -векторизации с использованием Numba также был бы простой и производительный метод, позволяющий ускорить код с помощью циклов. Пример

import numpy as np
import numba as nb

@nb.njit(fastmath=True,error_model="numpy",parallel=True)
def calc(a,M):
    d=np.empty((a.shape[0],a.shape[0]))
    for i in nb.prange(a.shape[0]):
        x = a[i]
        for j in range(a.shape[0]):
            y = a[j]
            z = M[i,j]
            d[i,j] = 2. * z/(y+x)
    return d

Сроки

M=np.random.rand(200,200)
a=np.random.rand(200)

d=calc(a,M) #first call takes longer due to compilation overhead
%timeit d=calc(a,M)
#parallel=True there is only quite limited speedup because of the small problem (200x200)
#11 µs ± 51 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
#parallel=False
#21.2 µs ± 191 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

#pure numpy solution (hpaulj)
%timeit d = 2 * M/(a[:,None]+a[None,:])
#75.7 µs ± 386 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

#without compilation
#20.8 ms ± 500 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 3

Вы могли бы сделать что-то вроде

d = 2*numpy.atleast_2d(a).T/(a+M)

Answer 4

С парой выборочных массивов (вы должны были предоставить их):

In [196]: a = np.arange(1,4); M = np.arange(1,10).reshape(3,3)                                                       
In [197]: a                                                                                                          
Out[197]: array([1, 2, 3])
In [198]: M                                                                                                          
Out[198]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
In [199]: d = 2 * M/(a[:,None]+a[None,:])                                                                            
In [200]: d                                                                                                          
Out[200]: 
array([[1.        , 1.33333333, 1.5       ],
       [2.66666667, 2.5       , 2.4       ],
       [3.5       , 3.2       , 3.        ]])

a[None,:] можно было бы упростить до a, но я хотел уточнить использование вещания для вычисления этого внешнего товар. В numpy есть различные инструменты для этого. Мне нравится индексирование None, потому что оно простое и идиоматическое c.

тестирование на ваш код (опять же, вы должны были предоставить такой результат):

In [202]: def foo(a): 
     ...:     d = np.zeros((3,3)) 
     ...:     for i in range(len(a)): 
     ...:         x = a[i] 
     ...:         for j in range(len(a)): 
     ...:             y = a[j] 
     ...:             z = M[i][j] 
     ...:             d[i][j] = 2 * z/(y+x) 
     ...:     return d 
     ...:                                                                                                            
In [203]: foo(a)                                                                                                     
Out[203]: 
array([[1.        , 1.33333333, 1.5       ],
       [2.66666667, 2.5       , 2.4       ],
       [3.5       , 3.2       , 3.        ]])