50. Matplotlib: Adding contour lines to a heatmap

Here will will add contour lines to a heat map.

We’ll use something a little more interesting for the array of values, we’ll define a Mandlebrot fractal function. We have build a 1,000 and 1,000 array and calculate z as a Mandlebrot function of x and y.

The heatmap is drawn with plt.imshow, and then contour lines are added with plt.contour.

See if you can follow how the arrays are built up, and the Mandlebrot function used to calculate Z, but the main purpose is to demonstrate adding contour lines to a heat map.

The Mandlebrot function may take up to 30 seconds to calculate (maybe longer on an old computer).

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

%matplotlib inline

# Mandlebrot function - a fractal

def iter_count(c,max_iter):
    x=c
    for n in range(max_iter):
        if abs(x)>2:
            return (n)
        x=x**2+c
    return (max_iter)

# Build a 512 * 512 array of x, y
    
n=1000
max_iter=75
# xmin,xmax,ymin,ymax = -0.32,0,0.8,01.1
xmin,xmax,ymin,ymax = -0.32,0.1,0.7,1.1
x_series=np.linspace(xmin,xmax,n)
y_series=np.linspace(ymin,ymax,n)

# Create an empty array for z (which will calculated from all x and y)

z=np.empty((n,n))

# Calculate z based on the Mandlebrot function

for i,y in enumerate(y_series):
    for j,x in enumerate(x_series):
        z[i,j]=iter_count(complex(x,y),max_iter)
        
# THE CHART

# Show a grey scale heatmap with imshow        

plt.imshow(z,
           cmap=cm.binary,
           interpolation='bicubic',
           origin='lower',
           extent=(xmin,xmax,ymin,ymax))

# Extent (above) defines the axes values which are not held in the
# results array (z) itself

# Add contours

levels=[2, 4, 6, 8, 10, 12, 14, 16] # Define levels 

ct=plt.contour(x_series,y_series,z,levels,cmap=cm.gray) 

# Add contour labels

plt.clabel(ct,fmt='%d')

plt.show()

plot_20

 

One thought on “50. Matplotlib: Adding contour lines to a heatmap

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google+ photo

You are commenting using your Google+ account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s