# -*- coding: utf-8 -*- """ Created on Tue Jan 10 12:35:28 2017 @author: Mariusz """ import numpy as np import matplotlib.pyplot as plt # solve equation x^3 - 1 = 0 using Newton's method def newtonsolve(x): N = 100 tol = 1.0e-10 i = 0 while i < N: rk = x**3-1 if abs(rk) < tol: break gk = 3*x**2 x = x - rk/gk if not np.isfinite(x): break i += 1 return x # Roots of the equation x^3 - 1 = 0 r1 = 1.0 r2 = 0.5*(-1 + 1j*np.sqrt(3)) r3 = 0.5*(-1 - 1j*np.sqrt(3)) # Tolerance tol = 1.0e-08 re = np.linspace(-2, 2, 200) im = np.linspace(-2, 2, 200) RE, IM = np.meshgrid(re, im) C = np.zeros(np.shape(RE)) for i1 in np.arange(0, np.size(RE, 0)): for i2 in np.arange(0, np.size(IM, 1)): r = newtonsolve(RE[i1, i2] + 1j*IM[i1, i2]) if abs(r - r1) < tol: C[i1, i2] = 1 if abs(r - r2) < tol: C[i1, i2] = 2 if abs(r - r3) < tol: C[i1, i2] = 3 img1 = plt.imshow(C, extent=[re[0], re[-1], im[0], im[-1]], clim = (0.0, 3.0), aspect='auto') img1.set_cmap('jet') plt.colorbar() plt.plot(np.real(r1), np.imag(r1), 'gx', lw=2, markersize=10) plt.plot(np.real(r2), np.imag(r2), 'gx', lw=2, markersize=10) plt.plot(np.real(r3), np.imag(r3), 'gx', lw=2, markersize=10) plt.xlabel('real') plt.ylabel('imag') plt.show()