Visualizing grid meshes

The purpose of this page is to document strategies to plot/visualize raw mesh files. The most common example is plotting mesh files in “SCRIP” format. These files contain information about the cell vertices, cell centers, and grid area for each cell in a mesh. One way of visualizing such a mesh is to draw polygons based on each cell vertex. The following is sufficient to do this using matplotlib, using a low-resolution MPAS mesh for illustrative purposes:

from xarray import open_dataset
from matplotlib.collections import PatchCollection, PolyCollection
import numpy
from matplotlib import pyplot, patches

# Fix longitudes to go from -180 to 180 instead of 0 to 360
def fix_lon(lon):
    lon = numpy.where(lon > 180, lon - 360, lon)
    return lon

# Wrap coords around edge of projection
# NOTE: this is a dumb way to do this, we should wrap in projection space
def wrap_coords(x, y):
    if (any(x < -90) and any(x > 90)):
        x = numpy.where(x < -90, 360 + x, x)
    if (any(y < -45) and any(y > 45)):   
        y = numpy.where(y < -45, 180 + y, y)
    return x, y

# Plot polygons defined by cell vertices
def plot_scrip(xv, yv, **kwargs):
    
    # Wrap edges
    for i in range(xv.shape[0]):
        xv[i,...], yv[i,...] = wrap_coords(xv[i,...], yv[i,...])
    
    # Create array of cell vertices, indexed [npoints, ncorners, 2]
    corners = numpy.stack([xv, yv], axis=2)
    
    # Create a PatchCollection from our aggregated list of PathPatches
    p = PolyCollection(corners, **kwargs)

    # Add the collection to the axes
    ax = pyplot.gca()
    ax.add_collection(p)

    # Set sane axes limits
    ax.set_xlim([xv.min(), xv.max()])
    ax.set_ylim([yv.min(), yv.max()])
    
    # Return collection of patches
    return p

figure = pyplot.figure(figsize=(20, 10))
ax = figure.add_subplot(111)
with open_dataset(grid_file) as ds:
    pl = plot_scrip(fix_lon(ds['grid_corner_lon']), ds['grid_corner_lat'], facecolor='none', edgecolor='dodgerblue')

This will probably be slow for large grids, since it draws each polygon in a loop, but serves for illustrative purposes. The resulting image for the QU240 grid:

Perhaps a better way to visualize these grids, especially for large grids, would be to plot the grid area. The following will plot native grid data given the scrip cell vertices:

from xarray import open_dataset
from matplotlib import pyplot
from matplotlib.collections import PolyCollection
import numpy
def plot_unstructured(
        xv, yv, data, vmin=None, vmax=None, **kwargs):
    """
    Plot unstructured data. xv and yv specify the x and y coordinates
    of the vertices of each cell and should have shape [ni,nv] where ni
    is the size of the grid (number of cells) and nv is the number of
    vertices for each cell. Data contains the values of the data you
    want to plot, and should have dimension [ni,]. The intent with this
    function is that you will need to read in an auxillary SCRIP-format
    file that describes the grid (unless this information is present in
    the same file that contains the data; unlikely) as well as the file
    that contains the data, and grab the cell corner information from
    the SCRIP file and the data from the data file. This function will
    then plot the data on the native grid by looping over each cell and
    drawing patches for each. Note that this will probably be really
    slow for big grids! Seems to work alright up to about ne120 or so,
    but really some more clever techniques should probably be used here
    (parallelism?).
    
    NOTE: To avoid artifacts due to antialiasing, you should probably pass
    antialiaseds=False to **kwargs.
    """

    # Wrap edges
    for i in range(xv.shape[0]):
        xv[i,...], yv[i,...] = wrap_coords(xv[i,...], yv[i,...])
        
    # Create array of cell vertices, indexed [npoints, ncorners, 2]
    corners = numpy.stack([xv, yv], axis=2)

    # Create a PatchCollection from our aggregated list of PathPatches
    p = PolyCollection(corners, array=data, **kwargs)

    # Set scale, mimicking vmin and vmax plot kwargs
    if vmin is not None and vmax is not None:
        p.set_clim([vmin, vmax])

    # Add the collection to the axes
    ax = pyplot.gca()
    ax.add_collection(p)

    # Set sane axes limits
    ax.set_xlim([xv.min(), xv.max()])
    ax.set_ylim([yv.min(), yv.max()])
    
    # Return collection of patches
    return p

with open_dataset(grid_file) as ds:
    figure = pyplot.figure(figsize=(20, 10))
    ax = figure.add_subplot(111)
    pl = plot_unstructured(fix_lon(ds['grid_corner_lon']), ds['grid_corner_lat'], ds['grid_area'], antialiaseds=False)