Calculus with uxarray

Calculus with uxarray#

(Previously was restricted to only 1D along grid dimension, but now can handle non-grid dimensions the way you would expect!)

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

# for visualizations
import holoviews as hv
hv.extension("bokeh")

# style: hide DataArray data values by default (click icon to expand)
xr.set_options(display_expand_data=False);
ⓘ

Setup:#

Helper for visualization#

def plot_vector_field(u, v, **kwargs):
    """plot a vector field from its two u, v components."""
    mag = np.hypot(u, v)
    angle = np.arctan2(v, u)
    return hv.VectorField(
        (u.uxgrid.face_lon, u.uxgrid.face_lat, angle, mag), **kwargs
    ).opts(magnitude="Magnitude", color="black")

Helpers for creating data#

(Realistically, most users will probably just ux.open_dataset() to load real data instead!)

# --- Constants ---
SCALE = 4.0                    # characteristic length scale to use as baseline.
LON0, LAT0 = 0.0, 30.0         # center not at 0, for fun :)
REGION = 7.0 * SCALE           # patch of 7 sigma x 7 sigma

# --- Grid: global HEALPix, then subset to a regional patch ---
grid_global = ux.Grid.from_healpix(zoom=6)
grid = grid_global.subset.bounding_box(
    lon_bounds=(LON0 - REGION/2, LON0 + REGION/2),
    lat_bounds=(LAT0 - REGION/2, LAT0 + REGION/2),
)
grid.sphere_radius = 1.0  # --> per-radian results. Use physical radius for per-meter.

# --- Helper for creating data ---
x = ux.UxDataArray(grid.face_lon.values - LON0, dims="n_face", uxgrid=grid, name="x")
y = ux.UxDataArray(grid.face_lat.values - LAT0, dims="n_face", uxgrid=grid, name="y")

def gaussian_vortex(xc, yc, sigma, sign, *, names=("u", "v"), x=x, y=y):
    """return u, v components of a gaussian vortex; clockwise when sign=-1.
    u = -sign*(y-yc)*g; v = +sign*(x-xc)*g; g = exp(-r^2 / 2 sigma^2)."""
    xd, yd = x - xc, y - yc   # distance from center
    g = np.exp(-(xd**2 + yd**2) / (2 * sigma**2))
    u = (-sign * yd * g).rename(names[0])
    v = (+sign * xd * g).rename(names[1])
    return u, v

grid

<uxarray.Grid>
Original Grid Type: HEALPix
Grid Dimensions:
  * n_node: 775
  * n_face: 702
  * n_max_face_nodes: 4
Grid Coordinates (Spherical):
  * node_lon: (775,)
  * node_lat: (775,)
  * face_lon: (702,)
  * face_lat: (702,)
Grid Coordinates (Cartesian):
  * node_x: (775,)
  * node_y: (775,)
  * node_z: (775,)
Grid Connectivity Variables:
  * face_node_connectivity: (702, 4)
Grid Descriptor Variables:
  * n_nodes_per_face: (702,)

grid.plot()

Basic case: curl of a single vortex#

# --- Make up some data: vortex with a Gaussian velocity field ---
# (Realistically, users will probably load real data with ds = ux.open_dataset(...),
# then use something like u = ds["zonal_wind"], v = ds["meridional_wind"], instead.)
u, v = gaussian_vortex(xc=0.0, yc=0.0, sigma=SCALE, sign=+1)
u

<xarray.UxDataArray 'u' (n_face: 702)> Size: 6kB
0.004534 8.175e-17 -0.004565 -0.01305 ... -0.5798 -0.3831 -0.3831 -0.2486
Dimensions without coordinates: n_face
data_plot = plot_vector_field(u, v).opts(aspect=1, arrow_heads=True)
data_plot

Take the curl, and plot results:#

# --- Curl along the grid dimension ---
curl = u.curl(v)
curl

<xarray.UxDataArray 'curl_u_v' (n_face: 702)> Size: 6kB
nan nan nan -7.849 nan -10.32 ... -13.37 -11.7 -11.89 -9.536 -9.536 -7.693
Dimensions without coordinates: n_face
Attributes:
    units:        1/m
    long_name:    Curl of (u, v)
    curl:         True
    description:  Curl of vector field computed as ∂v/∂x - ∂u/∂y + u·tan(φ)/a...
vmax = np.abs(curl).max().item()

curl_plot = curl.plot(cmap="RdBu_r", clim=(-vmax, vmax), aspect=1)

curl_plot * data_plot

There is a strong positive value (red) in the middle because the vortex is counterclockwise.

The white ring and light blue regions come from the underlying gaussian; more realistic vortices may fall off more slowly and just be entirely red.

The curl is NaN in the outer two layers of cells because they aren’t surrounded by neighbors with data. Rather than extrapolating, uxarray reports NaN here.

Data with more dimensions (“time”)#

Make up some data: two vortices, one moving downwards over time#

sigmaA = 1.0 * SCALE        # vortex A: larger
sigmaB = sigmaA / 1.5       # vortex B: smaller (1.5x ratio)

# --- Time axis and vortex paths ---
nt = 31  # number of timesteps for full dataset
t = np.linspace(0, 1, nt)
time = xr.DataArray(t, dims="time", coords={"time": t})

# Vortex A: stationary, middle left.
# Vortex B: starts top-right, moves down, ends bottom-right
xcA = -1.5 * SCALE
xcB = 1.7 * SCALE
ycA = 0.0
ycB = 0.3*REGION * (1 - 2 * time)    # dims (time,)

uA, vA = gaussian_vortex(xcA, ycA, sigmaA, +1)   # counter-clockwise
uB, vB = gaussian_vortex(xcB, ycB, sigmaB, -1)   # clockwise

u = uA + uB
v = vA + vB

u

<xarray.UxDataArray 'u' (n_face: 702, time: 31)> Size: 174kB
-0.002441 -0.004593 -0.008226 -0.01402 ... -0.0807 -0.0807 -0.0807 -0.0807
Coordinates:
  * time     (time) float64 248B 0.0 0.03333 0.06667 0.1 ... 0.9333 0.9667 1.0
Dimensions without coordinates: n_face

Curl at time=0#

(This feature has existed in uxarray for a while, same as above!)

u0, v0 = u.isel(time=0), v.isel(time=0)
curl_t0 = u0.curl(v0)
curl_t0

<xarray.UxDataArray 'curl_u_v' (n_face: 702)> Size: 6kB
nan nan nan 2.124 nan 4.687 5.747 ... -1.644 -9.456 -4.065 5.806 11.8 1.681 6.45
Coordinates:
    time     float64 8B 0.0
Dimensions without coordinates: n_face
Attributes:
    units:        1/m
    long_name:    Curl of (u, v)
    curl:         True
    description:  Curl of vector field computed as ∂v/∂x - ∂u/∂y + u·tan(φ)/a...
data_plot = plot_vector_field(u0, v0).opts(aspect=1, arrow_heads=True)
curl_plot = curl_t0.plot(cmap="RdBu_r", clim=(-vmax, vmax), aspect=1)
curl_plot * data_plot

Curl at all times#

(New feature, same syntax, loops across times like you would expect!)

curl = u.curl(v)
curl

<xarray.UxDataArray 'curl_u_v' (n_face: 702, time: 31)> Size: 174kB
nan nan nan nan nan nan nan ... -3.881 -3.881 -3.881 -3.881 -3.881 -3.881 -3.881
Coordinates:
  * time     (time) float64 248B 0.0 0.03333 0.06667 0.1 ... 0.9333 0.9667 1.0
Dimensions without coordinates: n_face
Attributes:
    units:        1/m
    long_name:    Curl of (u, v)
    curl:         True
    description:  Curl of vector field computed as ∂v/∂x - ∂u/∂y + u·tan(φ)/a...
# --- compare to a per-slice workaround ---
curl_workaround = ux.concat(
    [u.isel(time=i).curl(v.isel(time=i)) for i in range(nt)], dim="time"
)

# the values are the same:
diff = (curl_workaround - curl)
print('(min, max) of curl_workaround - curl: ', (diff.min().item(), diff.max().item()))
# (checking np.allclose would show False, though, because (NaN == NaN) is False.)

# here's what the workaround array looks like:
curl_workaround

(min, max) of curl_workaround - curl:  (0.0, 0.0)

<xarray.UxDataArray 'curl_u_v' (time: 31, n_face: 702)> Size: 174kB
nan nan nan 2.124 nan 4.687 5.747 ... -12.44 -9.271 -7.243 -4.403 -6.131 -3.881
Coordinates:
  * time     (time) float64 248B 0.0 0.03333 0.06667 0.1 ... 0.9333 0.9667 1.0
Dimensions without coordinates: n_face
Attributes:
    units:        1/m
    long_name:    Curl of (u, v)
    curl:         True
    description:  Curl of vector field computed as ∂v/∂x - ∂u/∂y + u·tan(φ)/a...

The workaround is easy enough to write for 1 dim, but imagine having multiple extra dims like “levels” and “time” too! Then it becomes more cumbersome. But, the u.curl(v) syntax will work for all cases, regardless of extra dims.

# --- plotting a few points in time:---
ncols = 2
frame_idx = [0, 7, 14, 21]
vmax = np.abs(curl).max().item()

panels = []
for k, i in enumerate(frame_idx):
    shading = curl.isel(time=i).plot(
        cmap="RdBu_r", clim=(-vmax, vmax), colorbar=(k % ncols == ncols - 1),
        title=f"t = {t[i]:.2f}", aspect=1,
    )
    arrows = plot_vector_field(u.isel(time=i), v.isel(time=i))
    panels.append(shading * arrows)

hv.Layout(panels).cols(ncols)

# --- Animation (with HoloViews' built-in playback controls) ---
vmax = np.abs(curl).max().item()

def frame(i):
    shading = curl.isel(time=i).plot(
        cmap="RdBu_r", clim=(-vmax, vmax), colorbar=True, aspect=1,
    )
    arrows = plot_vector_field(u.isel(time=i), v.isel(time=i))
    return shading * arrows

anim = hv.HoloMap({f"{t[i]:.3f}": frame(i) for i in range(nt)}, kdims="time")

hv.output(holomap="scrubber", fps=10)
anim

Can also apply other methods like gradient() and divergence(), too!#

Divergence#

div = u.divergence(v)
vmax = np.abs(div).max().item()

# --- Animation (with HoloViews' built-in playback controls) ---
def frame(i):
    shading = div.isel(time=i).plot(
        cmap="RdBu_r", clim=(-vmax, vmax), colorbar=True, aspect=1,
    )
    arrows = plot_vector_field(u.isel(time=i), v.isel(time=i))
    return shading * arrows

anim = hv.HoloMap({f"{t[i]:.3f}": frame(i) for i in range(nt)}, kdims="time")

hv.output(holomap="scrubber", fps=10)
anim

Note that the divergence limits here are much much smaller!

Why do the values decrease as the vortex moves towards the equator? It’s probably due to the metric term:

\[ \text{div}((u,v)) = \quad ... \quad - \frac{v \tan{\varphi}}{a} \]

where \(\varphi\) is the latitude and \(a\) is the sphere radius. Prompt: can you prove whether it is actually due to the metric term?

See the docs page for more details! https://uxarray.readthedocs.io/en/latest/user-guide/vector_calculus.html#divergence-f

Gradient#

Gradient takes a scalar field as an input and ouputs zonal and meridional components, so the visualization is a little different. For now, let’s just take the gradient of the magnitude of the field to see what kind of outputs it gives:

mag = np.sqrt(u**2 + v**2).rename("magnitude(u,v)")
mag

<xarray.UxDataArray 'magnitude(u,v)' (n_face: 702, time: 31)> Size: 174kB
0.002735 0.005388 0.01001 0.01766 0.02968 ... 0.0919 0.0919 0.0919 0.0919 0.0919
Coordinates:
  * time     (time) float64 248B 0.0 0.03333 0.06667 0.1 ... 0.9333 0.9667 1.0
Dimensions without coordinates: n_face
grad = mag.gradient()
grad

<xarray.UxDataset> Size: 348kB
Dimensions:              (time: 31, n_face: 702)
Coordinates:
  * time                 (time) float64 248B 0.0 0.03333 0.06667 ... 0.9667 1.0
Dimensions without coordinates: n_face
Data variables:
    zonal_gradient       (n_face, time) float64 174kB nan nan ... -2.563 -2.563
    meridional_gradient  (n_face, time) float64 174kB nan nan ... -3.655 -3.655
Attributes:
    gradient:  True

Now we once again have a vector field and a scalar field, but this time the calculus result is the vector field instead! We can plot it in a similar way, but let’s change the color and label things clearly to avoid confusion:

vmax_mag = np.abs(mag).max().item()

# --- Animation (with HoloViews' built-in playback controls) ---
def frame(i):
    mag_plot = mag.isel(time=i).plot(
        cmap="viridis", clim=(0, vmax_mag), colorbar=True, aspect=1,
        title="colors=magnitude; arrows=gradient",
    )
    zonal_grad_i = grad['zonal_gradient'].isel(time=i)
    merid_grad_i = grad['meridional_gradient'].isel(time=i)
    grad_arrows = plot_vector_field(zonal_grad_i, merid_grad_i).opts(line_color="red", line_width=1.5, scale=0.8)
    return mag_plot * grad_arrows

anim = hv.HoloMap({f"{t[i]:.3f}": frame(i) for i in range(nt)}, kdims="time")

hv.output(holomap="scrubber", fps=10)
anim

Working with real data#

Data#

Using the same example as in the maria-overview.ipynb notebook:

The simulation is described in Regional Multi-Week Convection-Permitting Simulations using the Model for Prediction Across Scales (MPAS-A) for use in the Caribbean, Mexico, and Central America Regions (Núñez Ocasio, Bacmeister, Xue and Morrison; NSF NCAR GDEX dataset d010077, doi:10.5065/PNB8-GD97)

%%time
grid_path = "/gdex/data/d010077/native_diag/MAAG.static.nc"
data_path = "/gdex/data/d010077/native_diag/diag.2017-09-20_06.00.00.nc"
ds_full = ux.open_dataset(grid_path, data_path)
ds = ds_full.isel(Time=0)  # just consider 1 snapshot here
ds

CPU times: user 1.36 s, sys: 57.8 s, total: 59.2 s
Wall time: 1min 19s

<xarray.UxDataset> Size: 10GB
Dimensions:                (n_face: 6087575, t_iso_levels: 27,
                            u_iso_levels: 27, n_node: 12177942, z_iso_levels: 27)
Coordinates:
  * t_iso_levels           (t_iso_levels) float32 108B 1e+04 1.25e+04 ... 1e+05
  * u_iso_levels           (u_iso_levels) float32 108B 1e+04 1.25e+04 ... 1e+05
  * z_iso_levels           (z_iso_levels) float32 108B 1e+04 1.25e+04 ... 1e+05
Dimensions without coordinates: n_face, n_node
Data variables: (12/82)
    cape                   (n_face) float32 24MB ...
    quantization_info      |S1 1B ...
    cin                    (n_face) float32 24MB ...
    dewpoint_200hPa        (n_face) float32 24MB ...
    dewpoint_250hPa        (n_face) float32 24MB ...
    dewpoint_500hPa        (n_face) float32 24MB ...
    ...                     ...
    w_700hPa               (n_face) float32 24MB ...
    w_850hPa               (n_face) float32 24MB ...
    w_925hPa               (n_face) float32 24MB ...
    w_isobaric             (n_face, u_iso_levels) float32 657MB ...
    xtime                  |S64 64B ...
    z_isobaric             (n_face, z_iso_levels) float32 657MB ...
Attributes: (12/122)
    model_name:                             mpas
    core_name:                              atmosphere
    source:                                 MPAS
    Conventions:                            MPAS
    git_version:                            v8.0.1-17-g32d1b6cfa
    on_a_sphere:                            YES
    ...                                     ...
    config_oml_hml0:                        30.0
    config_oml_gamma:                       0.14
    config_oml_relaxation_time:             864000.0
    config_sounding_interval:               none
    file_id:                                p5uyap8moz
    NCO:                                    netCDF Operators version 5.3.1 (H...
# check that individual data_vars are not too big to load into memory:
ds['uzonal_200hPa'].nbytes / 1e9

0.0243503

Load relevant data variables into memory.

We’ll compare the curl of the winds to the vorticity, and their divergence to the outgoing longwave radiation (OLR) at top of atmosphere.

Using two different pressures to check different regions: 850 hPa for “lower atmosphere”, and 200 hPa for “upper atmosphere”.

%%time
u200 = ds['uzonal_200hPa'].compute()
v200 = ds['umeridional_200hPa'].compute()
vort200 = ds['vorticity_200hPa'].compute().topological_mean("face")
# (the topological_mean("face") interpolates node-centered data to be face-centered instead)
u850 = ds['uzonal_850hPa'].compute()
v850 = ds['umeridional_850hPa'].compute()
vort850 = ds['vorticity_850hPa'].compute().topological_mean("face")
olr = ds['olrtoa'].compute()

CPU times: user 3.57 s, sys: 243 ms, total: 3.82 s
Wall time: 4.75 s

Compare curl and vorticity#

The data contains zonal and meridional components of winds, and vorticity.

Let’s do a sanity check: uxarray’s curl of the winds should match the data’s voriticity!

%%time
# (focusing just on the storm region, for demo purposes)
storm_box = dict(lon_bounds=(-71, -59), lat_bounds=(13, 22))
u200_storm = u200.subset.bounding_box(**storm_box)
v200_storm = v200.subset.bounding_box(**storm_box)
vort200_storm = vort200.subset.bounding_box(**storm_box)
u850_storm = u850.subset.bounding_box(**storm_box)
v850_storm = v850.subset.bounding_box(**storm_box)
vort850_storm = vort850.subset.bounding_box(**storm_box)

CPU times: user 49.1 s, sys: 662 ms, total: 49.8 s
Wall time: 36.2 s

%%time
curl200_storm = u200_storm.curl(v200_storm)
curl850_storm = u850_storm.curl(v850_storm)
vlim200 = max(np.abs(curl200_storm).quantile(0.99), np.abs(vort200_storm).quantile(0.99)).item()
vlim850 = max(np.abs(curl200_storm).quantile(0.99), np.abs(vort200_storm).quantile(0.99)).item()
vlim = max(vlim200, vlim850)
kws = dict(clim=(-vlim, vlim), data_aspect=1, cmap="RdBu_r")

vort200_plot = vort200_storm.plot(**kws, colorbar=False, title="vorticity (200 hPa)")
curl200_plot = curl200_storm.plot(**kws, title="curl (200 hPa)")
vort850_plot = vort850_storm.plot(**kws, colorbar=False, title="vorticity (850 hPa)")
curl850_plot = curl850_storm.plot(**kws, title="curl (850 hPa)")
layout = vort200_plot + curl200_plot + vort850_plot + curl850_plot
layout.cols(2)

CPU times: user 11.4 s, sys: 231 ms, total: 11.7 s
Wall time: 13.7 s

By eye (pun intended?), it looks like a very close match!

Compare divergence and top-of-atmosphere outgoing longwave radiation (OLR)#

Horizontal divergence often corresponds to updrafts or downdrafts, which can lead to clouds.

One might expect more clouds across regions of more divergence!

%%time
# (focusing just on the CONUS region, for demo purposes)
conus_box = dict(lon_bounds=(-125, -66), lat_bounds=(24, 50))
u200_conus = u200.subset.bounding_box(**conus_box)
v200_conus = v200.subset.bounding_box(**conus_box)
u850_conus = u850.subset.bounding_box(**conus_box)
v850_conus = v850.subset.bounding_box(**conus_box)
olr_conus = olr.subset.bounding_box(**conus_box)

CPU times: user 57.4 s, sys: 1.76 s, total: 59.2 s
Wall time: 1min 7s

%%time
div200_conus = u200_conus.divergence(v200_conus)
div850_conus = u850_conus.divergence(v850_conus)
dlim200 = np.abs(div200_conus).quantile(0.99).item()
dlim850 = np.abs(div850_conus).quantile(0.99).item()

CPU times: user 4.76 s, sys: 68.5 ms, total: 4.83 s
Wall time: 2.22 s

%%time
kws_div = dict(data_aspect=1, cmap="RdBu_r")

div200_plot = div200_conus.plot(**kws_div, clim=(-dlim200, dlim200), title="divergence at 200 hPa")
olr_plot = olr_conus.plot(data_aspect=1, cmap="Greys", title="top-of-atmosphere OLR")
# (use dlim200 for all limits to use same scale in both plots, or swap to dlim850 if you prefer!)
div850_plot = div850_conus.plot(**kws_div, clim=(-dlim200, dlim200), title="divergence at 850 hPa")
(div200_plot + olr_plot + div850_plot).cols(1)

CPU times: user 33.5 s, sys: 2.32 s, total: 35.8 s
Wall time: 43.7 s

The correlation isn’t perfect (divergence isn’t the only thing related to cloud-cover and OLR), but it’s clearly present, at least in the 200 hPa divergence example.

Prompt: why is some of the upper-atmosphere divergence correlated with OLR, but some is not?

Prompt: why is there so much lower-atmosphere divergence in the western US but not as much in the eastern US?