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

Question

这是我用牛顿方法制作分形的一个小脚本.

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数组的优势？

PS:如果你想在不等太久的情况下尝试代码,最好使用100×100.

额外背景:
请参阅牛顿方法以查找多项式的零.
分形的基本思想是测试复平面中的猜测并计算迭代次数以收敛到零.这就是递归的内容newton(),最终会返回步数.复平面中的猜测表示图像中的像素,通过收敛的步数来着色.从一个简单的算法,你得到这些美丽的分形.

Answer 1

我从Lauritz五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之前正在使用蒙面阵列,然后才意识到它们正在减慢速度.我在很长一段时间没有看过它们,但我确实记得蒙面阵列在过去是一个缓慢的来源.我相信这是因为它们主要是在纯粹的Python上通过numpy数组实现,而不是几乎完全用C语言写成numpy数组.