在Python中優化矩陣計算

Question

我正在嘗試在python中使用numpy數組制作Abelian Sandpile 。 對於較小的正方形矩陣，計算速度還可以，但是對於較大的正方形矩陣，計算速度會顯着降低（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

您可以為此目的使用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)

Answer 2

雖然使用numpy進行矢量化並不一定會降低您的算法復雜度，但可能會將您的開銷減少至少幾十倍。 作為一般經驗法則，如果發現自己編寫了顯式循環或使用了顯式if語句，則應考慮重新考慮您的方法。

在這里可以為您提供幫助的是簡單的掩蓋以實現下陷。 如果您在與x形狀相同的蒙版中具有標記為1的傾倒部位，則可以通過移動蒙版直接減去傾倒的樁並添加分布的沙子：

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來確定何時停止循環，則可能會發現更容易切換為計算有多少個topple網站：

count = np.sum(mask)

當此版本的count達到零（或原始版本達到x.size ）時，外部循環終止。