diff --git a/conda_package/mpas_tools/viz/__init__.py b/conda_package/mpas_tools/viz/__init__.py index e69de29bb..2e6d19f68 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 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..2b65da8d0 --- /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', 'nNodes', 'nInterp'), + nodeCellIndices) + dsTris['nodeCellWeights'] = (('nTriangles', 'nNodes', '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', 'nNodes'), 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', 'nNodes', 'nInterp'), + nodeCellIndices[triIndices, :, :]) + dsNew['nodeCellWeights'] = (('nTriangles', 'nNodes', 'nInterp'), + nodeCellWeights[triIndices, :, :]) + for prefix in ['x', 'y', 'z', 'lat']: + outVar = '{}Node'.format(prefix) + coord = dsTris[outVar].values + dsNew[outVar] = (('nTriangles', 'nNodes'), 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', 'nNodes'), coord) + dsTris = dsNew + + return dsTris