UXarray with CESM-HR on the ne120 Grid

UXarray with CESM-HR on the ne120 Grid#

One month of CESM high-resolution atmosphere output, on the model’s own ne120 grid (777,602 cells). Each section shows one UXarray feature, with plots side by side to compare its options:

  1. Reading the data

  2. Global maps

  3. Subsetting

  4. Points vs polygons

  5. Zonal statistics

  6. Grid topology

  7. Cross-sections

  8. Neighborhood filters, weighted means and remapping

import warnings

import cartopy.crs as ccrs
import cftime
import geoviews as gv
import geoviews.feature as gf
import holoviews as hv
import numpy as np
import xarray as xr
from holoviews.operation.datashader import rasterize

import uxarray as ux

warnings.filterwarnings("ignore", message=r".*(No spatial coordinate|np\.dot\(\)|IProgress not found)")

1. Reading the Data#

Monthly atmosphere output of member 010 of MESACLIP, the CESM-HR historical ensemble (GDEX d651007), and the SCRIP file with the shape of each of the model’s columns

# SCRIP grid file: the polygon around each of the 777,602 model columns
grid_path = "/glade/campaign/cesm/cesmdata/cseg/inputdata/share/scripgrids/ne120np4_pentagons_100310.nc"
# Kerchunk references to member 010's monthly atmosphere output, 1920-2005
data_path = "/gdex/data/d651007/kerchunk/b.e13.BHISTC5.ne120_t12.cesm-ihesp-hires1.0.46-1920-2005.010.atm.month_1.parq"

Xarray opens all 86 years lazily; only the selected month is read. CAM stamps monthly means at the end of the month, so 2005-08-01 is July 2005.

ux.open_dataset attaches the grid. The column dimension ncol becomes n_face, UXarray’s name for cells. Arithmetic keeps the grid attached, so converting the temperatures from kelvin to °C is one subtraction.

%%time
# Attach the grid to the data; ncol becomes n_face
uxds = ux.open_dataset(
    grid_path, 
    data_path, 
    engine="kerchunk", 
    chunks={}, 
    decode_timedelta=False
)

uxds = uxds[["TS", "T", "U"]].sel(time=cftime.DatetimeNoLeap(2005, 8, 1), method="nearest").load()

# The variables used below, with temperatures converted from kelvin to °C
ts = uxds["TS"] - 273.15  # surface temperature (°C)
t = uxds["T"] - 273.15  # air temperature on 30 levels (°C)
u = uxds["U"]  # eastward wind on 30 levels (m/s)
uxgrid = uxds.uxgrid  # the grid, shared by every variable

uxds

CPU times: user 2.72 s, sys: 522 ms, total: 3.24 s
Wall time: 5.24 s

<xarray.UxDataset> Size: 190MB
Dimensions:  (n_face: 777602, lev: 30)
Coordinates:
  * lev      (lev) float64 240B 3.643 7.595 14.36 24.61 ... 957.5 976.3 992.6
    time     object 8B 2005-08-01 00:00:00
Dimensions without coordinates: n_face
Data variables:
    TS       (n_face) float32 3MB 290.0 289.8 289.2 289.1 ... 298.3 298.5 297.9
    T        (lev, n_face) float32 93MB 226.2 226.1 225.9 ... 298.1 298.3 298.0
    U        (lev, n_face) float32 93MB 83.3 83.54 83.93 ... -0.349 -0.1629
Attributes:
    np:               4
    ne:               120
    Conventions:      CF-1.0
    source:           CAM
    case:             b.e13.BHISTC5.ne120_t12.cesm-ihesp-hires1.0.46-1920-200...
    title:            UNSET
    logname:          fredc
    host:             derecho1
    Version:          $Name$
    revision_Id:      $Id$
    initial_file:     B.E.13.BHISTC5.ne120_t12.sehires38.003.sunway.cam.i.192...
    topography_file:  /glade/campaign/cesm/cesmdata/inputdata/atm/cam/topo/US...

2. Global Maps#

.plot() shades every cell. Changing projection to Robinson

# Two map projections
robinson = ccrs.Robinson()
north_pole = ccrs.Orthographic(central_longitude=0, central_latitude=90)  # view from above the pole

# Plot styling shared by every global map of surface temperature
ts_range = (float(ts.min()), float(ts.max()))  # one color range
map_style = dict(cmap="inferno", clim=ts_range, clabel="Surface temperature (°C)")

%%time
# Map: whole globe, Robinson
robinson_map = ts.plot(
    projection=robinson,
    frame_width=560,
    colorbar=False,  # the polar map's colorbar covers both
    title="Robinson",
    **map_style)

CPU times: user 12.4 s, sys: 765 ms, total: 13.2 s
Wall time: 16 s

%%time
# Map: Northern Hemisphere, Orthographic
polar_map = ts.plot(
    projection=north_pole,
    frame_width=284,
    frame_height=284,
    title="Orthographic, over the North Pole",
    **map_style)

CPU times: user 6.51 s, sys: 362 ms, total: 6.87 s
Wall time: 7.08 s

%%time
# Add coastlines (* overlays one plot on another)
robinson_panel = robinson_map * gf.coastline(projection=robinson)
polar_panel = polar_map * gf.coastline(projection=north_pole)

# Show side by side (+) shared_axes=False gives each projection its own axes
(robinson_panel + polar_panel).opts(shared_axes=False)

CPU times: user 3.15 ms, sys: 247 μs, total: 3.39 ms
Wall time: 3.4 ms

3. Subsetting#

Three ways to select a region around Boulder, Colorado. Each returns the data on a smaller grid.

%%time
# Projection and region used in the rest of the notebook
plate_carree = ccrs.PlateCarree()  # longitude and latitude drawn as x and y
boulder = (-105.27, 40.01)  # (lon, lat)
box_bounds = dict(lon_bounds=(-111, -99.5), lat_bounds=(35.5, 44.5))

# Subset the data three ways
box = ts.subset.bounding_box(**box_bounds)  # cells inside a longitude-latitude box
circle = ts.subset.bounding_circle(center_coord=boulder, r=4)  # cells within 4° of Boulder
nearest = ts.subset.nearest_neighbor(center_coord=boulder, k=500)  # the 500 cells closest to Boulder

CPU times: user 25.6 s, sys: 514 ms, total: 26.1 s
Wall time: 35.4 s

%%time
# Plot styling shared by the three maps: same extent and color range
region_style = dict(
    projection=plate_carree,
    cmap="inferno",
    clim=(float(box.min()), float(box.max())),
    clabel="Surface temperature (°C)",
    xlim=(-112.5, -98),
    ylim=(34.5, 45.5),
    frame_width=260)

# Map: cells in the box
box_map = box.plot(
    title=f"bounding_box: {box.uxgrid.n_face:,} cells",
    colorbar=False,
    **region_style)

# Map: cells in the circle
circle_map = circle.plot(
    title=f"bounding_circle: {circle.uxgrid.n_face:,} cells",
    colorbar=False,
    **region_style)

# Map: the nearest cells
nearest_map = nearest.plot(
    title=f"nearest_neighbor: {nearest.uxgrid.n_face:,} cells",
    **region_style)

# Show side by side
box_map + circle_map + nearest_map

CPU times: user 2.44 s, sys: 208 ms, total: 2.65 s
Wall time: 3.11 s

4. Points vs Polygons#

.plot() draws polygons, the shape of each cell. .plot.points() draws one point per cell center, which is much quicker. Points leave gaps where there are more pixels than cells, which on the globe shows toward the poles.

# Plot styling shared by both maps
global_style = dict(projection=robinson, frame_width=420, colorbar=False, **map_style)
coastline = gf.coastline(projection=robinson)

%%time
# Map: every cell drawn as a polygon
polygons = ts.plot.polygons(
    title="Polygons",
    **global_style)

CPU times: user 6.41 s, sys: 439 ms, total: 6.85 s
Wall time: 6.85 s

%%time
# Map: one point per cell center, turned into an image the same way
points = ts.plot.points(
    rasterize=True,
    title="Points",
    **global_style)

CPU times: user 20.7 ms, sys: 11.8 ms, total: 32.5 ms
Wall time: 114 ms

%%time
# Show side by side, with coastlines
polygons * coastline + points * coastline

CPU times: user 2.39 ms, sys: 0 ns, total: 2.39 ms
Wall time: 2.4 ms

Zoomed in, each point covers one pixel of its cell; dynspread=True enlarges the points so they can be seen.

# Plot styling: about 5° by 4° around Boulder, same color range as the subsets
zoom_style = dict(
    projection=plate_carree,
    cmap="inferno",
    clim=region_style["clim"],
    clabel="Surface temperature (°C)",
    xlim=(-108, -102.5),
    ylim=(38, 42),
    frame_width=400,
    xlabel="Longitude",
    ylabel="Latitude")

# Map, zoomed in: polygons
zoomed_polygons = circle.plot.polygons(
    title="Polygons",
    colorbar=False,
    **zoom_style)

# Map, zoomed in: points
zoomed_points = circle.plot.points(
    rasterize=True,
    dynspread=True,  # enlarge each point by a few pixels
    title="Points",
    **zoom_style)

# Show side by side
zoomed_polygons + zoomed_points

5. Zonal Statistics#

zonal_mean averages along lines of latitude (here every 1°) or, with conservative=True, over latitude bands weighted by area (here 10° bands).

%%time
# Zonal means of surface temperature, two ways
ts_zonal = ts.zonal_mean(lat=(-90, 90, 1))  # along lines of latitude, every 1°
ts_bands = ts.zonal_mean(lat=(-90, 90, 10), conservative=True)  # area-weighted, over 10° bands

CPU times: user 39.4 s, sys: 377 ms, total: 39.8 s
Wall time: 23.2 s

# Axis labels for the HoloViews plots below
latitude = hv.Dimension("latitudes", label="Latitude")
temperature = hv.Dimension("TS", label="Surface temperature (°C)")

# Line plot: zonal mean every 1°
lines = hv.Curve(
    (ts_zonal.latitudes, ts_zonal),
    latitude,
    temperature,
    label="every 1°")

# Dot plot: zonal mean of each 10° band
bands = hv.Scatter(
    (ts_bands.latitudes, ts_bands),
    latitude,
    temperature,
    label="10° bands").opts(color="black", size=7)  # plot styling
# Overlay the two, with latitude on the vertical axis to line up with the map
zonal_profile = (lines * bands).opts(
    invert_axes=True,
    xlim=(-90, 90),
    frame_width=200,
    frame_height=300,
    legend_position="bottom",
    title="zonal_mean")

# Map: surface temperature, longitude and latitude as x and y
flat_map = ts.plot(
    projection=plate_carree,
    frame_width=600,
    colorbar=False,
    title="Surface temperature, July 2005",
    **map_style)

# Show the map with coastlines, and the zonal mean beside it
flat_map * gf.coastline(projection=plate_carree) + zonal_profile

zonal_anomaly subtracts the zonal mean from each cell, leaving the differences along each latitude, such as land against ocean. zonal_mean also works on fields with levels, giving a latitude–height section (rasterize draws it as an image).

%%time
# Zonal anomaly of surface temperature, and zonal mean of air temperature at every level
ts_anomaly = ts.zonal_anomaly(lat=(-90, 90, 1))  # each cell minus the mean of its latitude
t_zonal = t.zonal_mean(lat=(-90, 90, 2))  # every level at once: (lev, latitudes)

CPU times: user 15.1 s, sys: 123 ms, total: 15.2 s
Wall time: 6.43 s

# Map: zonal anomaly of surface temperature
anomaly_map = ts_anomaly.plot(
    projection=plate_carree,
    cmap="RdBu_r",
    clim=(-15, 15),
    clabel="Difference from the zonal mean (°C)",
    frame_width=500,
    title="zonal_anomaly")

# Latitude-height plot: zonal-mean air temperature, one box per latitude and level
level = hv.Dimension("lev", label="Model level (≈ pressure, hPa)")
t_mesh = hv.QuadMesh((t_zonal.latitudes, t_zonal.lev, t_zonal), [latitude, level], "T")
t_section = rasterize(t_mesh).opts(  # draw as an image, plus plot styling
    cmap="inferno",
    colorbar=True,
    clabel="Air temperature (°C)",
    invert_yaxis=True,  # surface at the bottom
    frame_width=300,
    frame_height=250,
    tools=["hover"],
    title="zonal_mean of air temperature")

# Show side by side
anomaly_map * gf.coastline(projection=plate_carree) + t_section

6. Grid Topology#

The grid is stored once in uxds.uxgrid and can be plotted and queried on its own.

uxgrid

<uxarray.Grid>
Original Grid Type: Scrip
Grid Dimensions:
  * n_node: 780456
  * n_face: 777602
  * n_max_face_nodes: 5
Grid Coordinates (Spherical):
  * node_lon: (780456,)
  * node_lat: (780456,)
  * face_lon: (777602,)
  * face_lat: (777602,)
Grid Coordinates (Cartesian):
  * node_x: (780456,)
  * node_y: (780456,)
  * node_z: (780456,)
Grid Connectivity Variables:
  * face_node_connectivity: (777602, 5)
Grid Descriptor Variables:
  * face_areas: (777602,)
  * n_nodes_per_face: (777602,)

At each corner of the cube, three faces meet. plot.edges, plot.nodes and plot.face_centers draw the parts of the mesh; shading by face_areas shows the pattern inside each element, whose smallest cells are at the element corners.

# Subset the grid itself around a corner of the cube
cube_corner = (45, 35.26)  # (lon, lat) where three faces of the cube meet
corner_grid = uxgrid.subset.bounding_circle(center_coord=cube_corner, r=1)

# Mesh plot: edges, nodes (cell corners) and face centers, overlaid
edges = corner_grid.plot.edges(projection=plate_carree, color="gray")
nodes = corner_grid.plot.nodes(projection=plate_carree, color="black", size=5)
centers = corner_grid.plot.face_centers(projection=plate_carree, color="crimson", size=5)
mesh = (edges * nodes * centers).opts(  # plot styling
    frame_width=330,
    title="Edges, nodes (black) and face centers (red)")

# Map: cell areas in km² (face_areas is on a sphere of radius 1), with the edges on top
cell_area = ux.UxDataArray(uxgrid.face_areas * 6371.0**2, uxgrid=uxgrid, name="area")
area_map = cell_area.subset.bounding_circle(center_coord=cube_corner, r=1).plot(
    projection=plate_carree,
    cmap="viridis",
    clabel="Cell area (km²)",
    frame_width=330,
    title="face_areas")
white_edges = corner_grid.plot.edges(projection=plate_carree, color="white")

# Show side by side
mesh + area_map * white_edges

plot.face_area_distribution shows all the cell areas at once. get_faces_containing_point finds the cell under a point, and face_face_connectivity lists the cells that share an edge with it.

%%time
# Find the cell under Boulder and the cells that share an edge with it
cell = uxgrid.get_faces_containing_point(boulder, return_counts=False)[0][0]
neighbors = uxgrid.face_face_connectivity.values[cell]
neighbors = neighbors[neighbors >= 0]  # drop unused slots (negative fill value)

CPU times: user 8.69 s, sys: 320 ms, total: 9.01 s
Wall time: 9.26 s

# Histogram: area of every cell in the grid
histogram = uxgrid.plot.face_area_distribution(
    bins=60,
    xlabel="Cell area (sphere of radius 1)",
    frame_width=380,
    frame_height=250,
    title="face_area_distribution")

# Grids to draw: the mesh around Boulder, the neighbors, and the cell
local_grid = uxgrid.subset.bounding_circle(center_coord=boulder, r=0.75)
neighbor_grid = uxgrid.isel(n_face=neighbors)  # isel selects cells by index
cell_grid = uxgrid.isel(n_face=[cell])

# Mesh plot: the cell (red) and its neighbors (blue) on the surrounding mesh (gray)
local_edges = local_grid.plot.edges(projection=plate_carree, color="lightgray")
neighbor_edges = neighbor_grid.plot.edges(projection=plate_carree, color="steelblue", line_width=3)
cell_edges = cell_grid.plot.edges(projection=plate_carree, color="crimson", line_width=3)
neighbor_map = (local_edges * neighbor_edges * cell_edges).opts(  # plot styling
    xlim=(-106.2, -104.3),
    ylim=(39.3, 40.7),
    frame_width=340,
    title=f"Cell {cell} (red) and its neighbors (blue)")

# Show side by side
histogram + neighbor_map

7. Cross-Sections#

cross_section samples a field along a line of latitude (lat=), a line of longitude (lon=) or the great circle between two points (start=, end=). For a field with levels the result is a vertical section.

%%time
# Three cross-sections: along 40°N, along 100°W, and from San Francisco to New York
along_40n = ts.cross_section(lat=40, steps=1441).sortby("lon")  # around the globe, every 0.25°
along_100w = u.cross_section(lon=-100, steps=721)  # pole to pole; U has levels: a vertical section
san_francisco, new_york = (-122.42, 37.77), (-74.01, 40.71)
sf_to_ny = ts.cross_section(start=san_francisco, end=new_york, steps=400)  # along the great circle

CPU times: user 1.86 s, sys: 254 ms, total: 2.11 s
Wall time: 2.79 s

# Plot styling and axis label shared by the panels below
panel_size = dict(frame_width=400, frame_height=220)
longitude = hv.Dimension("lon", label="Longitude")
coast = gf.coastline(projection=plate_carree)

# Map: the three lines drawn on the surface temperature map from section 5
path_40n = gv.Path((along_40n.lon, along_40n.lat)).opts(color="deepskyblue", line_width=2)
path_100w = gv.Path((along_100w.lon, along_100w.lat)).opts(color="white", line_width=2)
path_arc = gv.Path((sf_to_ny.lon, sf_to_ny.lat)).opts(color="lime", line_width=3)
path_map = (flat_map * coast * path_40n * path_100w * path_arc).opts(
    frame_width=400,
    title="The three lines")

# Line plot: surface temperature along 40°N
lat_line = hv.Curve((along_40n.lon, along_40n), longitude, temperature).opts(  # plot styling
    color="deepskyblue",
    title="lat=40: surface temperature",
    **panel_size)

# Latitude-height plot: eastward wind along 100°W, one box per sample and level
u_mesh = hv.QuadMesh((along_100w.lat, along_100w.lev, along_100w), [latitude, level], "U")
lon_section = rasterize(u_mesh).opts(  # draw as an image, plus plot styling
    cmap="RdBu_r",
    symmetric=True,  # white at 0 m/s
    colorbar=True,
    clabel="Eastward wind (m/s)",
    invert_yaxis=True,
    tools=["hover"],
    title="lon=-100: eastward wind",
    **panel_size)

# Line plot: surface temperature from San Francisco to New York
arc_line = hv.Curve((sf_to_ny.lon, sf_to_ny), longitude, temperature).opts(  # plot styling
    color="lime",
    title="start, end: San Francisco to New York",
    **panel_size)

# Show as a 2 x 2 grid; shared_axes=False keeps each panel's own longitude range
(path_map + lat_line + lon_section + arc_line).opts(shared_axes=False).cols(2)

8. More Functions#

Neighborhood Filters#

neighborhood(r=1) groups each cell with the cells within 1° of it; reductions such as .mean() (a smoothing filter) and .std() (local variability) then run over each group.

%%time
# Neighborhood filter over the contiguous US: mean and standard deviation within 1°
usa = ts.subset.bounding_box(lon_bounds=(-125, -66), lat_bounds=(24, 50))
nearby = usa.neighborhood(r=1)  # for each cell, the cells within 1°
smoothed = nearby.mean()  # smoothing filter
spread = nearby.std()  # local variability

CPU times: user 1.89 s, sys: 84.5 ms, total: 1.98 s
Wall time: 2.45 s

# Plot styling shared by the three maps
usa_style = dict(projection=plate_carree, clabel="°C", xlim=(-125, -66), ylim=(24, 50), frame_width=250)
usa_range = (float(usa.min()), float(usa.max()))  # shared by the first two panels

# Map: surface temperature
original_map = usa.plot(
    cmap="inferno",
    clim=usa_range,
    colorbar=False,
    title="Surface temperature",
    **usa_style)

# Map: neighborhood mean (smoothed)
smoothed_map = smoothed.plot(
    cmap="inferno",
    clim=usa_range,
    title="neighborhood(r=1).mean()",
    **usa_style)

# Map: neighborhood standard deviation (local variability)
spread_map = spread.plot(
    cmap="viridis",
    clim=(0, float(spread.max())),  # its own range, or it would share the temperatures'
    title="neighborhood(r=1).std()",
    **usa_style)

# Show side by side
original_map + smoothed_map + spread_map

Area-Weighted Means#

weighted_mean weights each cell by its area; combined with subset it gives regional means in one line each. For a field with levels it averages each level.

# Area-weighted mean surface temperature of five regions
all_lon = (-180, 180)
regions = {
    "Globe": ts,
    "Tropics": ts.subset.bounding_box(lon_bounds=all_lon, lat_bounds=(-23.5, 23.5)),
    "N. Hemisphere": ts.subset.bounding_box(lon_bounds=all_lon, lat_bounds=(0, 90)),
    "S. Hemisphere": ts.subset.bounding_box(lon_bounds=all_lon, lat_bounds=(-90, 0)),
    "Contiguous US": usa,
}
region_means = [(name, float(sub.weighted_mean())) for name, sub in regions.items()]

# Bar chart: weighted mean of each region
bars = hv.Bars(region_means, "Region", temperature).opts(  # plot styling
    frame_width=400,
    frame_height=250,
    title="weighted_mean of each region")

# Line plot: global-mean air temperature at each level
t_profile = t.weighted_mean()  # one weighted mean per level
air_temperature = hv.Dimension("T", label="Air temperature (°C)")
profile = hv.Curve((t_profile.lev, t_profile), level, air_temperature).opts(  # plot styling
    invert_xaxis=True,  # surface (about 1000 hPa) on the left, top of the model on the right
    padding=(0, 0.05),
    frame_width=320,
    frame_height=250,
    title="weighted_mean at each level")

# Show side by side
bars + profile

Remapping#

remap moves data onto any grid UXarray can build: here a 1° latitude–longitude grid and a HEALPix grid. nearest_neighbor takes the value of the closest ne120 cell (inverse_distance_weighted and bilinear are the other methods).

%%time
# Two destination grids: 1° latitude-longitude and HEALPix
one_degree = ux.Grid.from_structured(lon=np.arange(-179.5, 180, 1.0), lat=np.arange(-89.5, 90, 1.0))
healpix = ux.Grid.from_healpix(zoom=6, pixels_only=False)  # pixels_only=False builds cell corners

CPU times: user 170 ms, sys: 32 ms, total: 202 ms
Wall time: 1.12 s

%%time
# Remap surface temperature onto each
ts_1deg = ts.remap.nearest_neighbor(destination_grid=one_degree, remap_to="faces")
ts_healpix = ts.remap.nearest_neighbor(destination_grid=healpix, remap_to="faces")

CPU times: user 855 ms, sys: 34.6 ms, total: 890 ms
Wall time: 792 ms

# Map: the box from section 3 on the original grid
native_map = box.plot(
    title=f"ne120: {uxgrid.n_face:,} cells",
    colorbar=False,
    **region_style)

# Map: the same box on the 1° grid
one_degree_map = ts_1deg.subset.bounding_box(**box_bounds).plot(
    title=f"1° lat-lon: {one_degree.n_face:,} cells",
    colorbar=False,
    **region_style)

# Map: the same box on the HEALPix grid
healpix_map = ts_healpix.subset.bounding_box(**box_bounds).plot(
    title=f"HEALPix: {healpix.n_face:,} cells",
    **region_style)

# Show side by side
native_map + one_degree_map + healpix_map