#
# This program uses Newton's method on a _complex_ polynomial
# for different initial guesses for the root.
#
#

import numpy as np
from newton import newton
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib import colors

# Define extent of x0's real and imag part
xmin = -1.5
xmax = 1.5
ymin = -1.5
ymax = 1.5
xpix = 1024
ypix = 1024

# Create arrays of x0-values
xgrid = np.linspace(xmin, xmax, xpix);
ygrid = np.linspace(ymin, ymax, ypix);

# Images used for storing computation results
im1 = np.zeros( (xpix, ypix) )
im2 = np.zeros( (xpix, ypix) )

# Define function and its derivative
f = lambda x: x**3 - 1
df = lambda x: 3*x**2
    
# loop through all pixels and do Newton calculation
for j in range(ypix):
    print "processing line ", j, " of ", ypix
    for k in range(xpix):
        z0 = complex(xgrid[k], ygrid[j])
        (z, numits) = newton(f, df, z0, maxit = 30)
        im1[j,k] = numits
        im2[j,k] = np.angle(z)

# A function that fixes colors
norm = lambda im: colors.Normalize(vmax = np.max(im), vmin = np.min(im), clip = False)

# colorize according to this recipe:
# angle of root is the color hue.
# the number of iterations becomes the color intensity.
# the more iterations, the darker
pretty_im = cm.jet(norm(im2)(im2)) *  cm.gray(norm(-im1)(-im1))

# Draw the image and save
plt.imshow(pretty_im, extent=[xmin,xmax,ymin,ymax])
plt.title('Attraktorer for de tre roettene til f(x) = x^3 - 1')
plt.xlabel('Re x_0')
plt.ylabel('Im x_0')
plt.savefig('newton_fractal.pdf')
plt.show()