How can I have straight contourlines in matplotlib?

Question

I am plotting values with imshow, and I want to have one contourline at a certain value. However, pyplot.contour() uses some kind of interpolation which causes the contourlines to be diagonal around the point. How can I make sure that the lines are exactly lined up with my rectangular boxes (so only horizontal and vertical lines)?

(Anyone who wants to reproduce the picture I've got, the values are uploaded here)

A picture of the data looks like this:

produced with this code:

pyplot.imshow(KS_imshow, extent = [5. ,8., 0., 22., ], origin='lower', interpolation='nearest', aspect='auto', cmap = 'Blues', vmin = 0., vmax = 1.)
cbar = pyplot.colorbar()

CS2 = pyplot.contour(ri,phii,KS_imshow,levels=[0.5], colors='r')
cbar.add_lines(CS2)

pyplot.show()

The variables ri, phii and KS_imshow are in the linked document.

Answer 1

The problem is that imshow creates "pixels", but the underlying data are just points (at the centers). Thus contour does not know anything about the image which imshow creates. However, you can create a similar image by upscaling the original data and then use contour on that. It is certainly a hack, but it achieves what you want. There remains a problem at the edges though and I'm not sure how to solve that.

import matplotlib.pyplot as plt
import numpy as np
import scipy.ndimage

# data ranges
xr = [5., 8.]
yr = [0., 22.]

# pixel widths
x_pw = np.diff(xr) / (KS_imshow.shape[1])
y_pw = np.diff(yr) / (KS_imshow.shape[0])

# plot the image
plt.imshow(KS_imshow, extent=xr+yr, origin='lower', interpolation='nearest',
        aspect='auto', cmap='Blues', vmin=0., vmax=1.)
cbar = plt.colorbar()

# upscale by a factor of 50 (might be an issue for large arrays)
highres = scipy.ndimage.zoom(KS_imshow, 50, order=0, mode='nearest') 

# correct the extent by the pixel widths
extent = np.array(xr+yr) + np.array([x_pw, -x_pw, y_pw, -y_pw]).flatten()/2

# create the contours
CS2 = plt.contour(highres, levels=[0.5], extent=extent, origin='lower',
        colors='r', linewidths=2)
cbar.add_lines(CS2)

plt.show()

Result:

However, just to show a threshold of 0.5, I would suggest to customize the colormap instead of using a contour line:

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as mcolors

blues = plt.cm.Blues(np.linspace(0,1,200))
reds = plt.cm.Reds(np.linspace(0,1,200))
colors = np.vstack((blues[0:128,:], reds[-129:,:]))

i = np.linspace(0,1,256)
r = np.column_stack((i, colors[:-1,0], colors[1:,0]))
g = np.column_stack((i, colors[:-1,1], colors[1:,1]))
b = np.column_stack((i, colors[:-1,2], colors[1:,2]))
d = dict(red=r, green=g, blue=b)
mycmap = mcolors.LinearSegmentedColormap('mymap', d, N=256)

plt.imshow(KS_imshow, extent=[5, 8, 0, 22], origin='lower',
        interpolation='nearest', aspect='auto', cmap=mycmap,
        vmin=0., vmax=1.)

cbar = plt.colorbar()

plt.show()

Result:

Answer 2

As an addition to the nice answer by @hitzig I present some code that makes it simpler to draw straight contour lines. However, the underlying principle is exactly the same.

All we need are

• the extent of the data ...
• and np.kron

Then we can scale up our data using big_data = np.kron(data, np.ones((factor, factor))) and draw the contour lines using the big_data array. We make sure that the size of the image stays the same by passing the extent of the original data.

Example:

# Make up some data
data = np.zeros((10, 20))
data[2:4, 2:8] = 1 + np.random.random((2,6))

# Extent of the data into x and y directions
# (left, right, bottom, top)
extent = [0, 20, 0, 10]

# Plot the data a few times. Each time, the contours
# get drawn based on "enlarged" data to some factor
enlargement_factors = [1, 2, 10]
fig, axs = plt.subplots(len(enlargement_factors), 1)

for i, fac in enumerate(enlargement_factors):
    # Draw the data
    im = axs[i].imshow(data, origin='lower', aspect='auto', extent=extent)

    # Scale the data up (enlarge) ... or leave equal if fac==1
    big_data = np.kron(data, np.ones((fac, fac)))

    # Draw the contour lines of the data
    axs[i].contour(big_data, levels=[0.5], extent=extent, colors='w')
    axs[i].set_title('Enlargement factor: {}'.format(fac))

fig.tight_layout()