在Python中优化矩阵计算

时间:2018-04-26 19:31:18

标签: python matrix

我正在尝试使用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个多小时。

2 个答案:

答案 0 :(得分: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)

答案 1 :(得分:0)

虽然使用numpy进行矢量化并不一定会降低算法的复杂性,但它可能会将开销降低至少几十倍。作为一般经验法则,如果您发现自己编写显式循环或使用显式if语句,则应考虑重新考虑您的方法。

在这里可以帮助您实现推翻的简单masking。如果你在一个与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确定何时停止循环,您可能会发现更容易切换到计算有多少推翻网站:

count = np.sum(mask)

当此版本的计数达到零(或原始版本达到x.size)时,外部循环终止。