python - 使用 numpy 数组加速牛顿法

Question

我正在使用牛顿法生成分形，以可视化根和找到根所需的迭代次数。

我对完成该功能的速度不满意。有什么方法可以加快我的代码速度吗？

def f(z):
    return z**4-1

def f_prime(z):
    '''Analytic derivative'''
    return 4*z**3

def newton_raphson(x, y, max_iter=20, eps = 1.0e-20):
    z = x + y * 1j
    iter = np.zeros((len(z), len(z)))
    for i in range(max_iter):
        z_old = z
        z = z-(f(z)/f_prime(z))
        for k in range(len(z[:,0])): #this bit of the code is slow. Can I do this WITHOUT for loops?
            for j in range(len(z[:,1])):
                if iter[k,j] != 0:
                    continue
                if z[k,j] == z_old[k,j]:
                    iter[k,j] = i
    return np.angle(z), iter #return argument of root and iterations taken

n_points = 1000; xmin = -1.5; xmax = 1.5
xs = np.linspace(xmin, xmax, n_points)
X,Y = np.meshgrid(xs, xs)
dat = newton_raphson(X, Y)

Answer 1

您可以简单地将循环矢量化以获得相当大的速度增益:

def newton_raphson(x, y, max_iter=20, eps = 1.0e-20):
    z = x + y * 1j
    nz = len(z)
    iters = np.zeros((nz, nz))
    for i in range(max_iter):
        z_old = z
        z = z-(f(z)/f_prime(z))
        mask = (iters == 0) & (z == z_old)
        iters[mask] = i

    return np.angle(z), items

您提供的方程式相当简单；但是，我假设您的 f 和 f_prime 函数要复杂得多。在这些方程式中而不是在提出的问题中可能会发现进一步的加速。

我也会避免使用 iter 作为变量名，因为它是一个内置的 python 函数。