Plotting spectral element data

[mvertens]

1. Using uxarray to plot out the data from native output(thanks to Will Wieder)

import dask
import numpy as np
import pandas as pd
import xarray as xr
import uxarray as ux  #need npl 2024a

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.pyplot import cm
import cartopy.crs as ccrs
Then point to mesh and data files & plot

mesh = '/glade/campaign/cesm/cesmdata/inputdata/share/meshes/ne30pg3_ESMFmesh_cdf5_c20211018.nc'
fin='/glade/derecho/scratch/slevis/temp_work/new_rawdata/tools/mksurfdata_esmf/surfdata_ne30np4.pg3_hist_1850_78pfts_c240905.nc'
ds = ux.open_dataset(mesh, fin)
ds.mapunit.plot()

[mvertens]

@maritsandstad - could you add your documentation for plotting using XESMF for mapping as an alternative method above. That would be super helpful!

[maritsandstad]

2. xesmf based approach

Intial infrastructure setup (works on sigma2 machines)

Do this every time:

 module load Anaconda3/2022.05
module load ESMF/8.4.1-intel-2022a

First time setup and install your xesmf environment according to this guide: https://xesmf.readthedocs.io/en/stable/installation.html

After that just load it before you start

conda activate xesmf_env

Code library functions:

import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import numpy as np
import xarray as xr
import xesmf

def make_se_regridder(weight_file):
    weights = xr.open_dataset(weight_file)
    in_shape = weights.src_grid_dims.load().data

    # Since xESMF expects 2D vars, we'll insert a dummy dimension of size-1
    if len(in_shape) == 1:
        in_shape = [1, in_shape.item()]

    # output variable shape
    out_shape = weights.dst_grid_dims.load().data.tolist()[::-1]

    print(in_shape, out_shape)

    dummy_in = xr.Dataset(
        {
            "lat": ("lat", np.empty((in_shape[0],))),
            "lon": ("lon", np.empty((in_shape[1],))),
        }
    )
    dummy_out = xr.Dataset(
        {
            "lat": ("lat", weights.yc_b.data.reshape(out_shape)[:, 0]),
            "lon": ("lon", weights.xc_b.data.reshape(out_shape)[0, :]),
        }
    )

    regridder = xesmf.Regridder(
        dummy_in,
        dummy_out,
        weights=weight_file,
        method="bilinear",
        reuse_weights=True,
        periodic=True,
    )
    return regridder


def regrid_se_data(regridder, data_to_regrid):
    print(data_to_regrid.dims)
    if isinstance(data_to_regrid, xr.DataArray):
        print(type(data_to_regrid))
        updated = data_to_regrid.copy().transpose(..., "lndgrid").expand_dims("dummy", axis=-2)
    else:
        vars_with_ncol = [name for name in data_to_regrid.variables if "ncol" in data_to_regrid[name].dims]
        updated = data_to_regrid.copy().update(
            data_to_regrid[vars_with_ncol].transpose(..., "lndgrid").expand_dims("dummy", axis=-2)
        )
    regridded = regridder(updated.rename({"dummy": "lat", "lndgrid": "lon"}))
    return regridded

Code to use library functions:

And importing them from a different file or using them inside the same file I can get regridded data using

# Change this to correct weight file
weight_file = "/cluster/shared/noresm/inputdata/lnd/clm2/mappingdata/maps/ne30np4/map_ne30np4_to_0.5x0.5rtm_aave_da_110320.nc"
regridder = make_se_regridder(weight_file=weight_file)

ds = xr.open_dataset("/path/to/datafile")
# Beware that you might have to rename the spectral element dimension to "lndgrid" before this step if it's  not already called that
vardata = regrid_se_data(regridder, ds.varname)

# And now vardata is on 0.5by0.5 latlon regular grid and can be plotted with cartopy as usual

[adagj]

1. Using uxarray to plot out the data from native output(thanks to Will Wieder)

import dask
import numpy as np
import pandas as pd
import xarray as xr
import uxarray as ux  #need npl 2024a

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.pyplot import cm
import cartopy.crs as ccrs
Then point to mesh and data files & plot

mesh = '/glade/campaign/cesm/cesmdata/inputdata/share/meshes/ne30pg3_ESMFmesh_cdf5_c20211018.nc'
fin='/glade/derecho/scratch/slevis/temp_work/new_rawdata/tools/mksurfdata_esmf/surfdata_ne30np4.pg3_hist_1850_78pfts_c240905.nc'
ds = ux.open_dataset(mesh, fin)
ds.mapunit.plot()

This works great!