regional masks not supported?

[dylanschlichting]

I've been working on adding a regional mask for a Gulf of Mexico RRM to do some statistical analysis. My goal is to plot the masked output with mosaic as a QC check, which fails. See below a successful example with ax.scatter using normalized vorticity compared to mosaic. Does this capability need to be added to the code, or is this user error? One reason this capability would be useful is that ax.scatter is much slower and the user has to carefully prescribe the dot size and dpi, which often takes a few iterations to get right. As a result, the boundaries of the masked data aren't always clear with a scatter plot.

Preliminaries:

#Packages 
import numpy as np
import xarray as xr
import glob
import mosaic
import cartopy
import cartopy.crs as ccrs
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import cartopy.feature as cfeature
import cmocean.cm as cmo
import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib import ticker
from matplotlib.dates import DateFormatter
from matplotlib.ticker import AutoMinorLocator, MultipleLocator

land_10m = cfeature.NaturalEarthFeature('physical', 'land', '10m',
                                edgecolor='face',
                                facecolor=cfeature.COLORS['land'])

#HF output
path = '/lustre/scratch5/dschlichting/E3SM/scratch/chicoma-cpu/Tl319_GoM1pt5_GMPAS-IAF_norivers/run/Tl319_GoM1pt5_GMPAS-IAF.mpaso.hist.am.highFrequencyOutput.1997-0*-01_00.00.00.nc'
ds = xr.open_mfdataset(glob.glob(path), concat_dim="Time", combine = 'nested')
ds = ds.sortby(ds.xtime)
ds
#Initial condition
pathi = '/lustre/scratch5/dschlichting/runs/gom1pt5r2/ocean/global_ocean/GoM5/WOA23/init/initial_state/initial_state.nc'
dsi = xr.open_dataset(pathi)

lat_c = np.degrees(dsi.variables['latCell'][:])
lon_c = np.degrees(dsi.variables['lonCell'][:])
lat_rad = 1.0/np.cos(np.radians(lat_c.mean().values))

crs = ccrs.PlateCarree(central_longitude=0)
projection = crs
transform = ccrs.PlateCarree()

rv = ds.relativeVorticityAtSurface/dsi.fCell

#GoM regional mask
pathm = '/usr/projects/climate/dschlichting/repos/analysis_notebooks/gom1pt5/gom1pt5_gom_mask.nc'
mask = xr.open_dataset(pathm)
iGoM = np.where(mask.regionCellMasks[:,0] == 1)

The contents of the mask file look like:

xarray.Dataset
    Dimensions:
        nCells: 1258039 nRegions: 1 nVertices: 2536448 nRegionGroups: 2 maxRegionsInGroup: 1
    Coordinates: (0)
    Data variables:
        regionCellMasks
        (nCells, nRegions)
        int32
        ...
        nCellsInRegion
        (nRegions)
        int32
        ...
        regionVertexMasks
        (nVertices, nRegions)
        int32
        ...
        nVerticesInRegion
        (nRegions)
        int32
        ...
        regionNames
        (nRegions)
        |S64
        ...
        nRegionsInGroup
        (nRegionGroups)
        int32
        ...
        regionsInGroup
        (nRegionGroups, maxRegionsInGroup)
        int32
        ...
        regionGroupNames
        (nRegionGroups)
        |S64
        ...
    Indexes: (0)
    Attributes:

    parent_id :
        4u9cvgtcjd
        qgux71w9vn
        1zdkm06zwk
        4rdyc30dwg
    history :
        MpasMaskCreator.x /lustre/scratch5/dschlichting/runs/gom1pt5r2/ocean/global_ocean/GoM5/WOA23/init/initial_state/initial_state.nc gom1pt5_gom_mask.nc
        Mon Oct 28 11:58:02 2024: ncks --glb_att_add MPAS_Mesh_Short_Name=WC1.5to60E3r1 --glb_att_add MPAS_Mesh_Long_Name=WC1.5to60kmL80E3SMv3r1 --glb_att_add MPAS_Mesh_Prefix=WC --glb_att_add MPAS_Mesh_E3SM_Version=3 --glb_att_add MPAS_Mesh_Pull_Request=https://github.com/MPAS-Dev/MPAS-Model/pull/628 --glb_att_add MPAS_Mesh_WC_Revision=1 --glb_att_add MPAS_Mesh_WC_Version_Author=dylanschlichting --glb_att_add MPAS_Mesh_WC_Version_Author_Email=dylanschlichting --glb_att_add MPAS_Mesh_WC_Version_Creation_Date=20241028 --glb_att_add MPAS_Mesh_WC_Minimum_Resolution_km=1.5 --glb_att_add MPAS_Mesh_WC_Maximum_Resolution_km=60 --glb_att_add MPAS_Mesh_WC_Maximum_Depth_m=5500.0 --glb_att_add MPAS_Mesh_WC_Number_of_Levels=80 --glb_att_add MPAS_Mesh_Description=MPAS Gulf of Mexico submesoscale permitting mesh 3, with a focused 1.5-km resolution around North America and 80 vertical levels --glb_att_add MPAS_Mesh_Bathymetry=Bathymetry is from GEBCO 2023, combined with BedMachine Antarctica v3 around Antarctica. --glb_att_add MPAS_Initial_Condition=World Ocean Atlas 2023 climatology 1991-2020 --glb_att_add MPAS_Mesh_compass_Version=1.4.0a6 --glb_att_add MPAS_Mesh_JIGSAW_Version=not found --glb_att_add MPAS_Mesh_JIGSAW_Python_Version=1.0.0 --glb_att_add MPAS_Mesh_MPAS_Tools_Version=0.34.1 --glb_att_add MPAS_Mesh_NCO_Version=5.2.6 --glb_att_add MPAS_Mesh_ESMF_Version=8.6.1 --glb_att_add MPAS_Mesh_geometric_features_Version=1.4.0 --glb_att_add MPAS_Mesh_Metis_Version=5.1.0 --glb_att_add MPAS_Mesh_pyremap_Version=1.3.0 /lustre/scratch5/dschlichting/runs/gom1pt5r2/ocean/global_ocean/GoM5/WOA23/init/initial_state/initial_state.nc /lustre/scratch5/dschlichting/runs/gom1pt5r2/ocean/global_ocean/GoM5/WOA23/init/initial_state/initial_state_with_metadata.nc
        mpirun -n 512 ./ocean_model
    mesh_spec :
        1.0
    Conventions :
        MPAS
    source :
        MpasMeshConverter.x
    file_id :
        zgvntg2vxs

Working example with scatter:

t = 110;nrows=1;ncols=1
fig, ax = plt.subplots(nrows,ncols, figsize = (8,4), 
                       subplot_kw={'projection': crs}, 
                       constrained_layout = True, dpi = 300)

m1 = ax.scatter(lon_c[iGoM[0]],lat_c[iGoM[0]],c=rv[t,iGoM[0]], s=1,
                cmap = cmo.balance, vmin = -1, vmax = 1,transform=transform)

fig.colorbar(m1, ax = ax, label = r'$(\partial_x v - \partial_y u)f^{-1}$', extend = 'both')

ax.set_extent([-98, -70, 13, 35], ccrs.PlateCarree())
ax.add_feature(land_10m, facecolor='0.8')
ax.set_aspect(lat_rad)
ax.coastlines(resolution='10m', linewidth = 0.4)  

gl = ax.gridlines(linewidth=0.75, color='black', alpha=0.5, linestyle='dotted', draw_labels=True)
gl.xformatter = LONGITUDE_FORMATTER
gl.yformatter = LATITUDE_FORMATTER
gl.right_labels = False
gl.left_labels = True
gl.top_labels = False
gl.bottom_labels = True
    
ax.set_title(r'Surface $\zeta f^{-1}$: '+ str(ds.xtime[t].values)[12:-2])
fig.canvas.draw()

Plot with mosaic

t = 110;nrows=1;ncols=1
fig, ax = plt.subplots(nrows,ncols, figsize = (8,4), 
                       subplot_kw={'projection': crs}, 
                       constrained_layout = True, dpi = 300)

descriptor = mosaic.Descriptor(dsi, projection, transform)
m1 = mosaic.polypcolor(ax, descriptor, rv[t,iGoM[0]], 
                       antialiaseds=False, cmap = cmo.balance, vmin = -1, vmax = 1)

fig.colorbar(m1, ax = ax, label = r'$(\partial_x v - \partial_y u)f^{-1}$', extend = 'both')

ax.set_extent([-98, -70, 13, 35], ccrs.PlateCarree())
ax.add_feature(land_10m, facecolor='0.8')
ax.set_aspect(lat_rad)
ax.coastlines(resolution='10m', linewidth = 0.4)  

gl = ax.gridlines(linewidth=0.75, color='black', alpha=0.5, linestyle='dotted', draw_labels=True)
gl.xformatter = LONGITUDE_FORMATTER
gl.yformatter = LATITUDE_FORMATTER
gl.right_labels = False
gl.left_labels = True
gl.top_labels = False
gl.bottom_labels = True
    
ax.set_title(r'Surface $\zeta f^{-1}$: '+ str(ds.xtime[t].values)[12:-2])
fig.canvas.draw()

[andrewdnolan]

Hey @dylanschlichting! Thanks for using mosaic and putting it through it's paces.

Yes, mosaic is capable of doing the regional plotting you're looking for but, it does not per se support regional masks. Instead it supports regionally masked datasets. In short, you'll need to make the Descriptor object from a masked dataset, so that the descriptor and the data array being plotted have the same dimensions (i.e. same nVertices).

So, all you'd need to change is the reading in of the dataset bloc of code to be something like:

#GoM regional mask
pathm = '/usr/projects/climate/dschlichting/repos/analysis_notebooks/gom1pt5/gom1pt5_gom_mask.nc'
mask = xr.open_dataset(pathm)
iGoM = mask.regionCellMasks[:, 0] == 1
# indices of cells within regional mask
idxGoM = mask.nCells.where(iGoM, drop=True).astype(int)

#HF output
path = '/lustre/scratch5/dschlichting/E3SM/scratch/chicoma-cpu/Tl319_GoM1pt5_GMPAS-IAF_norivers/run/Tl319_GoM1pt5_GMPAS-IAF.mpaso.hist.am.highFrequencyOutput.1997-0*-01_00.00.00.nc'
ds = xr.open_mfdataset(glob.glob(path), concat_dim="Time", combine = 'nested')
ds = ds.sortby(ds.xtime)
ds = ds.isel(nCells = idxGoM)
ds
#Initial condition
pathi = '/lustre/scratch5/dschlichting/runs/gom1pt5r2/ocean/global_ocean/GoM5/WOA23/init/initial_state/initial_state.nc'
dsi = xr.open_dataset(pathi)
dsi = dsi.isel(nCells = idxGoM)

where I've moved the reading in of the mask to the top, and added one final step where the region mask gets converted to an index array of cells within the region (idxGoM). I then use that index array to subset the high frequency output and the initial condition (i.e. ds = ds.isel(nCells = idxGoM)).

Then the mosaic block you have above should pretty much work as expected, all you'd need to do is remove the masking of data array in the call to mosaic.polypcolor. In other words, that one call should now look like:

...
m1 = mosaic.polypcolor(ax, descriptor, rv.isel(Time=t),
                       antialiaseds=False, cmap = cmo.balance, vmin = -1, vmax = 1)
...

That should produce:

[andrewdnolan]

Thanks for bringing this up! While mosaic does support this functionality, it's not entirely clear how to go about doing it from the existing documentation. I think you're question hints towards two easy but important improvements we should make:

1. Raise an error when the data array based to mosaic.polypcolor does not match the dimension of the dataset used to create the Descriptor. This should catch the regional sub-setting confusion you ran into above, and more generally catch other possible mismatched dims errors.
2. If I can hunt down an existing region mask on the LCRC database, I think it'd be useful to explicitly advise users on how to go about regional plotting in the documentation.

[xylar]

Here are some existing region masks:
https://web.lcrc.anl.gov/public/e3sm/diagnostics/mpas_analysis/region_masks/
Let me know if I can help you find one that matches a mesh you want to test with.