使用 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 函数。

Answer 2

这应该是您正在寻找的，并且适用于所有形状的 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 次迭代。

Answer 3

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

保存 z[mask] 的结果以避免重复使用花哨的索引（快约 80%）
将 z - f(z) / f_prime(z) 合并为一个函数（在本例中为 ~25%）
如果可以接受较大的误差 (~90%)，则使用较低的精度 (dtype=np.complex64)
使用广播而不是 meshgrid 创建 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

使用 numpy 数组加速牛顿法

Speed up Newtons Method using numpy arrays

python

arrays

performance

numpy

newtons-method