Speed up Newtons Method using numpy arrays

Question

I am using Newton's method to generate fractals that visualise the roots and the number of iterations taken to find the roots.

I am not happy with the speed taken to complete the function. Is there are a way to speed up my code?

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

Possible optimisations in addition to vectorising the loops

• Save the result of z[mask] to avoid repeatedly using fancy indexing (~80% faster)
• Combine z - f(z) / f_prime(z) into one function (~25% in this case)
• Use lower precision (dtype=np.complex64) if the larger error is acceptable (~90%)
• Use broadcasting instead of meshgrid to create the z plane (negligible effect, but generally a good idea).

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