使用 numpy 数组加速牛顿法答案

【问题标题】：Speed up Newtons Method using numpy arrays使用 numpy 数组加速牛顿法
【发布时间】：2016-03-16 07:27:49
【问题描述】：

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

我对完成该功能的速度不满意。有没有办法加快我的代码？

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)

【问题讨论】：

for 循环在k 和j 上的目的是什么？这是一种测试收敛性的奇怪方法——而且你根本没有使用eps。
不要使用iter，你正在覆盖一个python构建
您想要 z 中每个元素的迭代次数？为什么你需要一个二维数组呢？

标签： python arrays performance numpy newtons-method

【解决方案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 函数。

【讨论】：

您应该返回迭代而不是项目（自动更正？）。此外，此方法返回的结果与他的原始代码不同。
mask = (iters == 0) & (z == z_old)?对吗？
@MylesBaker 你说得对，items 不是iters。是的，刚刚尝试再次自动更正它。

【解决方案2】：

除了矢量化循环之外的可能优化

保存z[mask] 的结果以避免重复使用花哨的索引（快约80%）
将 z - f(z) / f_prime(z) 合并为一个函数（在这种情况下约为 25%）
如果较大的误差可以接受 (~90%)，则使用较低的精度 (dtype=np.complex64)
使用广播而不是网格网格来创建 z 平面（效果可以忽略不计，但通常是个好主意）。

def iterstep(z):
    return 0.75*z + 0.25 / z**3

def newton_raphson(Re, Im, max_iter):
    z = Re + 1j * Im[:, np.newaxis]
    itercount = np.zeros_like(z, dtype=np.int32)
    mask = np.ones_like(z, dtype=np.bool)
    for i in range(max_iter):
        z_old = z.copy()
        z[mask] = iterstep(z[mask])
        mask = (z != z_old)
        itercount[mask] = i
    return np.angle(z), itercount

axis = np.linspace(-1.5, 1.5, num=1000, dtype=np.complex64)
angle, itercount = newton_raphson(Re=axis, Im=axis, max_iter=20)

What it looks like

【讨论】：

【解决方案3】：

这应该是您正在寻找并适用于所有形状的 numpy 数组的内容。

def newton_numpy(x, y, max_iter=1000, eps = 1e-5):
    z = x + y * 1j

    # deltas is to store the differences between to iterations
    deltas = np.ones_like(z)
    # this is to store how many iterations were used for each element in z
    needed_iterations = np.zeros_like(z, dtype=int)

    it = 0
    # we loop until max_iter or every element converged
    while it < max_iter and np.any(deltas > eps):
        z_old = z.copy()

        # recalculate only for values that have not yet converged
        mask = deltas > eps
        z[mask] = z[mask] - (f(z[mask]) / f_prime(z[mask]))

        needed_iterations[mask] += 1
        deltas[mask] = np.abs(z_old[mask] - z[mask])
        it += 1

   return np.angle(z), needed_iterations

使用此代码：

n_points = 25
xmin = -1.5
xmax = 1.5
xs = np.linspace(xmin, xmax, n_points)
X, Y = np.meshgrid(xs, xs)
ang, iters = newton_numpy(X, Y, eps=1e-5, max_iters=1000)

每个元素需要 0.09 秒，平均 85 次迭代。

【讨论】：