# -*- coding: utf-8 -*-
"""
Created on Tue Jan 10 15:47:33 2017

@author: Mariusz
"""
import numpy as np
import matplotlib.pyplot as plt
import random

# 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)):
        C[i1, i2] = random.randint(0, 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()