Search code examples
pythonnumpymatplotlibspatialcontour

Find contour lines from matplotlib.pyplot.contour()


I'm trying to find (but not draw!) contour lines for some data:

from pprint import pprint 
import matplotlib.pyplot 
z = [[0.350087, 0.0590954, 0.002165], [0.144522, 0.885409, 0.378515], 
     [0.027956, 0.777996, 0.602663], [0.138367, 0.182499, 0.460879], 
     [0.357434, 0.297271, 0.587715]] 
cn = matplotlib.pyplot.contour(z) 

I know cn contains the contour lines I want, but I can't seem to get to them. I've tried several things:

print dir(cn) 
pprint(cn.collections[0]) 
print dir(cn.collections[0]) 
pprint(cn.collections[0].figure) 
print dir(cn.collections[0].figure) 

to no avail. I know cn is a ContourSet, and cn.collections is an array of LineCollections. I would think a LineCollection is an array of line segments, but I can't figure out how to extract those segments.

My ultimate goal is to create a KML file that plots data on a world map, and the contours for that data as well.

However, since some of my data points are close together, and others are far away, I need the actual polygons (linestrings) that make up the contours, not just a rasterized image of the contours.

I'm somewhat surprised qhull doesn't do something like this.

Using Mathematica's ListContourPlot and then exporting as SVG works, but I want to use something open source.

I can't use the well-known CONREC algorithm because my data isn't on a mesh (there aren't always multiple y values for a given x value, and vice versa).

The solution doesn't have to python, but does have to be open source and runnable on Linux.


Solution

  • You can get the vertices back by looping over collections and paths and using the iter_segments() method of matplotlib.path.Path.

    Here's a function that returns the vertices as a set of nested lists of contour lines, contour sections and arrays of x,y vertices:

    import numpy as np
    
    def get_contour_verts(cn):
        contours = []
        # for each contour line
        for cc in cn.collections:
            paths = []
            # for each separate section of the contour line
            for pp in cc.get_paths():
                xy = []
                # for each segment of that section
                for vv in pp.iter_segments():
                    xy.append(vv[0])
                paths.append(np.vstack(xy))
            contours.append(paths)
    
        return contours
    

    Edit:

    It's also possible to compute the contours without plotting anything using the undocumented matplotlib._cntr C module:

    from matplotlib import pyplot as plt
    from matplotlib import _cntr as cntr
    
    z = np.array([[0.350087, 0.0590954, 0.002165],
                  [0.144522,  0.885409, 0.378515],
                  [0.027956,  0.777996, 0.602663],
                  [0.138367,  0.182499, 0.460879], 
                  [0.357434,  0.297271, 0.587715]])
    
    x, y = np.mgrid[:z.shape[0], :z.shape[1]]
    c = cntr.Cntr(x, y, z)
    
    # trace a contour at z == 0.5
    res = c.trace(0.5)
    
    # result is a list of arrays of vertices and path codes
    # (see docs for matplotlib.path.Path)
    nseg = len(res) // 2
    segments, codes = res[:nseg], res[nseg:]
    
    fig, ax = plt.subplots(1, 1)
    img = ax.imshow(z.T, origin='lower')
    plt.colorbar(img)
    ax.hold(True)
    p = plt.Polygon(segments[0], fill=False, color='w')
    ax.add_artist(p)
    plt.show()
    

    enter image description here