Как лучше всего оптимизировать вычисления, повторяемые по сетке NxM в Python

Question

Работая в Python, я делаю некоторые физические вычисления по сетке значений NxM, где N переходит от 1 до 3108, а M переходит от 1 до 2304 (это соответствует большому изображению). Мне нужно вычислить значение в каждой точке этого пространства, что составляет ~ 7 миллионов вычислений. Мой текущий подход мучительно медленный, и мне интересно, есть ли способ выполнить эту задачу, и это не займет несколько часов ...

Мой первый подход состоял в том, чтобы просто использовать вложенные циклы for, но это казалось наименее эффективным способом решения моей проблемы. Я попытался использовать nditer NumPy и выполнять итерации по каждой оси отдельно, но я читал, что это на самом деле не ускоряет мои вычисления. Вместо того чтобы циклически проходить каждую ось по отдельности, я также попытался создать трехмерный массив и цикл по внешней оси, как показано в ответе Брайана здесь Как я могу, в python, итерировать несколько множественных списков одновременно, чисто? . Вот текущее состояние моего кода:

import numpy as np
x,y = np.linspace(1,3108,num=3108),np.linspace(1,2304,num=2304) # x&y dimensions of image
X,Y = np.meshgrid(x,y,indexing='ij')
all_coords = np.dstack((X,Y)) # moves to 3-D
all_coords = all_coords.astype(int) # sets coords to int

Для справки, all_coords выглядит так:

array([[[1.000e+00, 1.000e+00],
        [1.000e+00, 2.000e+00],
        [1.000e+00, 3.000e+00],
        ...,
        [1.000e+00, 2.302e+03],
        [1.000e+00, 2.303e+03],
        [1.000e+00, 2.304e+03]],

       [[2.000e+00, 1.000e+00],
        [2.000e+00, 2.000e+00],
        [2.000e+00, 3.000e+00],
        ...,
        [2.000e+00, 2.302e+03],
        [2.000e+00, 2.303e+03],
        [2.000e+00, 2.304e+03]],

и так далее. Вернуться к моему коду ...

'''
- below is a function that does a calculation on the full grid using the distance between x0,y0 and each point on the grid.
- the function takes x0,y0 and returns the calculated values across the grid
'''
def do_calc(x0,y0):
    del_x, del_y = X-x0, Y-y0
    np.seterr(divide='ignore', invalid='ignore')
    dmx_ij = (del_x/((del_x**2)+(del_y**2))) # x component
    dmy_ij = (del_y/((del_x**2)+(del_y**2))) # y component
    return dmx_ij,dmy_ij

# now the actual loop

def do_loop():
    dmx,dmy = 0,0
    for pair in all_coords:
        for xi,yi in pair:
            DM = do_calc(xi,yi)
            dmx,dmy = dmx+DM[0],dmy+DM[1]
    return dmx,dmy

Как вы можете заметить, запуск этого кода занимает невероятно много времени ... Если есть какой-либо способ изменить мой код так, чтобы его выполнение не заняло несколько часов, мне было бы крайне интересно узнать, как это сделать тот. Заранее спасибо за помощь.

Answer 1

Вот метод, который дает ускорение в 10000 раз при N=310, M=230. Поскольку метод масштабируется лучше, чем исходный код, я ожидаю, что при полном размере задачи коэффициент будет превышать миллион.

Метод использует сдвиговую инвариантность задачи. Например, del_x**2 по сути одинаков с точностью до сдвига при каждом вызове do_calc, поэтому мы вычисляем его только один раз.

Если выходные данные do_calc взвешены до суммирования, проблема больше не является полностью инвариантной трансляцией, и этот метод больше не работает. Результат, однако, затем может быть выражен в терминах линейной свертки. На N=310, M=230 это все еще оставляет нам более чем 1000-кратное ускорение. И, опять же, это будет больше при полном размере проблемы

код для исходной задачи

import numpy as np

#N, M = 3108, 2304
N, M = 310, 230

### OP's code

x,y = np.linspace(1,N,num=N),np.linspace(1,M,num=M) # x&y dimensions of image
X,Y = np.meshgrid(x,y,indexing='ij')
all_coords = np.dstack((X,Y)) # moves to 3-D
all_coords = all_coords.astype(int) # sets coords to int

'''
- below is a function that does a calculation on the full grid using the distance between x0,y0 and each point on the grid.
- the function takes x0,y0 and returns the calculated values across the grid
'''
def do_calc(x0,y0):
    del_x, del_y = X-x0, Y-y0
    np.seterr(divide='ignore', invalid='ignore')
    dmx_ij = (del_x/((del_x**2)+(del_y**2))) # x component
    dmy_ij = (del_y/((del_x**2)+(del_y**2))) # y component
    return np.nan_to_num(dmx_ij), np.nan_to_num(dmy_ij)

# now the actual loop

def do_loop():
    dmx,dmy = 0,0
    for pair in all_coords:
        for xi,yi in pair:
            DM = do_calc(xi,yi)
            dmx,dmy = dmx+DM[0],dmy+DM[1]
    return dmx,dmy

from time import time

t = [time()]

### pp's code

x, y = np.ogrid[-N+1:N-1:2j*N - 1j, -M+1:M-1:2j*M - 1J]
den = x*x + y*y
den[N-1, M-1] = 1
xx = x / den
yy = y / den
for zz in xx, yy:
    zz[N:] -= zz[:N-1]
    zz[:, M:] -= zz[:, :M-1]
XX = xx.cumsum(0)[N-1:].cumsum(1)[:, M-1:]
YY = yy.cumsum(0)[N-1:].cumsum(1)[:, M-1:]
t.append(time())

### call OP's code for reference

X_OP, Y_OP = do_loop()
t.append(time())

# make sure results are equal

assert np.allclose(XX, X_OP)
assert np.allclose(YY, Y_OP)
print('pp {}\nOP {}'.format(*np.diff(t)))

Пример прогона:

pp 0.015251636505126953
OP 149.1642508506775

Код для взвешенной задачи:

import numpy as np

#N, M = 3108, 2304
N, M = 310, 230

values = np.random.random((N, M))
x,y = np.linspace(1,N,num=N),np.linspace(1,M,num=M) # x&y dimensions of image
X,Y = np.meshgrid(x,y,indexing='ij')
all_coords = np.dstack((X,Y)) # moves to 3-D
all_coords = all_coords.astype(int) # sets coords to int

'''
- below is a function that does a calculation on the full grid using the distance between x0,y0 and each point on the grid.
- the function takes x0,y0 and returns the calculated values across the grid
'''
def do_calc(x0,y0, v):
    del_x, del_y = X-x0, Y-y0
    np.seterr(divide='ignore', invalid='ignore')
    dmx_ij = (del_x/((del_x**2)+(del_y**2))) # x component
    dmy_ij = (del_y/((del_x**2)+(del_y**2))) # y component
    return v*np.nan_to_num(dmx_ij), v*np.nan_to_num(dmy_ij)

# now the actual loop

def do_loop():
    dmx,dmy = 0,0
    for pair, vv in zip(all_coords, values):
        for (xi,yi), v in zip(pair, vv):
            DM = do_calc(xi,yi, v)
            dmx,dmy = dmx+DM[0],dmy+DM[1]
    return dmx,dmy

from time import time
from scipy import signal

t = [time()]
x, y = np.ogrid[-N+1:N-1:2j*N - 1j, -M+1:M-1:2j*M - 1J]
den = x*x + y*y
den[N-1, M-1] = 1
xx = x / den
yy = y / den
XX, YY = (signal.fftconvolve(zz, values, 'valid') for zz in (xx, yy))

t.append(time())
X_OP, Y_OP = do_loop()
t.append(time())
assert np.allclose(XX, X_OP)
assert np.allclose(YY, Y_OP)
print('pp {}\nOP {}'.format(*np.diff(t)))

Пример прогона:

pp 0.12683939933776855
OP 158.35225439071655