Skip to article frontmatterSkip to article content

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

Prerequisites

ConceptsImportance
UXarrayNecessary
Unstructured GridHelpful

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
Loading...

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:

  1. 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.
  2. 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")
Skip..., either data is available in local or data_load_method is NOT 1
CPU times: user 218 μs, sys: 22 μs, total: 240 μs
Wall time: 193 μs

# 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")
CPU times: user 74.1 ms, sys: 23.2 ms, total: 97.3 ms
Wall time: 218 ms

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)
)
Loading...
Introduction
RRFSv2 Overview
MPAS(Regional) Analysis and Visualization
Advanced MPAS Analysis & Visualization with UXarray