如何利用 NumPy 的功能来修复和优化这段非常简单的“生命游戏”代码？

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from random import randint

arraySize = 50
Z = np.array([[randint(0, 1) for x in range(arraySize)] for y in range(arraySize)])


def computeNeighbours(Z):
    rows, cols = len(Z), len(Z[0])
    N = np.zeros(np.shape(Z))

    for x in range(rows):
        for y in range(cols):
            Q = [q for q in [x-1, x, x+1] if ((q >= 0) and (q < cols))]
            R = [r for r in [y-1, y, y+1] if ((r >= 0) and (r < rows))]
            S = [Z[q][r] for q in Q for r in R if (q, r) != (x, y)]
            N[x][y] = sum(S)

    return N

def iterate(Z):
    rows, cols = len(Z), len(Z[0])
    N = computeNeighbours(Z)

    for x in range(rows):
        for y in range(cols):
            if Z[x][y] == 1:
                if (N[x][y] < 2) or (N[x][y] > 3):
                    Z[x][y] = 0
            else:
                if (N[x][y] == 3):
                    Z[x][y] = 1

    return Z

fig = plt.figure()

Zs = [Z]
ims = []

for i in range(0, 100):
    im = plt.imshow(Zs[len(Zs)-1], interpolation = 'nearest', cmap='binary')
    ims.append([im])
    Zs.append(iterate(Z))

ani = animation.ArtistAnimation(fig, ims, interval=250, blit=True)
plt.show()

首先，我写了一个简单的实现《生命的游戏》 http://en.wikipedia.org/wiki/Conway%27s_Game_of_Life仅使用标准 Python 工具。我绘制了它，它运行正确，并且动画正确。

接下来我尝试将数组转换为 NumPy 数组，这就是代码现在所在的位置。然而，动画似乎不再起作用，我一直无法弄清楚为什么。 (Update:此错误已修复！）

接下来，我有兴趣利用 NumPy 来优化我的代码。到目前为止，我所做的是将我使用的 Python 数组转换为 NumPy 数组。虽然我确信性能有所提高，但无法识别。

我有兴趣知道对于此类应用程序将执行什么优化顺序，这样我就可以了解如何在我当前的项目中使用 NumPy 的强大功能，该项目只是一个（可能）具有许多规则的三维元胞自动机。

对代码进行以下更改修复了动画错误：

1）改变iterate创建 Z 的深层副本，然后对该深层副本进行更改。新的iterate:

def iterate(Z):
    Zprime = Z.copy()
    rows, cols = len(Zprime), len(Zprime[0])
    N = computeNeighbours(Zprime)

    for x in range(rows):
        for y in range(cols):
            if Zprime[x][y] == 1:
                if (N[x][y] < 2) or (N[x][y] > 3):
                    Zprime[x][y] = 0
            else:
                if (N[x][y] == 3):
                    Zprime[x][y] = 1

    return Zprime

2）根据1，更改这段代码：

for i in range(0, 100):
    im = plt.imshow(Zs[len(Zs)-1], interpolation = 'nearest', cmap='binary')
    ims.append([im])
    Zs.append(iterate(Z))

to:

for i in range(0, 100):
    im = plt.imshow(Zs[len(Zs)-1], interpolation = 'nearest', cmap='binary')
    ims.append([im])
    Zs.append(iterate(Zs[len(Zs)-1]))

下列：

for x in range(rows):
        for y in range(cols):
            if Z[x][y] == 1:
                if (N[x][y] < 2) or (N[x][y] > 3):
                    Z[x][y] = 0
            else:
                if (N[x][y] == 3):
                    Z[x][y] = 1

可以替换为：

set_zero_idxs = (Z==1) & ((N<2) | (N>3))
set_one_idxs = (Z!=1) & (N==3)
Z[set_zero_idxs] = 0
Z[set_one_idxs] = 1

这将比您拥有的循环需要更多的操作，但我希望它会更快。

EDIT:

所以我刚刚对这两个解决方案进行了基准测试，毫不奇怪，numpy 版本快了 180 倍：

In [49]: %timeit no_loop(z,n)
1000 loops, best of 3: 177 us per loop

In [50]: %timeit loop(z,n)
10 loops, best of 3: 31.2 ms per loop

EDIT2:

我认为这个循环：

for x in range(rows):
        for y in range(cols):
            Q = [q for q in [x-1, x, x+1] if ((q >= 0) and (q < cols))]
            R = [r for r in [y-1, y, y+1] if ((r >= 0) and (r < rows))]
            S = [Z[q][r] for q in Q for r in R if (q, r) != (x, y)]
            N[x][y] = sum(S)

可以替换为：

N = np.roll(Z,1,axis=1) + np.roll(Z,-1,axis=1) + np.roll(Z,1,axis=0) + np.roll(Z,-1,axis=0)

这里有一个隐含的假设，即数组没有界限，并且x[-1]旁边是x[0]。如果这是一个问题，您可以在数组周围添加一个零缓冲区：

shape = Z.shape
new_shape = (shape[0]+2,shape[1]+2)
b_z = np.zeros(new_shape)
b_z[1:-1,1:-1] = Z
b_n = np.roll(b_z,1,axis=1) + np.roll(b_z,-1,axis=1) + np.roll(b_z,1,axis=0) + np.roll(b_z,-1,axis=0)
N = b_n[1:-1,1:-1]

并作为基准：

In [4]: %timeit computeNeighbours(z)
10 loops, best of 3: 140 ms per loop 

In [5]: %timeit noloop_computeNeighbours(z)
10000 loops, best of 3: 133 us per loop

In [6]: %timeit noloop_with_buffer_computeNeighbours(z)
10000 loops, best of 3: 170 us per loop

所以只需稍微改进一个因素1052。 Numpy 万岁！