matplotlib draw a contour line on a colorbar plot

Question

I used below code to generate the colorbar plot of an image:

plt.imshow(distance)
cb = plt.colorbar()
plt.savefig(generate_filename("test_images.png"))
cb.remove()

The image looks likes this:

I want to draw a single contour line on this image where the signed distance value is equal to 0. I checked the doc of pyplot.contour but it needs a X and Y vector that represents the coordinates and a Z that represents heights. Is there a method to generate X, Y, and Z? Or is there a better function to achieve this? Thanks!

Answer 1

Reference: https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.contour.html

You can accomplish this by setting the [levels] parameter in contour([X, Y,] Z, [levels], **kwargs).

You can draw contour lines at the specified levels by giving an array that is in increasing order.

import matplotlib.pyplot as plt
import numpy as np

x = y = np.arange(-3.0, 3.0, 0.02)
X, Y = np.meshgrid(x, y)
Z1 = np.exp(-X ** 2 - Y ** 2)
Z2 = np.exp(-(X - 1) ** 2 - (Y - 1) ** 2)
Z3 = np.exp(-(X + 1) ** 2 - (Y + 1) ** 2)
Z = (Z1 - Z2 - Z3) * 2

fig, ax = plt.subplots()
im = ax.imshow(Z, interpolation='gaussian',
               origin='lower', extent=[-4, 4, -4, 4],
               vmax=abs(Z).max(), vmin=-abs(Z).max())
plt.colorbar(im)

CS = ax.contour(X, Y, Z, levels=[0.9], colors='black')
ax.clabel(CS, fmt='%1.1f', fontsize=12)

plt.show()

Result (levels=[0.9]):

Answer 2

If you leave out X and Y, by default, plt.contour uses the array indices (in this case the range 0-1023 in both x and y).

To only draw a contour line at a given level, you can use levels=[0]. The colors= parameter can fix one or more colors. Optionally, you can draw a line on the colorbar to indicate the value of the level.

import matplotlib.pyplot as plt
import numpy as np
from scipy import ndimage # to smooth a test image

# create a test image with similar properties as the given one

np.random.seed(20221230)
distance = np.pad(np.random.randn(1001, 1001), (11, 11), constant_values=-0.02)
distance = ndimage.filters.gaussian_filter(distance, 100)
distance -= distance.min()
distance = distance / distance.max() * 0.78 - 0.73

plt.imshow(distance)
cbar = plt.colorbar()

level = 0
color = 'red'
plt.contour(distance, levels=[level], colors=color)
cbar.ax.axhline(level, color=color) # show the level on the colorbar
plt.show()