Python get contour line vertices#
Software requirements:
Python 3
numpy
matplotlib
cartopy
Example script#
get_contour_line_vertices.py
#!/usr/bin/env python
# coding: utf-8
'''
DKRZ example
Contour line plot: Work with the coordinates of drawn contour lines
Description
In some cases, the contour line labels are positioned unfavorably so that
they are no longer readable.
In order to be able to add an additional contour line label, we need to go
a little deeper and first understand how to get the corresponding contour
line data.
Content
- _QuadContourSet_ object
- levels
- collections
- allsegs
- allkinds
- Vertices
- draw additional labels
- skip selected contour lines
- suppress/remove contour lines
-------------------------------------------------------------------------------
2024 copyright DKRZ licensed under CC BY-NC-SA 4.0 <br>
(https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en)
-------------------------------------------------------------------------------
'''
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
#--------------------------------------------------------------------------
# Function: add_contour_label()
#
# The following function combines everything we have worked out above. It
# need the axis, plot name, and a contour vale, and other parameter.
#--------------------------------------------------------------------------
def add_contour_label(ax, plot, cnvalue, ymin=5000, factor=2, color='red',
bbox=True, info=True, nowarning=False):
#-- get the index of a given contour line value cnvalue
index = np.where(plot.levels == cnvalue)[0][0]
#-- get the path from the plot PathCollection
cnpath = [ p.get_paths() for i, p in enumerate(plot.collections) if i == index ]
if len(cnpath[0]) != 0:
#-- get the vertices from the Path object
cnvertices = cnpath[0][0].vertices
#-- label position
pos = int(cnvertices.shape[0] / factor)
x = cnvertices[pos, 0]
y = cnvertices[pos, 1]
if (y > 0) and (y < ymin):
if nowarning == False:
print(f'add_contour_label: WARNING -- no segments found for cnvalue = {cnvalue}')
return
##- draw label text
t = ax.text(x, y, int(cnvalue), fontsize=10, weight='bold',
ha='center', va='center', color=color)
if bbox: t.set_bbox(dict(facecolor='yellow', edgecolor='black', alpha=1))
if info:
print(f'add_contour_label: cnvalue = {cnvalue}')
print(f'add_contour_label: index = {index}')
print(f'add_contour_label: n-paths = {len(cnpath[0][0])}')
print(f'add_contour_label: extent = {cnpath[0][0].get_extents()}')
print(f'add_contour_label: vertices = {cnvertices.shape}')
print(f'add_contour_label: pos = {pos}')
return t
else:
if nowarning == False:
print(f'add_contour_label: WARNING - no segments found for cnvalue={cnvalue}')
return
#------------------------------------------------------------------------------
# Main
#------------------------------------------------------------------------------
def main():
infile = '../../data/ua_Amon_t1.nc'
ds = xr.open_dataset(infile)
data = ds.ua.isel(time=0).sel(lat=47., method='nearest')
#-- Create contour plot
#
# A good example is a longitude vs. pressure level slice, as the contour lines
# are very close together in the upper pressure levels and automatically drawn
# labels are no longer readable.
plt.switch_backend('agg')
fig, ax = plt.subplots(figsize=(6,3))
ax.invert_yaxis()
cnplot = ax.contour(ds.lon, ds.plev, data,
levels=20,
colors='black',
linewidths=0.8)
plt.clabel(cnplot, fontsize=10, inline=False, fmt = '%1.0f', colors='black')
#plt.show()
#-- What returns `ax.contour`
#
# The call of **ax.contour()** returns a **matplotlib.contour.QuadContourSet**
# object, here assigned to the variable _cnplot_.
print('ax.contour returns an object: ', cnplot)
#-- The object _cnplot_ has the following attributes:
#
# levels, collections, allsegs, allkinds
#
# levels
#
# The levels attribute gives you all used contour level values.
#
# cnplot.levels = [level0, level1, ...]
cnlevels = cnplot.levels
print(cnlevels)
#-- collections
#
# The collections attribute stored the PathCollection object, that contains the
# paths for each contour line.
#
# cmplot.collections = PathCollection object
print(cnplot.collections)
#-- allsegs
#
# The allsegs attribute can be used to retrieve the vertices for each contour line.
#
# cnplot.allsegs = [level0segs, level1segs, ...]
#
# cnplot.allsegs[level][element][vertex coordinates]
#print(cnplot.allsegs)
#-- allkinds
#
# cnplot.allkinds = [level0kinds, level1kinds, ...]
PRINT = False
if PRINT:
for level_index in range(len(cnlevels)):
print(f'level value: {cnlevels[level_index]}')
print(f'level index: {level_index}')
num_levels = len(cnplot.allsegs)
print(f' number of levels: {num_levels}')
num_element = len(cnplot.allsegs[level_index]) #-- in level 0
print(f' number of elements of level index {level_index}: {num_element}')
if num_element != 0:
num_vertices = len(cnplot.allsegs[level_index][0]) #-- of element 0, in level 0
print(f' number of vertices: {num_vertices}')
num_coord = len(cnplot.allsegs[level_index][0][0]) #-- of vertex 0, in element 0, in level 0
print(f' number of coordinates: {num_coord}')
print('-----------------------------------------------------')
if PRINT:
print(cnplot.allkinds)
#-- Plot contour line by line
#
# You can select one or more contour lines by using the _allsegs_ segments.
#
# The next code creates two figures where the upper figure shows all contour lines,
# and the lower only shows the first 8 contour lines.
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(6,4), layout='tight')
for i, ax in enumerate(axes.flat):
ax.set_ylim([-10000, 100000])
ax.invert_yaxis()
for j in range(len(cnplot.levels)):
segments = cnplot.allsegs[j]
if (i == 0):
for seg in segments:
ax.plot(seg[:,0], seg[:,1], ls='-', lw=1)
ax.xaxis.set_tick_params(labelbottom=False)
elif (i == 1) and (j < 8):
for seg in segments:
ax.plot(seg[:,0], seg[:,1], ls='-', lw=1)
#plt.show()
#-- Vertices
print(cnplot.collections)
print(cnplot.collections[0])
PRINT = False
if PRINT:
print(cnplot.collections[1].get_paths())
#-- choose contour value
cnvalue = -6
#-- get the index of a given contour line value
index = np.where(cnplot.levels == cnvalue)[0][0]
#cnvalue = cnplot.levels[index]
#-- get the path from the PathCollection
cnpath = [ p.get_paths() for i, p in enumerate(cnplot.collections) if i == index ]
#-- get the vertices from the Path object
cnvertices = cnpath[0][0].vertices
print(f'Contour value: contour_value = {cnvalue}\n',
f' index = {index}\n',
f' cnplot.levels[{index}] = {cnplot.levels[index]}')
print('Number of vertices: ', len(cnpath[0][0]))
print('Bbox (extend): ', cnpath[0][0].get_extents())
print('Vertices shape: ', cnvertices.shape)
#-- Draw additional label
#
# To add an additional label to the contour line of value -6 from above, we have
# to choose the postion without knowing where. Therefore, we take the coordinates
# of the contour line center, whose index can be obtained using the vertices shape
# `cnvertices.shape[0] / 2`.
#-- get index of the point in the middle of the contour line
pos = int(cnvertices.shape[0] / 2)
#-- have a look at the coordinates of the middle point
x = cnvertices[pos, 0]
y = cnvertices[pos, 1]
print(f'First coordinates of the contour line: x={x:.4f}, y={y:.4f}')
# Now that we have the coordinates of the center position, we can add a text
# label with a yellow bounding box to the plot.
fig, ax = plt.subplots(figsize=(10,5))
ax.set_ylim([-10000, 100000])
ax.invert_yaxis()
cnplot = ax.contour(ds.lon, ds.plev, data,
levels=20,
colors='black',
linewidths=0.8)
plt.clabel(cnplot, fontsize=10, inline=False, fmt = '%1.0f', colors='black')
t = ax.text(x, y, int(cnvalue), fontsize=10, weight='bold', color='red')
t.set_bbox(dict(facecolor='yellow', edgecolor='black', alpha=1))
#plt.show()
#-- Add additional labels with different locations and colors
#
# And now we want to test the function and play with different contour line values
# and colors.
fig, ax = plt.subplots(figsize=(10,5))
ax.set_ylim([-10000, 100000])
ax.invert_yaxis()
cnplot = plt.contour(ds.lon, ds.plev, data,
levels=20,
colors='black',
linewidths=0.8)
plt.clabel(cnplot, fontsize=10, inline=False, fmt = '%1.0f', colors='black')
infoargs = dict(info=False, nowarning=True)
for i in range(len(cnplot.levels)):
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=4., color='red', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=2., color='blue', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=1.1, color='green', **infoargs)
#plt.show()
#-- Draw only selected contour labels
#
# This time we do not want to have all contour lines drawn automatically, but
# only a few selected ones. This makes it easier to see what the factor parameter does.
fig, ax = plt.subplots(figsize=(10,5))
ax.set_ylim([-10000, 100000])
ax.invert_yaxis()
cnplot = plt.contour(ds.lon, ds.plev, data,
levels=20,
colors='black',
linewidths=0.8)
infoargs = dict(info=False, nowarning=True)
for i in range(len(cnplot.levels)):
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=4., color='red', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=2., color='blue', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=1.1, color='green', **infoargs)
#plt.show()
#-- save graphic output to PNG file
fig.savefig('plot_get_contour_lines.png', bbox_inches='tight', facecolor='white')
#-- Draw selected contour labels suppress drawing the lines
#
# The contour lines are drawn in the same way as above, but this time we don't
# want to see any contour lines in the plot. Thanks to the `cnplot.collection`,
# you can delete every single contour line at the end.
fig, ax = plt.subplots(figsize=(10,5))
ax.set_ylim([-10000, 100000])
ax.invert_yaxis()
cnplot = plt.contour(ds.lon, ds.plev, data,
levels=20,
colors='black',
linewidths=0.8)
infoargs = dict(info=False, nowarning=True)
for i in range(len(cnplot.levels)):
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=4., color='red', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=2., color='blue', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=1.1, color='green', **infoargs)
#-- remove the contour lines
for coll in cnplot.collections:
coll.remove()
#plt.show()
#-- Close the plot
#
# This will delete the figure but the cnplot data will be kept.
plt.close()
# Use `cnplot` data after closing the figure
#
# If you close a figure, the content of the figure gets lost, but the assigned
# data to `cnplot` still exists.
cnplot = plt.contour(ds.lon, ds.plev, data,
levels=20,
colors='black',
linewidths=0.8)
cnpaths = [ p.get_paths() for i, p in enumerate(cnplot.collections)]
plt.close() #-- this suppress the contour line drawing
# To get the same plot as above, we need to set the x-axis limit, but then we
# can create the plot the same way.
fig, ax = plt.subplots(figsize=(10,5))
ax.set_xlim([0, 360])
ax.set_ylim([-10000, 100000])
ax.invert_yaxis()
infoargs = dict(info=False, nowarning=True)
for i in range(len(cnplot.levels)):
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=4., color='red', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=2., color='blue', **infoargs)
add_contour_label(ax, cnplot, cnplot.levels[i], ymin=2000, factor=1.1, color='green', **infoargs)
#plt.show()
if __name__ == '__main__':
main()
Plot result:
