diff --git a/docs/developers_guide/api.md b/docs/developers_guide/api.md index 9066a8f..57104e7 100644 --- a/docs/developers_guide/api.md +++ b/docs/developers_guide/api.md @@ -1,16 +1,42 @@ (dev-api)= -# API Reference - ```{eval-rst} .. currentmodule:: mosaic +.. _api: + +API Reference +============= + +This page provides an auto-generated summary of mosaic's API. + +Top Level Utilities +------------------- .. autosummary:: :toctree: generated/ polypcolor Descriptor datasets.open_dataset -``` +Descriptor +---------- + +Constructor +~~~~~~~~~~~ +.. autosummary:: + :toctree: generated/ + Descriptor + +Properties +~~~~~~~~~~ +.. autosummary:: + :toctree: generated/ + + Descriptor.cell_patches + Descriptor.edge_patches + Descriptor.vertex_patches + Descriptor.transform + +``` diff --git a/mosaic/descriptor.py b/mosaic/descriptor.py index 6dcd3bb..970ff0b 100644 --- a/mosaic/descriptor.py +++ b/mosaic/descriptor.py @@ -2,6 +2,9 @@ import numpy as np import xarray as xr +from cartopy.crs import CRS +from numpy import ndarray +from xarray.core.dataset import Dataset renaming_dict = {"lonCell": "xCell", "latCell": "yCell", @@ -17,75 +20,119 @@ "edgesOnVertex"] -class Descriptor: +def attr_to_bool(attr: str): + """Format attribute strings and return a boolean value """ - Class describing unstructured MPAS meshes in order to support plotting - within ``matplotlib``. The class contains various methods to create - :py:class:`~matplotlib.collections.PolyCollection` objects for - variables defined at cell centers, vertices, and edges. - - Parameters - ---------- - - ds : DataArray - A valid mpas mesh dataset... - - projection : cartopy.crs.Projection - The coordinate system .... + match attr.strip().upper(): + case "YES": + return True + case "NO": + return False + case _: + raise ValueError("f{attr} was unable to be parsed as YES/NO") - transform : cartopy.crs.Projection - The coordinate system .... - - use_latlon : bool - Whether to use the lat/lon coordinate arrays to construct the patches +class Descriptor: + """Data structure describing unstructured MPAS meshes. + + Enables visualization of fields defined at cell centers, vertices, + and edges within ``matplotlib`` through the creation of + :py:class:`~matplotlib.collections.PolyCollection` objects. Separate patch + arrays, which are subsequently passed to the + :py:class:`~matplotlib.collections.PolyCollection` class by the + :py:class:`polypcolor` method, + are lazily loaded for each variable location (i.e. cells, edges, and + vertices) when attribute is first looked up. We use lazily loaded arrays + because we are rarely plotting variables located at all three locations + and patch array creation is the most expensive + process in our visualization procedure. + + The resulting patch arrays properly handle culled mesh boundaries (i.e. + culled land boundaries for spherical meshes and/or mesh boundary for planar + non-periodic meshes). Additionally, x and/or y periodicity is properly + handled for planar periodic mesh and we can correct for patch wrapping over + the antimeridian for projected spherical meshes. - Attributes + Parameters ---------- - - latlon : bool - pass - projection : cartopy.crs.Projection - pass - transform : cartopy.crs.Projection - pass - + mesh_ds : DataSet + A valid MPAS mesh dataset that contains the basic mesh variables (i.e. + coordinate and connectivity arrays) needed for creating patch arrays. + projection : cartopy.crs.Projection, optional + The target projection for plotting. + transform : cartopy.crs.Projection, optional + The coordinate system in which the parent mesh coordinates are defined. + use_latlon : bool, optional + Whether to use the lat/lon coordinate arrays to construct the patches. + + Notes + ----- + If both the ``projection`` and ``transform`` parameters are passed, then + the coordinate arrays in the :attr:`.Descriptor.ds` will be transformed + prior to patch construction. We do this as matter of efficiency, since + transforming the one-dimensional coordinate arrays is much faster than + transforming the multi-dimensional patch arrays. This means the + :attr:`.Descriptor.transform` will **not** be the values passed to the + constructor, but instead will be equal to :attr:`.Descriptor.projection` + + Examples + -------- + >>> import cartopy.crs as ccrs + >>> import mosaic + >>> + >>> ds = mosaic.datasets.open_dataset("QU.240km") + >>> + >>> transform = ccrs.PlateCarree() + >>> projection = ccrs.NorthPolarStereo() + >>> + >>> # set `use_latlon` to True b/c our transform expects lat/lon coords + >>> descriptor = mosaic.Descriptor( + >>> ds, projection, transform, use_latlon=True + >>> ) """ - def __init__(self, ds, projection=None, transform=None, use_latlon=False): - """ - """ - self.latlon = use_latlon - self.projection = projection + def __init__( + self, + mesh_ds: Dataset, + projection: CRS | None = None, + transform: CRS | None = None, + use_latlon: bool = False + ) -> None: + #: The coordinate system in which patch coordinates are defined. + #: + #: This *could* have a different value than ``transform`` parameter + #: passed to the constructor, because we transform the one-dimensional + #: coordinate arrays at initialization, if both the ``transform`` and + #: ``projection`` kwargs are provided. 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': - self.latlon = True - # also check if projection requires lat/lon coords + #: Boolean whether parent mesh is spherical + self.is_spherical = attr_to_bool(mesh_ds.on_a_sphere) - # create a minimal dataset, stored as an attr, for patch creation - self.ds = self.create_minimal_dataset(ds) + # calls attribute setter method + self.latlon = use_latlon + + #: :py:class:`~xarray.Dataset` that contains the minimal subset of + #: coordinate and connectivity arrays from the parent mesh needed to + #: create patches arrays. + self.ds = self._create_minimal_dataset(mesh_ds) - # reproject the minimal dataset, even for non-spherical meshes - if projection and transform: - self._transform_coordinates(projection, transform) - self._pre_projected = True + # calls ``projection.setter`` method, which will transform coordinates + # if both a projection and transform were provided to the constructor + self.projection = projection - def create_minimal_dataset(self, ds): + def _create_minimal_dataset(self, ds: Dataset) -> Dataset: """ Create a xarray.Dataset that contains the minimal subset of - coordinate / connectivity arrays needed to create pathces for plotting + coordinate / connectivity arrays needed to create patches for plotting Parameters ---------- - ds : DataArray - A valid mpas mesh dataset.... + A valid MPAS mesh dataset """ - def fix_outofbounds_indices(ds, array_name): + def fix_outofbounds_indices(ds: Dataset, array_name: str) -> Dataset: """ Some meshes (e.g. QU240km) don't do masking of ragged indices with 0. Instead they use `nInidices+1` as the invalid value, @@ -94,7 +141,7 @@ def fix_outofbounds_indices(ds, array_name): NOTE: Assumes connectivity arrays are zero indexed """ - # progamatically create the appropriate dimension name + # programmatically create the appropriate dimension name dim = "n" + array_name.split("On")[0].title() # get the maximum valid size for the array to be indexed too maxSize = ds.sizes[dim] @@ -120,65 +167,142 @@ def fix_outofbounds_indices(ds, array_name): minimal_ds.attrs.clear() for array in connectivity_arrays: - # zero index all the connectivty arrays + # zero index all the connectivity arrays minimal_ds[array] = minimal_ds[array] - 1 - # fix any out of bounds indicies + # fix any out of bounds indices minimal_ds = fix_outofbounds_indices(minimal_ds, array) if self.latlon: - # convert lat/lon coordinates from radian to degrees for loc in ["Cell", "Edge", "Vertex"]: minimal_ds[f"lon{loc}"] = np.rad2deg(minimal_ds[f"lon{loc}"]) minimal_ds[f"lat{loc}"] = np.rad2deg(minimal_ds[f"lat{loc}"]) - # rename the coordinate arrays to all be named x.../y... - # irrespective of whether spherical or cartesian coords are used + # rename the coordinate arrays to be named x.../y... irrespective + # of whether spherical or Cartesian coordinates are used minimal_ds = minimal_ds.rename(renaming_dict) return minimal_ds + @property + def projection(self) -> CRS: + """The target projection for plotting. + """ + return self._projection + + @projection.setter + def projection(self, projection: CRS) -> None: + # Issue warning if changing the projection after initialization + # TODO: Add heuristic size (i.e. ``self.ds.nbytes``) above which the + # warning is raised + if hasattr(self, "_projection"): + print(("Reprojecting the descriptor can be inefficient" + "for large meshes")) + + # If both a projection and a transform are provided then + if projection and self.transform: + # reproject coordinate arrays in the minimal dataset + self._transform_coordinates(projection, self.transform) + # update the transform attribute following the reprojection + self.transform = projection + # Then loop over patch attributes + for loc in ["cell", "edge", "vertex"]: + attr = "f{loc}_patches" + # and only delete attributes that have previously been cached + if attr in self.__dict__: + del self.__dict__[attr] + + self._projection = projection + + @property + def latlon(self) -> bool: + """Boolean whether the lat/lon coordinate arrays should be used for + patch construction. + """ + return self._latlon + + @latlon.setter + def latlon(self, value) -> None: + """TODO: check that the passed value is consistent with transform + """ + if self.is_spherical: + value = True + + self._latlon = value + @cached_property - def cell_patches(self): - """ Lazily loaded cell patches + def cell_patches(self) -> ndarray: + """:py:class:`~numpy.ndarray` of patch coordinates for cell centered + values + + **Dimensions**: ``(nCells, maxEdges)`` + + Notes + ----- + The second dimension of the patch array is length ``maxEdges``, even if + ``nEdgesOnCell`` for the ``i-th`` cell is less than ``maxEdges``. We've + chosen to deal with ragged indices (i.e. nodes indices greater than + ``nEdgesOnCell`` for the given cell) by repeating the first node of the + patch. Nodes are ordered counter clockwise around the cell center. """ patches = _compute_cell_patches(self.ds) patches = self._fix_antimeridian(patches, "Cell") return patches @cached_property - def edge_patches(self): + def edge_patches(self) -> ndarray: + """:py:class:`~numpy.ndarray` of patch coordinates for edge centered + values + + **Dimensions**: ``(nEdges, 4)`` + + Notes + ----- + 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. + """ patches = _compute_edge_patches(self.ds) patches = self._fix_antimeridian(patches, "Edge") return patches @cached_property - def vertex_patches(self): + def vertex_patches(self) -> ndarray: + """:py:class:`~numpy.ndarray` of patch coordinates for vertex centered + values + + **Dimensions**: ``(nVertices, 6)`` + + Notes + ----- + 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 to simplify vertex patches creation 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). + + 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. + """ patches = _compute_vertex_patches(self.ds) 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): - """ - """ for loc in ["Cell", "Edge", "Vertex"]: transformed_coords = projection.transform_points( transform, self.ds[f"x{loc}"], self.ds[f"y{loc}"]) - # transformed_coords is a np array so need to assign to the values + # ``transformed_coords`` is a numpy array so needs to assigned to + # the values of the dataarray self.ds[f"x{loc}"].values = transformed_coords[:, 0] self.ds[f"y{loc}"].values = transformed_coords[:, 1] @@ -246,29 +370,9 @@ def _fix_antimeridian(self, patches, loc, projection=None): return patches - def transform_patches(self, patches, projection, transform): - """ - """ - - raise NotImplementedError("This is a place holder. Do not use.") - transformed_patches = projection.transform_points( - transform, patches[..., 0], patches[..., 1]) - - # transformation will return x,y,z values. Only need x and y - patches.data[...] = transformed_patches[..., 0:2] - - return patches - - -def _compute_cell_patches(ds): +def _compute_cell_patches(ds: Dataset) -> ndarray: """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"] @@ -291,14 +395,8 @@ def _compute_cell_patches(ds): return nodes -def _compute_edge_patches(ds): +def _compute_edge_patches(ds: Dataset) -> ndarray: """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 @@ -331,20 +429,8 @@ def _compute_edge_patches(ds): return nodes -def _compute_vertex_patches(ds): +def _compute_vertex_patches(ds: Dataset) -> ndarray: """Create vertex patches (i.e. Dual Cells) for an MPAS mesh. - - 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). - - 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"] @@ -382,7 +468,7 @@ def _compute_vertex_patches(ds): # ------------------------------------------------------------------------- # 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. + # the full kite area and is properly 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]) diff --git a/mosaic/polypcolor.py b/mosaic/polypcolor.py index 1b6e6c3..843c722 100644 --- a/mosaic/polypcolor.py +++ b/mosaic/polypcolor.py @@ -59,7 +59,7 @@ def polypcolor( elif "nVertices" in c.dims: verts = descriptor.vertex_patches - transform = descriptor.get_transform() + transform = descriptor.transform collection = PolyCollection(verts, alpha=alpha, array=c, cmap=cmap, norm=norm, **kwargs)