Plotting Heat Map Using Python

For Machine Learning, practitioners commonly use  Python and R because they are open source languages. I was able to learn Python using "Pycharm" which I would strongly recommend for any Python beginners.

I was given a challenge to create 
(i) 2-dimensional array of size 100x100 and populate it with random floating points between 0 and 1, inclusive (i.e., [0,1]); 
(ii) plot the 2d array using any python library, to create a visual “heat map” representation of the data; 
(iii) write a loop that refreshes the numbers in the array and replots the heatmap each time the array is repopulated.  

Stretch assignment: Create a movie of the changing heat maps by playing each heat map frame by frame in a sequence. 

I was able to generate a heat map as shown in the picture with the following code:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

def animate(data, im):
    im.set_data(data)

def step():
    while True:
        data = np.random.rand(100, 100)
        yield data

fig, ax = plt.subplots()
im = ax.imshow(np.random.rand(100, 100), interpolation='nearest')
ani = animation.FuncAnimation(
    fig, animate, step, interval=100, repeat=True, fargs=(im, ))
plt.show()