Visualize MPAS mesh in python

[bradyrx]

I've been working on a method to visualize MPAS-O output on a static plot in python as an alternative to the more time-consuming route of rendering viz in ParaView. I'm looking for a way to faithfully represent the grid without regridding.

It seems like datashader/holoviews/geoviews is the way to go, but I'm having trouble getting it to work fully. I'm wondering if anyone has used this route or has a solution.

Method 1: Datashader points and shade:

from colorcet import cwr
import datashader 
import datashader.transfer_functions as tf


# Single time slice of air-sea CO2 flux.
ds = xr.open_dataset('MPAS-O.fgco2.nc')
verts = pd.DataFrame({'x': ds.lonCell, 'y': ds.latCell, 'z': ds.fgco2})

cvs = datashader.Canvas(plot_height=400,plot_width=400)
tf.shade(cvs.points(verts, 'x', 'y', agg=datashader.mean('z')), cmap=cwr,
         )

This produces a map by coloring each pixel with an RGB value. This isn't a faithful representation of the grid, and I haven't found an easy way to plot this on a map projection.

Method 2: Trimesh

This is ultimately what I'd like to get working. This can easily be drawn on a map projection using geoviews once it works.

import datashader
import datashader.transfer_functions as tf
import pandas as pd
import xarray as xr

ds = xr.open_dataset('MPAS-O.fgco2.nc')
mesh = xr.open_dataset('oRRS30to10v3.171128.nc')

# Set up table with coordinates for each `nCell` index as well as the data being plotted ('z')
verts = pd.DataFrame({'x': ds.lonCell, 'y': ds.latCell, 'z': ds.fgco2})

# Construct indices for each triangle.
cell0 = mesh.nCells  # Base cell
cell1 = mesh.cellsOnCell.isel(maxEdges=1) - 1  # NE cell
cell2 = mesh.cellsOnCell.isel(maxEdges=2) - 1  # N cell

# Points to the triangle vertices (`verts`) for each triangle being plotted.
tris = pd.DataFrame({'v0': cell0, 'v1': cell1, 'v2': cell2})

cvs = datashader.Canvas(plot_height=400, plot_width=400)
tf.shade(cvs.trimesh(verts, tris))

I tried also including all possible triangles, without success.

# Pairs of neighboring cell indices to construct triangles from.
PAIRS = [
    (0,1),
    (1,2),
    (2,3),
    (3,4),
    (4,5),
    (5,0)
]

cell0 = mesh.nCells
df = pd.DataFrame()
for idx in PAIRS:
    cell1 = mesh.cellsOnCell.isel(maxEdges=idx[0]) - 1
    cell2 = mesh.cellsOnCell.isel(maxEdges=idx[1]) - 1
    tris = pd.DataFrame({'v0': cell0,
                         'v1': cell1,
                         'v2': cell2})
    df = pd.concat([df, tris], ignore_index=True)

cvs = datashader.Canvas(plot_height=400,plot_width=400)
tf.shade(cvs.trimesh(verts, df), cmap=cwr)

Any thoughts? The tris argument should be a three-column list of pointers to indices in verts to form a non-overlapping set of all triangles on the mesh.

Resources

https://datashader.org/user_guide/Trimesh.html

https://holoviews.org/reference/elements/bokeh/TriMesh.html

https://twitter.com/oceanographer/status/1126803420579545094

[bradyrx]

CC @xylar, @pwolfram

Also I can host these example files on a Google Bucket or something. This should be straightforward to assess with any arbitrary MPAS-O variable on any mesh.

[xylar]

@bradyrx, I'd love to look at this more sometime soon but don't have time today. For now, I figured I'd point you to some viz stuff I have for another context that might be helpful. Here, I create a patch collection to reuse for many plots:
https://github.com/MPAS-Dev/MPAS-Model/blob/3e4d773acc4e4c76627124e1319e921dee8a908d/testing_and_setup/compass/ocean/isomip_plus/viz/plot.py#L667-L687
Here is an example of plotting that patch collection:
https://github.com/MPAS-Dev/MPAS-Model/blob/3e4d773acc4e4c76627124e1319e921dee8a908d/testing_and_setup/compass/ocean/isomip_plus/viz/plot.py#L583-L614
This method shows hexagons (etc.) of constant shade, which is what I wanted. Here's an example plot with apologies for the rainbow color map:

[bradyrx]

Thanks @xylar for giving this some thought and your working solution here. I'll give that a try as a potential avenue, but it would be good for us to find some solutions with e.g. holoviews to add to user documentation. I'll try your solution when I have a moment, but need to focus on analysis for now.

Here's a consolidated notebook with my attempts and the example output pulled from Google Storage. I have some running threads on gitter/twitter with folks familiar with this sort of thing so I'll point them here.

Notebook: https://nbviewer.jupyter.org/gist/bradyrx/fe329f47e9a85fb0adacefa113178b44

[bradyrx]

I've updated it with a working version. I had to drop -1 partial cells. This still isn't exactly what I want but is getting in the right direction:

https://gist.github.com/bradyrx/fe329f47e9a85fb0adacefa113178b44

[xylar]

@bradyrx, it looks like you're essentially plotting the dual mesh (with cell centers as the vertices of triangles), is that correct? That's also what ParaView does but it drives me nuts for the reasons that I think you have found, too, that it leaves boundary gaps in the mesh that, in my view, shouldn't be there. I'd personally only be interested in pursuing a viz approach on the native grid if it plots cells as full polygons rather than the triangles on the dual mesh.

[milenaveneziani]

just a quick note to say that I am following this and that it is great you are working on it @bradyrx.
I also do not like to tinker with paraview too much, so my simple solution at the moment is to just do a scatter plot..

[bradyrx]

@bradyrx, it looks like you're essentially plotting the dual mesh (with cell centers as the vertices of triangles), is that correct? That's also what ParaView does but it drives me nuts for the reasons that I think you have found, too, that it leaves boundary gaps in the mesh that, in my view, shouldn't be there.

Correct -- I'm following this approach since TriMesh is the only unstructured mesh option I've found in plotting packages so far. But I agree that it's not perfect.

I'd personally only be interested in pursuing a viz approach on the native grid if it plots cells as full polygons rather than the triangles on the dual mesh.

Agreed. In short, I want a static solution in python to plot the full polygons as we do in ParaView. The VTK extractor you wrote is awesome, but requires spending time to work with the GUI. In theory you could write some macros or use a VTK package on python, though. But I also stylistically would like to be able to leverage something like matplotlib/cartopy to make plots for papers for instance. It's really nice to use the matplotlib API to zoom into a region, tile a bunch of small multiples, add the coastline, colorbar, etc. This is a lot of manual work to make it look nice from ParaView. (But note I use the ParaView route a lot to make nice movies, etc. and it will continue to be useful)

Glad to hear it @milenaveneziani! Hopefully we can find a path forward here. CC @andsery005 who is working on this at NCAR for MPAS-A as well.

[milenaveneziani]

Yes, I agree that it would be great to use cartopy. Here is a zoom-in scatter plot I did recently for the Southern Ocean, using cartopy. The only thing missing is the labeling of meridians and parallels, but that is something that is coming up in the next released version of cartopy:

[bradyrx]

@milenaveneziani, that looks great. I think that's probably the best solution until we can find one that faithfully depicts the polygons. I had used scatter in the past for global checks, but hadn't tried zooming in. I imagine you just tweak the marker size until it doesn't overlap?

[milenaveneziani]

I imagine you just tweak the marker size until it doesn't overlap?

exactly, and that's the problem. For one, every time you plot a different mesh, you have to re-do the tweaking. For two, the rendering isn't uniform, of course, and there is bound to be an overlap of cells/dots where resolution is higher.
So, just a temporary solution for sure.

[bradyrx]

I spent a little bit of time with @xylar's patches approach:

import copy

import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection

import numpy
import xarray as xr

def _compute_cell_patches(dsMesh, cmap):
    patches = []
    nVerticesOnCell = dsMesh.nEdgesOnCell.values
    verticesOnCell = dsMesh.verticesOnCell.values - 1
    xVertex = dsMesh.xVertex.values
    yVertex = dsMesh.yVertex.values
    for iCell in range(dsMesh.sizes['nCells']):
        nVert = nVerticesOnCell[iCell]
        vertexIndices = verticesOnCell[iCell, :nVert]
        vertices = numpy.zeros((nVert, 2))
        vertices[:, 0] = 1e-3*xVertex[vertexIndices]
        vertices[:, 1] = 1e-3*yVertex[vertexIndices]

        polygon = Polygon(vertices, True)
        patches.append(polygon)

    p = PatchCollection(patches, cmap=cmap, alpha=1.)

    return p

dsMesh = xr.open_dataset('oRRS30to10v3.171128.nc')
ds = xr.open_zarr(fsspec.get_mapper('gcs://unstructured_mesh/mpas-o.fgco2')).load()

# Runtime ~2.5min
oceanPatches = _compute_cell_patches(dsMesh, 'RdBu')
# Copy since plotting seems to do some overwriting.
patches = copy.copy(oceanPatches)

# Set up specifically for fgco2
patches.set_array(ds.fgco2)
patches.set_edgecolor('face')
patches.set_clim(vmin=-0.0001, vmax=0.0001)

# plot
plt.figure(figsize=[9, 3])
ax = plt.subplot('111')
ax.add_collection(patches)
plt.colorbar(patches)
plt.axis([0, 500, 0, 1000])
ax.set_aspect('equal')
ax.autoscale(tight=True)

This looks like it's going down the right road, but it does take 2.5 minutes to render the patch collection. Perhaps it can be saved out and reloaded for the viz? Also I need to figure out how to encode the lat/lon information in the patches collection so I can project, zoom in, etc.

[xylar]

Okay, thanks for looking into it. That does sound too slow to be useful. I've only used it for idealized, small test cases.

[milenaveneziani]

as @bradyrx says, we could save it though, right? Since it's only mesh dependent. 2.5min for the 30to10 doesn't sound like much to me if it's done only once.