'''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()