From 99ef287f85a2321f88039d51f9d9a29b38e36a42 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 4 Jul 2020 20:46:51 +0200 Subject: [PATCH] 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 | 1 + .../mpas_tools/viz/mesh_to_triangles.py | 147 ++++++++++++++++++ 2 files changed, 148 insertions(+) create mode 100644 conda_package/mpas_tools/viz/mesh_to_triangles.py diff --git a/conda_package/mpas_tools/viz/__init__.py b/conda_package/mpas_tools/viz/__init__.py index e69de29bb..d297b5536 100644 --- a/conda_package/mpas_tools/viz/__init__.py +++ b/conda_package/mpas_tools/viz/__init__.py @@ -0,0 +1 @@ +from mpas_tools.viz.mesh_to_triangles import mesh_to_triangles \ No newline at end of file diff --git a/conda_package/mpas_tools/viz/mesh_to_triangles.py b/conda_package/mpas_tools/viz/mesh_to_triangles.py new file mode 100644 index 000000000..3a1a3bda5 --- /dev/null +++ b/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