Add function for converting MPAS meshes to viz triangles

Each pair of adjacent vertices on an MPAS cell is joined to the
cell center to create a triangle.  Interpolation weights and cell
indices are computed so that MPAS fields at cell centers can either
be plotted as "flat" to show the individual cell polygons or
"gouraud" to smoothly interpolate between values from the closest
cell centers.

conda_package/mpas_tools/viz/__init__.py

@@ -0,0 +1 @@
+from mpas_tools.viz.mesh_to_triangles import mesh_to_triangles

conda_package/mpas_tools/viz/mesh_to_triangles.py

@@ -0,0 +1,147 @@
+import xarray
+import numpy
+
+
+def mesh_to_triangles(dsMesh, periodicCopy=False):
+    """
+    Construct a dataset in which each MPAS cell is divided into the triangles
+    connecting pairs of adjacent vertices to cell centers.
+
+    Parameters
+    ----------
+    dsMesh : xarray.Dataset
+        An MPAS mesh
+
+    periodicCopy : bool, optional
+        Whether to make a periodic copy of triangles that cross -180/180 degrees
+        longitude.  This is helpful when plotting triangles in a lon/lat space.
+
+    Returns
+    -------
+    dsTris : xarray.Dataset
+        A dataset that defines triangles connecting pairs of adjacent vertices
+        to cell centers as well as the cell index that each triangle is in and
+        cell indices and weights for interpolating data defined at cell centers
+        to triangle nodes.  ``dsTris`` includes variables ``triCellIndices``,
+        the cell that each triangle is part of; ``nodeCellIndices`` and
+        ``nodeCellWeights``, the indices and weights used to interpolate from
+        MPAS cell centers to triangle nodes; Cartesian coordinates ``xNode``,
+        ``yNode``, and ``zNode``; and ``lonNode``` and ``latNode`` in radians.
+        ``lonNode`` is guaranteed to be within 180 degrees of the cell center
+        corresponding to ``triCellIndices``.  Nodes always have a
+        counterclockwise winding.
+
+    """
+    nVerticesOnCell = dsMesh.nEdgesOnCell.values
+    verticesOnCell = dsMesh.verticesOnCell.values - 1
+    cellsOnVertex = dsMesh.cellsOnVertex.values - 1
+
+    kiteAreasOnVertex = dsMesh.kiteAreasOnVertex.values
+
+    nTriangles = numpy.sum(nVerticesOnCell)
+
+    maxEdges = dsMesh.sizes['maxEdges']
+    nCells = dsMesh.sizes['nCells']
+    if dsMesh.sizes['vertexDegree'] != 3:
+        raise ValueError('mesh_to_triangles only supports meshes with '
+                         'vertexDegree = 3')
+
+    # find the third vertex for each triangle
+    nextVertex = -1*numpy.ones(verticesOnCell.shape, int)
+    for iVertex in range(maxEdges):
+        valid = iVertex < nVerticesOnCell
+        invalid = numpy.logical_not(valid)
+        verticesOnCell[invalid, iVertex] = -1
+        nv = nVerticesOnCell[valid]
+        cellIndices = numpy.arange(0, nCells)[valid]
+        iNext = numpy.where(iVertex < nv - 1, iVertex + 1, 0)
+        nextVertex[:, iVertex][valid] = verticesOnCell[cellIndices, iNext]
+
+    valid = verticesOnCell >= 0
+    verticesOnCell = verticesOnCell[valid]
+    nextVertex = nextVertex[valid]
+
+    # find the cell index for each triangle
+    triCellIndices, _ = numpy.meshgrid(numpy.arange(0, nCells),
+                                       numpy.arange(0, maxEdges),
+                                       indexing='ij')
+    triCellIndices = triCellIndices[valid]
+
+    # find list of cells and weights for each triangle node
+    nodeCellIndices = -1*numpy.ones((nTriangles, 3, 3), int)
+    nodeCellWeights = numpy.zeros((nTriangles, 3, 3))
+
+    # the first node is at the cell center, so the value is just the one from
+    # that cell
+    nodeCellIndices[:, 0, 0] = triCellIndices
+    nodeCellWeights[:, 0, 0] = 1.
+
+    # the other 2 nodes are associated with vertices
+    nodeCellIndices[:, 1, :] = cellsOnVertex[verticesOnCell, :]
+    nodeCellWeights[:, 1, :] = kiteAreasOnVertex[verticesOnCell, :]
+    nodeCellIndices[:, 2, :] = cellsOnVertex[nextVertex, :]
+    nodeCellWeights[:, 2, :] = kiteAreasOnVertex[nextVertex, :]
+
+    weightSum = numpy.sum(nodeCellWeights, axis=2)
+    for iNode in range(3):
+        nodeCellWeights[:, :, iNode] = nodeCellWeights[:, :, iNode]/weightSum
+
+    dsTris = xarray.Dataset()
+    dsTris['triCellIndices'] = ('nTriangles', triCellIndices)
+    dsTris['nodeCellIndices'] = (('nTriangles', 'nNode', 'nInterp'),
+                                 nodeCellIndices)
+    dsTris['nodeCellWeights'] = (('nTriangles', 'nNode', 'nInterp'),
+                                 nodeCellWeights)
+
+    # get Cartesian and lon/lat coordinates of each node
+    for prefix in ['x', 'y', 'z', 'lat', 'lon']:
+        outVar = '{}Node'.format(prefix)
+        cellVar = '{}Cell'.format(prefix)
+        vertexVar = '{}Vertex'.format(prefix)
+        coord = numpy.zeros((nTriangles, 3))
+        coord[:, 0] = dsMesh[cellVar].values[triCellIndices]
+        coord[:, 1] = dsMesh[vertexVar].values[verticesOnCell]
+        coord[:, 2] = dsMesh[vertexVar].values[nextVertex]
+        dsTris[outVar] = (('nTriangles', 'nNode'), coord)
+
+    # nothing obvious we can do about triangles containing the poles
+
+    # make sure node longitudes are within 180 degrees of the cell center
+    lonNode = dsTris.lonNode.values
+    lonCell = lonNode[:, 0]
+    copyEast = numpy.zeros(lonCell.shape, bool)
+    copyWest = numpy.zeros(lonCell.shape, bool)
+    for iNode in [1, 2]:
+        mask = lonNode[:, iNode] - lonCell > numpy.pi
+        copyEast = numpy.logical_or(copyEast, mask)
+        lonNode[:, iNode][mask] = lonNode[:, iNode][mask] - 2*numpy.pi
+        mask = lonNode[:, iNode] - lonCell < -numpy.pi
+        copyWest = numpy.logical_or(copyWest, mask)
+        lonNode[:, iNode][mask] = lonNode[:, iNode][mask] + 2*numpy.pi
+    if periodicCopy:
+        dsNew = xarray.Dataset()
+        eastIndices = numpy.nonzero(copyEast)[0]
+        westIndices = numpy.nonzero(copyWest)[0]
+        triIndices = numpy.append(numpy.append(numpy.arange(0, nTriangles),
+                                               eastIndices), westIndices)
+
+        dsNew['triCellIndices'] = ('nTriangles', triCellIndices[triIndices])
+        dsNew['nodeCellIndices'] = (('nTriangles', 'nNode', 'nInterp'),
+                                    nodeCellIndices[triIndices, :, :])
+        dsNew['nodeCellWeights'] = (('nTriangles', 'nNode', 'nInterp'),
+                                    nodeCellWeights[triIndices, :, :])
+        for prefix in ['x', 'y', 'z', 'lat']:
+            outVar = '{}Node'.format(prefix)
+            coord = dsTris[outVar].values
+            dsNew[outVar] = (('nTriangles', 'nNode'), coord[triIndices, :])
+
+        coord = dsTris['lonNode'].values[triIndices, :]
+        eastSlice = slice(nTriangles, nTriangles+len(eastIndices))
+        coord[eastSlice, :] = coord[eastSlice, :] + 2*numpy.pi
+        westSlice = slice(nTriangles+len(eastIndices),
+                          nTriangles + len(eastIndices) + len(westIndices))
+        coord[westSlice, :] = coord[westSlice, :] - 2 * numpy.pi
+        dsNew['lonNode'] = (('nTriangles', 'nNode'), coord)
+        dsTris = dsNew
+
+    return dsTris