Merge pull request #57 from andrewdnolan/coastlines

Coastlines (outlines) based on MPAS mesh connectivity info

Author: andrewdnolan

docs/developers_guide/api.md

@@ -16,6 +16,7 @@ Plotting
     :toctree: generated/
 
     polypcolor
+    coastlines
     contour
     contourf
 

docs/user_guide/quick_start.md

@@ -48,26 +48,25 @@ fig, ax = plt.subplots(
     subplot_kw={"projection": projection},
 )
 
-cmap = cmocean.tools.crop(cmocean.cm.topo, -5e3, 0, 0.0)
-
 # create a `Descriptor` object which takes the mesh information and creates
 # the polygon coordinate arrays needed for `matplotlib.collections.PolyCollection`.
 descriptor = mosaic.Descriptor(ds, projection, transform)
 
 # using the `Descriptor` object we just created, make a pseudocolor plot of
-# the "indexToCellID" variable, which is defined at cell centers.
+# the "bottomDepth" variable, which is defined at cell centers.
 collection = mosaic.polypcolor(
-    ax, descriptor, -ds.bottomDepth, antialiaseds=True, cmap=cmap
+    ax, descriptor, ds.bottomDepth, antialiaseds=True, cmap=cmocean.cm.deep
 )
 
-ax.coastlines(lw=0.5)
-ax.add_feature(cfeature.LAND, fc="grey", zorder=-1, alpha=0.8)
+# plot the coastlines as resolved by the MPAS mesh
+mosaic.coastlines(ax, descriptor)
+
 fig.colorbar(
     collection,
     fraction=0.1,
     shrink=0.4,
     extend="both",
-    label="Topography [m a.s.l.]",
+    label="Bottom Depth [m]",
 );
 ```
 

mosaic/__init__.py

@@ -1,8 +1,16 @@
 from __future__ import annotations
 
 from mosaic import datasets
+from mosaic.coastlines import coastlines
 from mosaic.contour import contour, contourf
 from mosaic.descriptor import Descriptor
 from mosaic.polypcolor import polypcolor
 
-__all__ = ["Descriptor", "contour", "contourf", "datasets", "polypcolor"]
+__all__ = [
+    "Descriptor",
+    "coastlines",
+    "contour",
+    "contourf",
+    "datasets",
+    "polypcolor",
+]

mosaic/coastlines.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import warnings
+
+import cartopy.feature as cfeature
+import numpy as np
+import shapely
+from cartopy.mpl.geoaxes import GeoAxes
+
+from mosaic.contour import ContourGraph, MPASContourGenerator
+from mosaic.descriptor import Descriptor
+
+
+def coastlines(
+    ax: GeoAxes, descriptor: Descriptor, color: str = "black", **kwargs
+) -> None:
+    """
+    Plot coastal **outlines** using the connectivity info from the MPAS mesh
+
+    Parameters
+    ----------
+    ax : cartopy.mpl.geoaxes.GeoAxes
+        The cartopy axes to add the coastlines to.
+    descriptor : Descriptor
+        The descriptor containing the projection and dataset information.
+    **kwargs
+       Additional keyword arguments. See
+       :py:class:`matplotlib.collections.Collection` for supported options.
+    """
+
+    if not isinstance(ax, GeoAxes):
+        msg = (
+            "Must provide a `cartopy.mpl.geoaxes` instance for "
+            "`mosaic.coastlines` to work. "
+        )
+        raise TypeError(msg)
+
+    if "edgecolor" in kwargs and "ec" in kwargs:
+        msg = "Cannot specify both 'edgecolor' and 'ec'."
+        raise TypeError(msg)
+    if "edgecolor" in kwargs:
+        color = kwargs.pop("edgecolor")
+    elif "ec" in kwargs:
+        color = kwargs.pop("ec")
+
+    if "facecolor" in kwargs and "fc" in kwargs:
+        msg = "Cannot specify both 'facecolor' and 'fc'."
+        raise TypeError(msg)
+    if "facecolor" in kwargs or "fc" in kwargs:
+        warnings.warn(
+            "'facecolor (fc)' is not supported for `mosaic.coastlines` "
+            "and will be ignored.",
+            stacklevel=2,
+        )
+        kwargs.pop("facecolor", None)
+        kwargs.pop("fc", None)
+
+    kwargs["edgecolor"] = color
+    kwargs["facecolor"] = "none"
+
+    generator = MPASCoastlineGenerator(descriptor)
+    coastlines = generator.create_coastlines()
+
+    geometries = shapely.GeometryCollection(
+        [shapely.LineString(cl) for cl in coastlines]
+    )
+
+    feature = cfeature.ShapelyFeature(geometries, descriptor.projection)
+    ax.add_feature(feature, **kwargs)
+
+
+class MPASCoastlineGenerator(MPASContourGenerator):
+    def __init__(self, descriptor: Descriptor):
+        # pass a dummy field to the parent class
+        super().__init__(descriptor, descriptor.ds.nCells)
+
+        self.domain = descriptor.projection.domain
+        self.boundary = descriptor.projection.boundary
+
+        shapely.prepare(self.domain)
+
+    def create_coastlines(self) -> list[np.ndarray]:
+        graph = self._create_coastline_graph()
+        lines = self._split_and_order_graph(graph)
+
+        return self._snap_lines_to_boundary(lines)
+
+    def _create_coastline_graph(self) -> ContourGraph:
+        edge_mask = (self.ds.cellsOnEdge == -1).any("TWO").values
+
+        vertex_1 = self.ds.verticesOnEdge[edge_mask].isel(TWO=0).values
+        vertex_2 = self.ds.verticesOnEdge[edge_mask].isel(TWO=1).values
+
+        return ContourGraph(vertex_1, vertex_2)
+
+    def _snap_lines_to_boundary(
+        self, lines: list[np.ndarray]
+    ) -> list[np.ndarray]:
+        def snap(point: np.ndarray):
+            return self.boundary.interpolate(
+                self.boundary.project(shapely.Point(point))
+            )
+
+        complete_lines = []
+        for line in lines:
+            # only need to snap lines that are not already closed loops
+            if np.array_equal(line[0], line[-1]):
+                complete_lines.append(line)
+                continue
+
+            contain_mask = shapely.contains_xy(self.domain, *line.T)
+            if not contain_mask.any():
+                continue
+
+            clipped = line[contain_mask]
+
+            if len(clipped) == 1:
+                # if only one point inside domain,
+                # all snapped points will lie along the same line
+                continue
+
+            # TODO: For coastlines with end points outside domain it would be
+            # better to cut at boundary intersection point rather than snapping
+            p0, p1 = snap(clipped[0]), snap(clipped[-1])
+
+            complete_lines.append(
+                np.concatenate([np.array(p0.xy).T, clipped, np.array(p1.xy).T])
+            )
+
+        return complete_lines

mosaic/contour.py

@@ -251,7 +251,7 @@ def __init__(self, descriptor: Descriptor, z: ArrayLike):
         self.ds = descriptor.ds
         self._z = np.asarray(array)
 
-        self.boundary_edge_mask = (self.ds.cellsOnEdge == -1).any("TWO").values
+        self.boundary_edge_mask = (self.ds.cellsOnEdge < 0).any("TWO").values
         self.boundary_vertices = np.unique(
             self.ds.verticesOnEdge[self.boundary_edge_mask]
         )
@@ -291,9 +291,9 @@ def _create_contour_graph(
         """ """
         ds = self.ds
 
-        padded_mask = np.r_[False, mask]
+        padded_mask = np.r_[False, False, mask]
         # mark mask as False for all cells outside domain
-        cells_on_edge_mask = np.asarray(padded_mask[ds.cellsOnEdge + 1])
+        cells_on_edge_mask = np.asarray(padded_mask[ds.cellsOnEdge + 2])
 
         # boolean mask for edges along contour
         edge_mask = np.logical_xor.reduce(cells_on_edge_mask, axis=1)

mosaic/descriptor.py

@@ -725,7 +725,7 @@ def _compute_cell_patches(ds: Dataset) -> ndarray:
     # connectivity arrays have already been zero indexed
     verticesOnCell = ds.verticesOnCell
     # get a mask of the active vertices
-    mask = verticesOnCell == -1
+    mask = verticesOnCell < 0
 
     # tile the first vertices index
     firstVertex = np.tile(verticesOnCell[:, 0], (maxEdges, 1)).T
@@ -782,8 +782,8 @@ def _compute_vertex_patches(ds: Dataset) -> ndarray:
     cellsOnVertex = ds.cellsOnVertex.values
     edgesOnVertex = ds.edgesOnVertex.values
     # get a mask of active nodes
-    cellMask = cellsOnVertex == -1
-    edgeMask = edgesOnVertex == -1
+    cellMask = cellsOnVertex < 0
+    edgeMask = edgesOnVertex < 0
 
     # get the coordinates needed to patch construction
     xCell = ds.xCell.values

mosaic/utils.py

@@ -44,8 +44,10 @@ def _make_lookup_table(mask: np.ndarray[bool]) -> np.ndarray[np.int64]:
     old_idx = np.flatnonzero(mask)  # 0..N-1
     new_idx = np.arange(old_idx.size, dtype=np.int64)
 
-    # lut[0] reserved for old=-1
-    lut = np.full(mask.size + 1, -1, dtype=np.int64)
+    # make culling along projection boundary with -2
+    lut = np.full(mask.size + 1, -2, dtype=np.int64)
+    # culled land boundaries stay -1
+    lut[0] = -1
     lut[old_idx + 1] = new_idx
 
     return lut

pyproject.toml

@@ -64,7 +64,10 @@ build-backend = "setuptools.build_meta"
 minversion = "6"
 addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
 xfail_strict = true
-filterwarnings = ["error"]
+filterwarnings = [
+    "error",
+    "ignore:numpy.ndarray size changed:RuntimeWarning",
+]
 log_cli_level = "INFO"
 testpaths = ["tests"]
 

tests/conftest.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import inspect
+
+import cartopy.crs as ccrs
+import pytest
+
+import mosaic
+
+# get the names, as strings, of unsupported projections for spherical meshes
+unsupported = [
+    p.__name__ for p in mosaic.descriptor.UNSUPPORTED_SPHERICAL_PROJECTIONS
+]
+
+PROJECTIONS = [
+    obj()
+    for name, obj in inspect.getmembers(ccrs)
+    if inspect.isclass(obj)
+    and issubclass(obj, ccrs.Projection)
+    and not name.startswith("_")  # skip internal classes
+    and obj is not ccrs.Projection  # skip base Projection class
+    and name not in unsupported  # skip unsupported projections
+]
+
+
+def id_func(projection):
+    return type(projection).__name__
+
+
+@pytest.fixture(scope="module", params=PROJECTIONS, ids=id_func)
+def iterate_supported_projections(request):
+    def _factory(ds):
+        return mosaic.Descriptor(ds, request.param, ccrs.Geodetic())
+
+    return _factory