Оптимизация вычисления матриц в Python

Question

Я пытаюсь создать Абелеву песочную кучу , используя массивы numpy в python. Скорость расчета приемлема для меньших квадратных матриц, но для больших она значительно замедляется (матрица 200x200, с 20000 исходных частиц песка, занимающих до 20-30 минут). Есть ли способ ускорить / оптимизировать расчет матрицы? Пороговое значение равно 3.

Основной код прямо сейчас -

import numpy as np
n = 200
size = (n,n)
x = np.zeros(size)
m = 0 # mean
if n%2 == 0:
    m = int((n+1)/2)
else :
    m = int(n/2)

x[m][m] = 100000
z = int(x[m][m])

def f(x):
    count = 0
    for i in range(0,n):
        for j in range(0,n):
            if x[i][j] > 3:
                x[i][j] = x[i][j] - 4
                if i-1 >= 0 :
                    x[i-1][j] = x[i-1][j] + 1 
                if i+1 < n :
                    x[i+1][j] = x[i+1][j] + 1
                if j-1 >= 0 :
                    x[i][j-1] = x[i][j-1] + 1 
                if j+1 < n :    
                    x[i][j+1] = x[i][j+1] + 1
            elif x[i][j] <= 3:
                count = count + 1
    return x, count
for k in range(0,z):
    y, count = f(x)
    if count == n**2 :
        break
    elif count < n**2:
        continue
print(y)

Я пытался запустить матрицу 500x500 с 100 000 исходных частиц, но это заняло более 6 часов.

Answer 1

Хотя векторизация с помощью numpy не обязательно снизит вашу алгоритмическую сложность, она, вероятно, сократит ваши накладные расходы как минимум в несколько десятков раз. Как правило, если вы пишете явные циклы или используете явные операторы if, вам следует подумать о переосмыслении вашего подхода.

Здесь вам может помочь простая маскировка для реализации свертывания. Если у вас есть места свертывания, помеченные 1 в маске той же формы, что и x, вы можете напрямую вычесть свернутые сваи и добавить распределенный песок, просто сдвинув маску:

mask = (x >= 4)
x[mask] -= 4
x[:, :-1] += mask[:, 1:] # topple left
x[:, 1:] += mask[:, :-1] # topple right
x[:-1, :] += mask[1:, :] # topple up
x[1:, :] += mask[:-1, :] # topple down

Если count - это количество неиспользованных сайтов, вы можете использовать np.count_nonzero, чтобы получить его из маски:

count = np.count_nonzero(mask)

Если, с другой стороны, вы используете count, чтобы определить, когда остановить цикл, вам может быть проще переключиться на подсчет количества сайтов свержения:

count = np.sum(mask)

Внешний цикл завершается, когда эта версия счета достигает нуля (или исходная версия достигает x.size).

Answer 2

Вы можете использовать numba для этой цели (вы можете добавить nopython = True или использовать статические типы для большего ускорения):

from numba import jit
import numpy as np

n = 200
size = (n,n)
x = np.zeros(size)
m = 0 # mean
if n%2 == 0:
    m = int((n+1)/2)
else :
    m = int(n/2)

x[m][m] = 100000
z = int(x[m][m])

def f(x):
    count = 0
    for i in range(0,n):
        for j in range(0,n):
            if x[i][j] > 3:
                x[i][j] = x[i][j] - 4
                if i-1 >= 0 :
                    x[i-1][j] = x[i-1][j] + 1 
                if i+1 < n :
                    x[i+1][j] = x[i+1][j] + 1
                if j-1 >= 0 :
                    x[i][j-1] = x[i][j-1] + 1 
                if j+1 < n :    
                    x[i][j+1] = x[i][j+1] + 1
            elif x[i][j] <= 3:
                count = count + 1
    return x, count

@jit
def f_jit(x):
    count = 0
    for i in range(0,n):
        for j in range(0,n):
            if x[i][j] > 3:
                x[i][j] = x[i][j] - 4
                if i-1 >= 0 :
                    x[i-1][j] = x[i-1][j] + 1 
                if i+1 < n :
                    x[i+1][j] = x[i+1][j] + 1
                if j-1 >= 0 :
                    x[i][j-1] = x[i][j-1] + 1 
                if j+1 < n :    
                    x[i][j+1] = x[i][j+1] + 1
            elif x[i][j] <= 3:
                count = count + 1
    return x, count

%%timeit
f(x)
28.7 ms ± 602 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
f_jit(x)
59.9 µs ± 7.22 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)