'''This program visualises the regions of attraction when computing 
the zeros of z**3 - 1 = 0 by means of Newton's method (cf. Chapter 9.4).

The axis window is defined by means of the following parameters:
ux = smallest x-coordinate
ox = largest x-coordinate
nx = number of grid points on the x-axis
uy = smallest x-coordinate
oy = largest x-coordinate
ny = number of grid points on the y-axis

the parameter kmax determines the number of newton iterations.
'''

import numpy as np
import matplotlib.pyplot as plt

ux, ox, nx = -2, 2, 401
uy, oy, ny = -1.5, 1.5, 301
x = np.linspace(ux, ox, nx)
y = np.linspace(uy, oy, ny)
xx, yy = np.meshgrid(x, y)

z = xx + yy*1j  # z is initialized as the complex plane

kmax = 50  # Increase this number to improve the quality of the plot

'''This will cause some overflow errors, which can be ignored. The computation
might take a few seconds to compute.
'''
for k in range(kmax):
    z -= (np.power(z,3) - 1) / (3 * np.power(z,2))

'''We want to know to which root each point in the complex plane converges. We
create an array c to keep track of that. 
    The entry 1 means we converge to 1
    The entry 2 means we converge to -0.5 + 0.5j*np.sqrt(3)
    The entry 3 means we converge to -0.5 - 0.5j*np.sqrt(3))
'''
c = ( 1 * np.isclose(z,1) 
    + 2 * np.isclose(z, -0.5 + 0.5j*np.sqrt(3))
    + 3 * np.isclose(z, -0.5 - 0.5j*np.sqrt(3)))

fig = plt.figure()
plt.pcolormesh(xx, yy, c, edgecolor='white', linewidth='0')
plt.show()