Basic MPAS Analysis & Visualization with UXarray

In this section, you’ll learn:¶

How to interact with MPAS data: retrieving, subsetting and storing data
How to visualize horizontal data, i.e. data from a single vertical level

UxDataset Data Structure

Prerequisites¶

Concepts	Importance
UXarray	Necessary
Unstructured Grid	Helpful

Time to learn: 15 minutes

Import packages¶

We use Cartopy, GeoViews, Holoviews and Matplotlib packages for visualization. Xarray, UXarray, NumPy and GeoPandas are used for the data structures we use to represent and manipulate the data.

# autoload external python modules if they changed
%load_ext autoreload
%autoreload 2

# add ../funcs to the current path
import sys, os
sys.path.append(os.path.join(os.getcwd(), "..")) 

# import modules
import warnings
import math

import cartopy.crs as ccrs
import geoviews as gv
import geoviews.feature as gf
import holoviews as hv
import hvplot.xarray
from holoviews import opts
import matplotlib.pyplot as plt

import s3fs

import geopandas as gp
import numpy as np
import uxarray as ux
import xarray as xr

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[1], line 14
     10 import warnings
     11 import math
     12 
     13 import cartopy.crs as ccrs
---> 14 import geoviews as gv
     15 import geoviews.feature as gf
     16 import holoviews as hv
     17 import hvplot.xarray

ModuleNotFoundError: No module named 'geoviews'

Retrieving the state border lines and coastlines¶

coast_lines = gf.coastline(projection=ccrs.PlateCarree(), line_width=1, scale="50m")
state_lines = gf.states(projection=ccrs.PlateCarree(), line_width=1, line_color='gray', scale="50m")

Helper functions¶

The following functions are used for visualizing the data. The horizontal_contour function generates the contour map for a given slice of data.

# Generates a contour plot for a horizontal slice
def horizontal_contour(ux_hslice, title, cmin=None, cmax=None, width=800, height=500, clevs=20, cmap="coolwarm", symmetric_cmap=False):

    # Get min and max
    amin = ux_hslice.min().item()
    amax = ux_hslice.max().item()
    
    cmin = math.floor(amin) if(cmin is None) else cmin
    cmax = math.ceil(amax) if(cmax is None) else cmax
    
    if symmetric_cmap:  # get a symmetric color map when cmin < 0, cmax >0
        cmax = max(abs(cmin), cmax)
        cmin = -cmax

    if isinstance(cmap, str):
        cmap = plt.get_cmap(cmap)

    # Generate contour plot
    title = f" min={amin:.1f} max={amax:.1f}"
    
    contour_plot = hv.operation.contours(
        ux_hslice.plot(),
        levels=np.linspace(cmin, cmax, num=clevs),  # levels=np.arange(cmin, cmax, 0.5)
        filled=True
    ).opts(
        line_color=None,  # line_width=0.001
        width=width, height=height,
        cmap='coolwarm', clim=(cmin, cmax),
        colorbar=True, show_legend=False,
        tools=['hover'], title=title
    )

    return contour_plot

Retrieve/load MPAS/JEDI data¶

The example MPAS/JEDI data are stored at jetstream2. We need to retreive those data first. There are two ways to retrieve MPAS data:

Download all example data from JetStream2 to local and them load them locally. This approach allows downloading the data once per machine and reuse it in notebooks.
Stream the JetStream2 S3 objects on demand. In this case, each notebook (including restarting a notebook) will retrieve the required data separately as needed.

# choose the data_load_method, check the above cell for details. Default to method 2; choose 1 if running from you own machine.
data_load_method = 2  # 1 or 2

Method 1: Download all example data once and reuse it in mulptile notebooks¶

%%time
local_dir="/tmp"

if data_load_method == 1 and not os.path.exists(local_dir + "/conus12km/bkg/mpasout.2024-05-06_01.00.00.nc"):
    jetstream_url = 'https://js2.jetstream-cloud.org:8001/'
    fs = s3fs.S3FileSystem(anon=True, asynchronous=False,client_kwargs=dict(endpoint_url=jetstream_url))
    conus12_path = 's3://pythia/mpas/conus12km'
    fs.get(conus12_path, local_dir, recursive=True)
    print("Data downloading completed")
else:
    print("Skip..., either data is available in local or data_load_method is NOT 1")

# Set file path
if data_load_method == 1:
    grid_file = local_dir + "/conus12km/conus12km.invariant.nc_L60_GFS"
    ana_file = local_dir + "/conus12km/bkg/mpasout.2024-05-06_01.00.00.nc"
    bkg_file = local_dir + "/conus12km/ana/mpasout.2024-05-06_01.00.00.nc"
    # jdiag_file = local_dir + "/conus12km/jdiag_aircar_t133.nc"  #q133.nc or uv233.nc

Method 2: Stream the JetStream2 S3 objects on demand¶

%%time
if data_load_method == 2:
    jetstream_url = 'https://js2.jetstream-cloud.org:8001/'
    fs = s3fs.S3FileSystem(anon=True, asynchronous=False,client_kwargs=dict(endpoint_url=jetstream_url))
    conus12_path = 's3://pythia/mpas/conus12km'
    
    grid_url=f"{conus12_path}/conus12km.invariant.nc_L60_GFS"
    bkg_url=f"{conus12_path}/bkg/mpasout.2024-05-06_01.00.00.nc"
    ana_url=f"{conus12_path}/ana/mpasout.2024-05-06_01.00.00.nc"
    # jdiag_url=f"{conus12_path}/jdiag_aircar_t133.nc"
    
    grid_file = fs.open(grid_url)
    ana_file = fs.open(ana_url)
    bkg_file = fs.open(bkg_url)
    # jdiag_file = fs.open(jdiag_url)
else:
    print("Skip..., data_load_method is NOT 2")

Loading the data into UXarray datasets¶

We use the UXarray data structures to work with the data. This package supports data defined over unstructured grid and provides utilities for modifying and visualizing it. The available functionality are discussed in UxDataset documentation.

uxds = ux.open_dataset(grid_file, bkg_file)

Regional temperature contour¶

We use the theta (potential temperature) variable from this dataset, which has a (Time, n_face, nVertLevels), i.e. Time * Number of grid faces * Number of vertical levels dimensionality to have a look at a regional, horizontal plot.

uxvar = uxds['theta'] - 273.15   ## Kelvin to Celsius

The data has multiple vertical levels, representing the values for different elevations. In this section, we are focusing on a single level. Next, we are plotting the contour of theta values in the data for this fixed level.

i_lev = 0   # `nVertLevels` index, for picking the 0th level
i_time = 0  # `Time` index
lat_val, long_val = 42.63, -83.3 # The latitude and longitude levels we will use in the next section
lat_line_lims, long_line_lims = [-135, -60], [20, 55] # limits for the cross-section lines
plot_x_lims, plot_y_lims = (-140, -55), (15,60) # limits for the plot

plot = horizontal_contour(uxvar.isel(Time=0, nVertLevels=i_lev), title=f'Contour plot for potential temperature over a region at vertical level 0') #, symmetric_cmap=True)
lat_line = hv.Curve((lat_line_lims, [lat_val, lat_val]))
long_line = hv.Curve(([long_val, long_val], long_line_lims))
overlay = (plot * coast_lines * state_lines * lat_line * long_line).redim.range(x=plot_x_lims, y=plot_y_lims)
overlay.opts(
    opts.Curve(color='black', line_width=3)
)

Basic MPAS Analysis & Visualization with UXarray

In this section, you’ll learn:¶

Related Documentation¶

Prerequisites¶

Import packages¶

Retrieving the state border lines and coastlines¶

Helper functions¶

Retrieve/load MPAS/JEDI data¶

Method 1: Download all example data once and reuse it in mulptile notebooks¶

Method 2: Stream the JetStream2 S3 objects on demand¶

Loading the data into UXarray datasets¶

Regional temperature contour¶