Switch to mosaic for all plotting functionality

deploy/conda-dev-spec.template

@@ -1,6 +1,5 @@
 # Base
 python>=3.9,<3.13
-antimeridian
 cartopy
 cartopy_offlinedata
 cmocean
@@ -9,9 +8,6 @@ dask <2025.1.0
 esmf={{ esmf }}={{ mpi_prefix }}_*
 ffmpeg
 geometric_features={{ geometric_features }}
-geoviews
-holoviews
-hvplot
 importlib_resources
 ipython
 jupyter
@@ -22,6 +18,7 @@ mache={{ mache }}
 matplotlib-base>=3.9.0
 metis={{ metis }}
 moab={{ moab }}=*_tempest_*
+mosaic>=1.1.0,<2.0.0
 mpas_tools={{ mpas_tools }}
 nco
 netcdf4=*=nompi_*
@@ -37,8 +34,6 @@ ruamel.yaml
 requests
 scipy>=1.8.0
 shapely>=2.0,<3.0
-spatialpandas
-uxarray <2025.01.0
 xarray
 
 # Static typing

docs/conf.py

@@ -71,6 +71,7 @@
         ('https://mpas-dev.github.io/geometric_features/main', None),
     'matplotlib': ('https://matplotlib.org/stable', None),
     'mpas_tools': ('https://mpas-dev.github.io/MPAS-Tools/master', None),
+    'mosaic': ('https://docs.e3sm.org/mosaic', None),
     'numpy': ('https://numpy.org/doc/stable', None),
     'python': ('https://docs.python.org', None),
     'scipy': ('https://docs.scipy.org/doc/scipy/reference', None),

docs/developers_guide/api.md

@@ -297,6 +297,7 @@ seaice/api
    :toctree: generated/
 
    area_for_field
+   cell_mask_to_edge_mask
    time_index_from_xtime
 ```
 

docs/developers_guide/framework/visualization.md

@@ -6,9 +6,10 @@ Visualization is an optional, but desirable aspect of tasks. Often,
 visualization is an optional step of a task but can also be included
 as part of other steps such as `init` or `analysis`.
 
-While developers can write their own visualization scripts associated with
-individual tasks, the following shared visualization routines are
-provided in `polaris.viz`:
+Horizontal visualization of MPAS fields is enabled through the use of
+[`mosaic`](https://docs.e3sm.org/mosaic/). While developers can write their
+own visualization scripts associated with individual tasks, the following
+shared visualization routines are provided in `polaris.viz`:
 
 (dev-visualization-style)=
 
@@ -25,13 +26,15 @@ before creating a `matplotlib` figure.
 
 ## horizontal fields from planar meshes
 
-{py:func}`polaris.viz.plot_horiz_field()` produces a patches-style
-visualization of x-y fields across a single vertical level at a single time
-step. The image file (png) is saved to the directory from which
-{py:func}`polaris.viz.plot_horiz_field()` is called. The function
-automatically detects whether the field specified by its variable name is
-a cell-centered variable or an edge-variable and generates the patches, the
-polygons characterized by the field values, accordingly.
+{py:func}`polaris.viz.plot_horiz_field()` produces a visualization of 
+horizontal fields at their native mesh location (i.e. cells, edges, or
+vertices) at a single vertical level and a single time step. The image file
+(png) is saved to the directory from which 
+{py:func}`polaris.viz.plot_horiz_field()` is called. 
+{py:func}`polaris.viz.plot_horiz_field()` is jut a wrapper for
+{py:func}`mosaic.polypcolor()`, which automatically detects whether the field
+to be plotted is defined at cells, edges, or vertices and generates the patches
+(i.e. the polygons characterized by the field values) accordingly.
 
 ```{image} images/baroclinic_channel_cell_patches.png
 :align: center
@@ -47,39 +50,38 @@ An example function call that uses the default vertical level (top) is:
 
 ```python
 cell_mask = ds_init.maxLevelCell >= 1
-plot_horiz_field(config, ds, ds_mesh, 'normalVelocity',
-                 'final_normalVelocity.png',
-                 t_index=t_index,
-                 vmin=-max_velocity, vmax=max_velocity,
-                 cmap='cmo.balance', show_patch_edges=True,
-                 cell_mask=cell_mask)
-```
+edge_mask = cell_mask_to_edge_mask(ds_init, cell_mask)
 
-The `cell_mask` argument can be any field indicating which horizontal cells
-are valid and which are not.  A typical value for ocean plots is as shown 
-above: whether there are any active cells in the water column.
+plot_horiz_field(ds_mesh, ds['normalVelocity'], 'final_normalVelocity.png',
+                 t_index=t_index, vmin=-max_velocity, vmax=max_velocity,
+                 cmap='cmo.balance', show_patch_edges=True, 
+                 field_mask=edge_mask)
+```
 
-For increased efficiency, you can store the `patches` and `patch_mask` from
-one call to `plot_horiz_field()` and reuse them in subsequent calls.  The
-`patches` and `patch_mask` are specific to the dimension (`nCell` or `nEdges`) 
-of the field to plot and the `cell_mask`.  So separate `patches` and 
-`patch_mask` variables should be stored for as needed:
+The `field_mask` argument can be any field indicating which horizontal mesh
+locations are valid and which are not, but it must be the same shape as data
+array being plotted.  A typical value for ocean plots is as shown 
+above: whether there are any active cells in the water column and then the cell
+mask is converted to an edges mask using the
+{py:func}`polaris.mpas.cell_mask_to_edge_mask()` function.
 
+For increased efficiency, you can store the instance of 
+{py:class}`mosaic.Descriptor` returned by `plot_horiz_field()` and reuse it in
+subsequent calls; assuming you are plotting with the same mesh.
 ```python
 cell_mask = ds_init.maxLevelCell >= 1
-cell_patches, cell_patch_mask = plot_horiz_field(
-    ds=ds, ds_mesh=ds_mesh, field_name='ssh', out_file_name='plots/ssh.png',
-    vmin=-720, vmax=0, figsize=figsize, cell_mask=cell_mask)
+descriptor = plot_horiz_field(ds_mesh, ds['ssh'], 'plots/ssh.png',
+                              vmin=-720, vmax=0, figsize=figsize,
+                              field_mask=cell_mask)
 
-plot_horiz_field(ds=ds, ds_mesh=ds_mesh, field_name='bottomDepth', 
-                 out_file_name='plots/bottomDepth.png', vmin=0, vmax=720,
-                 figsize=figsize, patches=cell_patches, 
-                 patch_mask=cell_patch_mask)
+plot_horiz_field(ds_mesh, ds['bottomDepth'], 'plots/bottomDepth.png',
+                 vmin=0, vmax=720, figsize=figsize, field_mask=cell_mask,
+                 descriptor=descriptor)
 
-edge_patches, edge_patch_mask = plot_horiz_field(
-    ds=ds, ds_mesh=ds_mesh, field_name='normalVelocity', 
-    out_file_name='plots/normalVelocity.png', t_index=t_index, 
-    vmin=-0.1, vmax=0.1, cmap='cmo.balance', cell_mask=cell_mask)
+edge_mask = cell_mask_to_edge_mask(ds_mesh, cell_mask)
+plot_horiz_field(ds_mesh, ds['normalVelocity'], 'plots/normalVelocity.png',
+                 t_index=t_index, vmin=-0.1, vmax=0.1, cmap='cmo.balance',
+                 field_mask=edge_mask, descriptor=descriptor)
 
 ...
 ```
@@ -91,7 +93,14 @@ edge_patches, edge_patch_mask = plot_horiz_field(
 ### plotting from spherical MPAS meshes
 
 You can use {py:func}`polaris.viz.plot_global_mpas_field()` to plot a field on
-a spherical MPAS mesh.
+a spherical MPAS mesh. Like the planar visualization function, this is also
+just a wrapper to {py:func}`mosaic.polypcolor()`. Thanks to `mosaic` variables
+defined at cells, edges, and vertices are all support as well as meshes with
+culled land boundaries are also supported. While `mosaic`
+[supports](https://docs.e3sm.org/mosaic/user_guide/wrapping.html) a variety
+of map projection for spherical meshes,
+{py:func}`polaris.viz.plot_global_mpas_field()` currently only supports
+[`cartopy.crs.PlateCarree`](https://scitools.org.uk/cartopy/docs/latest/reference/projections.html#cartopy.crs.PlateCarree).
 
 ```{image} images/cosine_bell_final_mpas.png
 :align: center
@@ -127,7 +136,7 @@ The `central_longitude` defaults to `0.0` and can be set to another value
 (typically 180 degrees) for visualizing quantities that would otherwise be
 divided across the antimeridian.
 
-The `colormap_section` of the config file must contain config options for
+The `<task>_viz` section of the config file must contain config options for
 specifying the colormap:
 
 ```cfg
@@ -141,26 +150,25 @@ colormap_name = viridis
 # the type of norm used in the colormap
 norm_type = linear
 
-# colorbar limits
-colorbar_limits = 0.0, 1.0
+# A dictionary with keywords for the norm
+norm_args = {'vmin': 0., 'vmax': 1.}
 ```
 
 `colormap_name` can be any available matplotlib colormap.  For ocean test
 cases, we recommend importing [cmocean](https://matplotlib.org/cmocean/) so
 the standard ocean colormaps are available.
 
-The `norm_type` is one of `linear` (a linear colormap) or `log` (a logarithmic 
-colormap).
+The `norm_type` is one of `linear` (a linear colormap), `symlog` (a 
+[symmetric log](https://matplotlib.org/stable/gallery/images_contours_and_fields/colormap_normalizations_symlognorm.html)
+colormap with a central linear region), or `log` (a logarithmic colormap).
 
-The `colorbar_limits` are the lower and upper bound of the colorbar range.
+The `norm_args` depend on the `norm_typ` and are the arguments to
+{py:class}`matplotlib.colors.Normalize`, {py:class}`matplotlib.colors.SymLogNorm`,
+and {py:class}`matplotlib.colors.LogNorm`, respectively. 
 
-There are also two optional config options used to set the colors on either end of the colormap:
+The config option `colorbar_ticks` (if it is defined) specifies tick locations
+along the colorbar. If it is not specified, they are determined automatically.
 
-```cfg
-# [optional] colormap set_under and set_over options
-under_color = k
-over_color = orange
-```
 ### plotting from lat/lon grids
 
 You can use {py:func}`polaris.viz.plot_global_lat_lon_field()` to plot a field 
@@ -200,38 +208,5 @@ class Viz(Step):
             title='Tracer at init', plot_land=False)
 ```
 
-The `colormap_section` of the config file must contain config options for
-specifying the colormap:
-
-```cfg
-# options for visualization for the cosine bell convergence task
-[cosine_bell_viz]
-
-# colormap options
-# colormap
-colormap_name = viridis
-
-# the type of norm used in the colormap
-norm_type = linear
-
-# A dictionary with keywords for the norm
-norm_args = {'vmin': 0., 'vmax': 1.}
-
-# We could provide colorbar tick marks but we'll leave the defaults
-# colorbar_ticks = np.linspace(0., 1., 9)
-```
-
-`colormap_name` can be any available matplotlib colormap.  For ocean test
-cases, we recommend importing [cmocean](https://matplotlib.org/cmocean/) so
-the standard ocean colormaps are available.
-
-The `norm_type` is one of `linear` (a linear colormap), `symlog` (a 
-[symmetric log](https://matplotlib.org/stable/gallery/images_contours_and_fields/colormap_normalizations_symlognorm.html)
-colormap with a central linear region), or `log` (a logarithmic colormap).
-
-The `norm_args` depend on the `norm_typ` and are the arguments to
-{py:class}`matplotlib.colors.Normalize`, {py:class}`matplotlib.colors.SymLogNorm`,
-and {py:class}`matplotlib.colors.LogNorm`, respectively. 
-
-The config option `colorbar_ticks` (if it is defined) specifies tick locations
-along the colorbar. If it is not specified, they are determined automatically.
+The `<task>_viz` of the config file is the same as what's used by
+{py:func}`polaris.viz.plot_global_mpas_field()`.

polaris/mpas/__init__.py

@@ -1,2 +1,3 @@
 from polaris.mpas.area import area_for_field
+from polaris.mpas.mask import cell_mask_to_edge_mask
 from polaris.mpas.time import time_index_from_xtime, time_since_start

polaris/mpas/mask.py

@@ -0,0 +1,35 @@
+
+def cell_mask_to_edge_mask(ds_mesh, cell_mask):
+    """Convert a cell mask to edge mask using mesh connectivity information
+
+    True corresponds to valid cells and False are invalid cells
+
+    Parameters
+    ----------
+    ds_mesh : xarray.Dataset
+        The MPAS mesh
+
+    cell_mask : xarray.DataArray
+        The cell mask we want to convert to an edge mask
+
+
+    Returns
+    -------
+    edge_mask : xarray.DataArray
+        The edge mask corresponding to the input cell mask
+    """
+
+    # test if any are False
+    if ~cell_mask.any():
+        return ds_mesh.nEdges > -1
+
+    # zero index the connectivity array
+    cells_on_edge = (ds_mesh.cellsOnEdge - 1)
+
+    # using nCells (dim) instead of indexToCellID since it's already 0 indexed
+    masked_cells = ds_mesh.nCells.where(~cell_mask, drop=True).astype(int)
+
+    # use inverse so True/False convention matches input cell_mask
+    edge_mask = ~cells_on_edge.isin(masked_cells).any("TWO")
+
+    return edge_mask

polaris/ocean/tasks/baroclinic_channel/init.py

@@ -7,6 +7,7 @@
 
 from polaris import Step
 from polaris.mesh.planar import compute_planar_hex_nx_ny
+from polaris.mpas import cell_mask_to_edge_mask
 from polaris.ocean.vertical import init_vertical_coord
 from polaris.ocean.viz import compute_transect, plot_transect
 from polaris.viz import plot_horiz_field
@@ -163,10 +164,11 @@ def run(self):
         write_netcdf(ds, 'initial_state.nc')
 
         cell_mask = ds.maxLevelCell >= 1
+        edge_mask = cell_mask_to_edge_mask(ds, cell_mask)
 
-        plot_horiz_field(ds, ds_mesh, 'normalVelocity',
+        plot_horiz_field(ds_mesh, ds['normalVelocity'],
                          'initial_normal_velocity.png', cmap='cmo.balance',
-                         show_patch_edges=True, cell_mask=cell_mask)
+                         show_patch_edges=True, field_mask=edge_mask)
 
         y_min = ds_mesh.yVertex.min().values
         y_max = ds_mesh.yVertex.max().values
@@ -191,6 +193,6 @@ def run(self):
                       vmin=vmin, vmax=vmax, cmap='cmo.thermal',
                       colorbar_label=r'$^\circ$C', color_start_and_end=True)
 
-        plot_horiz_field(ds, ds_mesh, 'temperature', 'initial_temperature.png',
+        plot_horiz_field(ds_mesh, ds['temperature'], 'initial_temperature.png',
                          vmin=vmin, vmax=vmax, cmap='cmo.thermal',
-                         cell_mask=cell_mask, transect_x=x, transect_y=y)
+                         field_mask=cell_mask, transect_x=x, transect_y=y)

polaris/ocean/tasks/baroclinic_channel/rpe/analysis.py

@@ -104,10 +104,10 @@ def run(self):
             time_index = np.argmin(np.abs(times - time))
 
             cell_mask = ds_init.maxLevelCell >= 1
-            plot_horiz_field(ds, ds_mesh, 'temperature', ax=ax,
+            plot_horiz_field(ds_mesh, ds['temperature'], ax=ax,
                              cmap='cmo.thermal', t_index=time_index,
                              vmin=min_temp, vmax=max_temp,
-                             cmap_title='SST (C)', cell_mask=cell_mask)
+                             cmap_title='SST (C)', field_mask=cell_mask)
             ax.set_title(f'day {times[time_index]:g}, $\\nu_h=${nu:g}')
 
         plt.savefig(output_filename)

polaris/ocean/tasks/baroclinic_channel/viz.py

@@ -3,6 +3,7 @@
 import xarray as xr
 
 from polaris import Step
+from polaris.mpas import cell_mask_to_edge_mask
 from polaris.ocean.viz import compute_transect, plot_transect
 from polaris.viz import plot_horiz_field
 
@@ -43,13 +44,14 @@ def run(self):
         ds = xr.load_dataset('output.nc')
         t_index = ds.sizes['Time'] - 1
         cell_mask = ds_init.maxLevelCell >= 1
+        edge_mask = cell_mask_to_edge_mask(ds_init, cell_mask)
         max_velocity = np.max(np.abs(ds.normalVelocity.values))
-        plot_horiz_field(ds, ds_mesh, 'normalVelocity',
+        plot_horiz_field(ds_mesh, ds['normalVelocity'],
                          'final_normalVelocity.png',
                          t_index=t_index,
                          vmin=-max_velocity, vmax=max_velocity,
                          cmap='cmo.balance', show_patch_edges=True,
-                         cell_mask=cell_mask)
+                         field_mask=edge_mask)
 
         y_min = ds_mesh.yVertex.min().values
         y_max = ds_mesh.yVertex.max().values
@@ -76,7 +78,7 @@ def run(self):
                       vmin=vmin, vmax=vmax, cmap='cmo.thermal',
                       colorbar_label=r'$^\circ$C', color_start_and_end=True)
 
-        plot_horiz_field(ds, ds_mesh, 'temperature', 'final_temperature.png',
+        plot_horiz_field(ds_mesh, ds['temperature'], 'final_temperature.png',
                          t_index=t_index, vmin=vmin, vmax=vmax,
-                         cmap='cmo.thermal', cell_mask=cell_mask, transect_x=x,
-                         transect_y=y)
+                         cmap='cmo.thermal', field_mask=cell_mask,
+                         transect_x=x, transect_y=y)