如何使用 numpy 数组加速分形生成？

这是我为使用牛顿方法制作分形而编写的一个小脚本。

import numpy as np
import matplotlib.pyplot as plt

f = np.poly1d([1,0,0,-1]) # x^3 - 1
fp = np.polyder(f)

def newton(i, guess):
    if abs(f(guess)) > .00001:
        return newton(i+1, guess - f(guess)/fp(guess))
    else:
        return i

pic = []
for y in np.linspace(-10,10, 1000):
    pic.append( [newton(0,x+y*1j) for x in np.linspace(-10,10,1000)] )

plt.imshow(pic)
plt.show()

我正在使用 numpy 数组，但仍然循环遍历 1000×1000 linspaces 的每个元素来应用newton()函数，它作用于单个猜测而不是整个数组。

我的问题是这样的：如何改变我的方法以更好地利用 numpy 数组的优势？

P.S.：如果您想尝试代码而不等待太久，最好使用 100×100。

额外背景：
请参阅牛顿法查找多项式的零点。
分形的基本思想是测试复平面中的猜测并计算收敛到零的迭代次数。这就是递归的意义所在newton()，最终返回步数。复平面中的猜测代表图片中的一个像素，由收敛步骤数着色。通过简单的算法，您可以得到这些美丽的分形。

我使用了 Lauritz V. Thaulow 的代码，并且能够通过以下代码获得相当显着的加速：

import numpy as np
import matplotlib.pyplot as plt
from itertools import count

def newton_fractal(xmin, xmax, ymin, ymax, xres, yres):
    yarr, xarr = np.meshgrid(np.linspace(xmin, xmax, xres), \
                             np.linspace(ymin, ymax, yres) * 1j)
    arr = yarr + xarr
    ydim, xdim = arr.shape
    arr = arr.flatten()
    f = np.poly1d([1,0,0,-1]) # x^3 - 1
    fp = np.polyder(f)
    counts = np.zeros(shape=arr.shape)
    unconverged = np.ones(shape=arr.shape, dtype=bool)
    indices = np.arange(len(arr))
    for i in count():
        f_g = f(arr[unconverged])
        new_unconverged = np.abs(f_g) > 0.00001
        counts[indices[unconverged][~new_unconverged]] = i
        if not np.any(new_unconverged):
            return counts.reshape((ydim, xdim))
        unconverged[unconverged] = new_unconverged
        arr[unconverged] -= f_g[new_unconverged] / fp(arr[unconverged])

N = 1000
pic = newton_fractal(-10, 10, -10, 10, N, N)

plt.imshow(pic)
plt.show()

对于 N=1000，使用 Lauritz 代码得到的时间为 11.1 秒，使用此代码得到的时间为 1.7 秒。

这里有两个主要的加速。首先，我使用 meshgrid 来加速输入值的 numpy 数组的创建。当 N=1000 时，这实际上是加速的一个非常重要的部分。

第二个加速来自于仅对未收敛部分进行计算。劳里茨（Lauritz）为此使用了屏蔽阵列，然后才意识到它们正在减慢速度。我已经有一段时间没有看过它们了，但我确实记得屏蔽数组在过去是缓慢的根源。我相信这是因为它们主要是通过 numpy 数组在纯 Python 中实现的，而不是像 numpy 数组一样几乎完全用 C 编写。