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)

What is the purpose of the `for`-loops over `k` and `j`? This is a strange way to test for convergence --- and you're not using the `eps` at-all. — DilithiumMatrix, Dec 10 '15 at 16:11
You want the number of iterations for each element in z? Why do you need a 2d array for this? — MaxNoe, Dec 10 '15 at 16:56

score 2 · Answer 1 · answered Dec 10 '15 at 16:50

2

You can simply vectorize the loops for fairly large speed gains:

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

Your presented equations are fairly simple; however, I would assume that your f and f_prime functions are significantly more complex. Further speedups can likely be found in those equations rather than in the presented problem.

I would also avoid the use of iter as a variable name as it is a built in python function.

answered Dec 10 '15 at 16:50

Daniel

17,131
7
51
70

You should return iters not items (autocorrect?). Also, this method returns a different result from his original code. – Myles Baker Dec 10 '15 at 16:53
`mask = (iters == 0) & (z == z_old)`? Is this right? – erip Dec 10 '15 at 16:55
@MylesBaker Whoops your right, `items` not `iters`. Yup, just tried to autocorrect it again. – Daniel Dec 10 '15 at 21:48

score 1 · Answer 2 · answered Dec 10 '15 at 17:12

This should be what you are looking for and be working for all shapes of numpy arrays.

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

With this code:

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)

It takes 0.09 seconds and a mean of 85 iterations per element.

score 1 · Accepted Answer · answered Dec 11 '15 at 01:58

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

Speed up Newtons Method using numpy arrays

3 Answers3