Merge pull request #12 from andrewdnolan/improved_patch_creation

- Improve edge/vertex patch creation along culled mesh boundaries
- Replaced numpy masked arrays, with ragged array to better integrate with cartopy
- Improved and optimized support datasets comprised of dask arrays
- Fixed bug in Polar Stereographic projections for spherical meshes

Author: andrewdnolan

mosaic/descriptor.py

@@ -13,7 +13,8 @@
 connectivity_arrays = ["cellsOnEdge",
                        "cellsOnVertex",
                        "verticesOnEdge",
-                       "verticesOnCell"]
+                       "verticesOnCell",
+                       "edgesOnVertex"]
 
 
 class Descriptor:
@@ -48,6 +49,7 @@ def __init__(self, ds, projection=None, transform=None, use_latlon=False):
         self.latlon = use_latlon
         self.projection = projection
         self.transform = transform
+        self._pre_projected = False
 
         # if mesh is on a sphere, force the use of lat lon coords
         if ds.attrs["on_a_sphere"].strip().upper() == 'YES':
@@ -60,6 +62,7 @@ def __init__(self, ds, projection=None, transform=None, use_latlon=False):
         # reproject the minimal dataset, even for non-spherical meshes
         if projection and transform:
             self._transform_coordinates(projection, transform)
+            self._pre_projected = True
 
     def create_minimal_dataset(self, ds):
         """
@@ -95,8 +98,8 @@ def fix_outofbounds_indices(ds, array_name):
         # list of coordinate / connectivity arrays needed to create patches
         mesh_arrays = coordinate_arrays + connectivity_arrays
 
-        # get the subset of arrays from the mesh dataset
-        minimal_ds = ds[mesh_arrays]
+        # get subset of arrays from mesh and load into memory if dask arrays
+        minimal_ds = ds[mesh_arrays].load()
 
         # delete the attributes in the minimal dataset to avoid confusion
         minimal_ds.attrs.clear()
@@ -138,6 +141,17 @@ def vertex_patches(self):
         patches = self._fix_antimeridian(patches, "Vertex")
         return patches
 
+    def get_transform(self):
+        """
+        """
+
+        if self._pre_projected:
+            transform = self.projection
+        else:
+            transform = self.transform
+
+        return transform
+
     def _transform_coordinates(self, projection, transform):
         """
         """
@@ -231,74 +245,128 @@ def transform_patches(self, patches, projection, transform):
 
 
 def _compute_cell_patches(ds):
+    """Create cell patches (i.e. Primary cells) for an MPAS mesh.
 
+    All cell patches have `ds.sizes["maxEdges"]` nodes, even if `nEdgesOnCell`
+    for the given cell is less than maxEdges. We choosed to deal with ragged
+    indices (i.e. node indices greater than `nEdgesOnCell` for a given cell)
+    by repeating the first node of the patch. Nodes are ordered counter
+    clockwise aroudn the cell center.
+    """
+    # get the maximum number of edges on a cell
+    maxEdges = ds.sizes["maxEdges"]
     # connectivity arrays have already been zero indexed
     verticesOnCell = ds.verticesOnCell
     # get a mask of the active vertices
     mask = verticesOnCell == -1
 
-    # get the coordinates needed to patch construction
-    xVertex = ds.xVertex
-    yVertex = ds.yVertex
+    # tile the first vertices index
+    firstVertex = np.tile(verticesOnCell[:, 0], (maxEdges, 1)).T
+    # set masked vertices to the first vertex of the cell
+    verticesOnCell = np.where(mask, firstVertex, verticesOnCell)
 
     # reshape/expand the vertices coordinate arrays
-    x_vert = np.ma.MaskedArray(xVertex[verticesOnCell], mask=mask)
-    y_vert = np.ma.MaskedArray(yVertex[verticesOnCell], mask=mask)
+    x_nodes = ds.xVertex.values[verticesOnCell]
+    y_nodes = ds.yVertex.values[verticesOnCell]
 
-    verts = np.ma.stack((x_vert, y_vert), axis=-1)
+    nodes = np.stack((x_nodes, y_nodes), axis=-1)
 
-    return verts
+    return nodes
 
 
-def _compute_edge_patches(ds, latlon=False):
+def _compute_edge_patches(ds):
+    """Create edge patches for an MPAS mesh.
+
+    Edge patches have four nodes which typically correspond to the two cell
+    centers of the `cellsOnEdge` and the two vertices which make up the edge.
+    For an edge patch along a culled mesh boundary one of the cell centers
+    usually used to construct the patch will be missing, so the corresponding
+    node will be collapsed to the edge coordinate.
+    """
 
     # connectivity arrays have already been zero indexed
     cellsOnEdge = ds.cellsOnEdge
     verticesOnEdge = ds.verticesOnEdge
-
-    # is this masking correct ??
+    # condition should only be true once per row or else wouldn't be an edge
     cellMask = cellsOnEdge < 0
-    vertexMask = verticesOnEdge < 0
-
-    # get the coordinates needed to patch construction
-    xCell = ds.xCell
-    yCell = ds.yCell
-    xVertex = ds.xVertex
-    yVertex = ds.yVertex
 
     # get subset of cell coordinate arrays corresponding to edge patches
-    xCell = np.ma.MaskedArray(xCell[cellsOnEdge], mask=cellMask)
-    yCell = np.ma.MaskedArray(yCell[cellsOnEdge], mask=cellMask)
+    xCell = ds.xCell.values[cellsOnEdge]
+    yCell = ds.yCell.values[cellsOnEdge]
     # get subset of vertex coordinate arrays corresponding to edge patches
-    xVertex = np.ma.MaskedArray(xVertex[verticesOnEdge], mask=vertexMask)
-    yVertex = np.ma.MaskedArray(yVertex[verticesOnEdge], mask=vertexMask)
+    xVertex = ds.xVertex.values[verticesOnEdge]
+    yVertex = ds.yVertex.values[verticesOnEdge]
 
-    x_vert = np.ma.stack((xCell[:, 0], xVertex[:, 0],
-                          xCell[:, 1], xVertex[:, 1]), axis=-1)
+    # if only one cell on edge (i.e. along a culled boundary), then collapse
+    # the node corresponding to the missing cell back the edge location
+    if np.any(cellMask):
+        xCell = np.where(cellMask, ds.xEdge.values[:, np.newaxis], xCell)
+        yCell = np.where(cellMask, ds.yEdge.values[:, np.newaxis], yCell)
 
-    y_vert = np.ma.stack((yCell[:, 0], yVertex[:, 0],
-                          yCell[:, 1], yVertex[:, 1]), axis=-1)
+    x_nodes = np.stack((xCell[:, 0], xVertex[:, 0],
+                        xCell[:, 1], xVertex[:, 1]), axis=-1)
 
-    verts = np.ma.stack((x_vert, y_vert), axis=-1)
+    y_nodes = np.stack((yCell[:, 0], yVertex[:, 0],
+                        yCell[:, 1], yVertex[:, 1]), axis=-1)
 
-    return verts
+    nodes = np.stack((x_nodes, y_nodes), axis=-1)
 
+    return nodes
 
-def _compute_vertex_patches(ds, latlon=False):
 
-    # connectivity arrays have already been zero indexed
-    cellsOnVertex = ds.cellsOnVertex
-    # get a mask of the active vertices
-    mask = ds.cellsOnVertex == -1
+def _compute_vertex_patches(ds):
+    """Create vertex patches (i.e. Dual Cells) for an MPAS mesh.
 
-    # get the coordinates needed to patch construction
-    xCell = ds.xCell
-    yCell = ds.yCell
+    Vertex patches have 6 nodes despite the typical dual cell only having
+    three nodes (i.e. the cell centers of three cells on the vertex) in order
+    ease the creation of vertex patches along culled mesh boundaries.
+    The typical vertex patch will be comprised of the edges and cell centers
+    of the `cellsOnVertex`. As the MPAS Mesh Specs (version 1.0: Section 5.3)
+    outlines: "Edges lead cells as they move around vertex". So, the first node
+    in a vertex patch will correspond to an edge (if present).
 
-    # reshape/expand the vertices coordinate arrays
-    x_vert = np.ma.MaskedArray(xCell[cellsOnVertex], mask=mask)
-    y_vert = np.ma.MaskedArray(yCell[cellsOnVertex], mask=mask)
+    For patches along culled boundaries, if an edge and/or cell center is
+    missing the corresponding node will be collapsed to the patches vertex
+    position.
+    """
+    nVertices = ds.sizes["nVertices"]
+    vertexDegree = ds.sizes["vertexDegree"]
 
-    verts = np.ma.stack((x_vert, y_vert), axis=-1)
+    nodes = np.zeros((nVertices, vertexDegree * 2, 2))
+    # connectivity arrays have already been zero indexed
+    cellsOnVertex = ds.cellsOnVertex.values
+    edgesOnVertex = ds.edgesOnVertex.values
+    # get a mask of active nodes
+    cellMask = cellsOnVertex == -1
+    edgeMask = edgesOnVertex == -1
+    unionMask = cellMask & edgeMask
 
-    return verts
+    # get the coordinates needed to patch construction
+    xCell = ds.xCell.values
+    yCell = ds.yCell.values
+    xEdge = ds.xEdge.values
+    yEdge = ds.yEdge.values
+    # convert vertex coordinates to column vectors for broadcasting below
+    xVertex = ds.xVertex.values[:, np.newaxis]
+    yVertex = ds.yVertex.values[:, np.newaxis]
+
+    # if edge is missing collapse edge node to vertex, else leave at edge
+    nodes[:, ::2, 0] = np.where(edgeMask, xVertex, xEdge[edgesOnVertex])
+    nodes[:, ::2, 1] = np.where(edgeMask, yVertex, yEdge[edgesOnVertex])
+
+    # if cell is missing collapse cell node to vertex, else leave at cell
+    nodes[:, 1::2, 0] = np.where(cellMask, xVertex, xCell[cellsOnVertex])
+    nodes[:, 1::2, 1] = np.where(cellMask, yVertex, yCell[cellsOnVertex])
+
+    # -------------------------------------------------------------------------
+    # NOTE: While the condition below probably only applies to the final edge
+    #       node we apply it to all, since the conditions above ensure the
+    #       patches will still be created correctly
+    # -------------------------------------------------------------------------
+    # if cell and edge missing collapse edge node to the first edge.
+    # Because edges lead the vertices this ensures the patch encompasses
+    # the full kite area and is propely closed.
+    nodes[:, ::2, 0] = np.where(unionMask, nodes[:, 0:1, 0], nodes[:, ::2, 0])
+    nodes[:, ::2, 1] = np.where(unionMask, nodes[:, 0:1, 1], nodes[:, ::2, 1])
+
+    return nodes

mosaic/polypcolor.py

@@ -1,3 +1,4 @@
+from cartopy.mpl.geoaxes import GeoAxes
 from matplotlib.axes import Axes
 from matplotlib.collections import PolyCollection
 from matplotlib.colors import Normalize
@@ -56,34 +57,23 @@ def polypcolor(
     elif "nVertices" in c.dims:
         verts = descriptor.vertex_patches
 
+    transform = descriptor.get_transform()
+
     collection = PolyCollection(verts, alpha=alpha, array=c,
                                 cmap=cmap, norm=norm, **kwargs)
 
-    collection._scale_norm(norm, vmin, vmax)
-
-    # get the limits of the **data**, which could exceed the valid
-    # axis limits of the transform
-    minx = verts[..., 0].min()
-    maxx = verts[..., 0].max()
-    miny = verts[..., 1].min()
-    maxy = verts[..., 1].max()
+    # only set the transform if GeoAxes
+    if isinstance(ax, GeoAxes):
+        collection.set_transform(transform)
 
-    corners = (minx, miny), (maxx, maxy)
-    ax.update_datalim(corners)
-    ax._request_autoscale_view()
+    collection._scale_norm(norm, vmin, vmax)
     ax.add_collection(collection)
 
-    # if data has been transformed use the transforms x-limits.
-    # patches have vertices that exceed the transfors x-limits to visually
-    # correct the antimeridian problem
-    if descriptor.projection:
-        minx = descriptor.projection.x_limits[0]
-        maxx = descriptor.projection.x_limits[1]
-
-        miny = descriptor.projection.y_limits[0]
-        maxy = descriptor.projection.y_limits[1]
-
-        ax.set_xbound(minx, maxx)
-        ax.set_ybound(miny, maxy)
+    # Update the datalim for this polycollection
+    limits = collection.get_datalim(ax.transData)
+    # TODO: account for nodes of patches that lie outside of projection bounds
+    #       (i.e. as a result of patch wrapping at the antimeridian)
+    ax.update_datalim(limits)
+    ax.autoscale_view()
 
     return collection