Destaggering and vertically interpolating dask data using xgcm

In this tutorial, we will show you how to leverage the xWRF accessors and the xWRF-provided COMODO-compliant attributes in order to destagger the WRF output and interpolate it vertically using dask and xgcm.

Loading the data

First of all, we load the data and use the simple .xwrf.postprocess() API and dask-enable the dataset by passing open_dataset the chunkgs kwarg. In a real-world scenario, you might want to spawn a Cluster in order to speed up calculations.

import xwrf
ds = xwrf.tutorial.open_dataset("wrfout", chunks='auto').xwrf.postprocess()
ds

<xarray.Dataset> Size: 104MB
Dimensions:                    (y: 340, x: 270, Time: 1, z: 39, x_stag: 271,
                                y_stag: 341)
Coordinates: (12/13)
    XLAT                       (y, x) float32 367kB dask.array<chunksize=(340, 270), meta=np.ndarray>
    XLONG                      (y, x) float32 367kB dask.array<chunksize=(340, 270), meta=np.ndarray>
    XTIME                      (Time) datetime64[ns] 8B dask.array<chunksize=(1,), meta=np.ndarray>
    XLAT_U                     (y, x_stag) float32 369kB dask.array<chunksize=(340, 271), meta=np.ndarray>
    XLONG_U                    (y, x_stag) float32 369kB dask.array<chunksize=(340, 271), meta=np.ndarray>
    XLAT_V                     (y_stag, x) float32 368kB dask.array<chunksize=(341, 270), meta=np.ndarray>
    ...                         ...
  * z                          (z) float32 156B 0.9969 0.9899 ... 0.002948
  * Time                       (Time) datetime64[ns] 8B 2099-10-01
  * y                          (y) float64 3kB -3.341e+05 ... 2.717e+06
  * x_stag                     (x_stag) float64 2kB -4.733e+06 ... -2.303e+06
  * y_stag                     (y_stag) float64 3kB -3.386e+05 ... 2.721e+06
  * x                          (x) float64 2kB -4.728e+06 ... -2.307e+06
Data variables:
    Times                      (Time) |S19 19B dask.array<chunksize=(1,), meta=np.ndarray>
    U                          (Time, z, y, x_stag) float32 14MB dask.array<chunksize=(1, 39, 340, 271), meta=np.ndarray>
    V                          (Time, z, y_stag, x) float32 14MB dask.array<chunksize=(1, 39, 341, 270), meta=np.ndarray>
    SINALPHA                   (Time, y, x) float32 367kB dask.array<chunksize=(1, 340, 270), meta=np.ndarray>
    COSALPHA                   (Time, y, x) float32 367kB dask.array<chunksize=(1, 340, 270), meta=np.ndarray>
    QVAPOR                     (Time, z, y, x) float32 14MB dask.array<chunksize=(1, 39, 340, 270), meta=np.ndarray>
    PSN                        (Time, y, x) float32 367kB dask.array<chunksize=(1, 340, 270), meta=np.ndarray>
    air_potential_temperature  (Time, z, y, x) float32 14MB dask.array<chunksize=(1, 39, 340, 270), meta=np.ndarray>
    air_pressure               (Time, z, y, x) float32 14MB dask.array<chunksize=(1, 39, 340, 270), meta=np.ndarray>
    wind_east                  (Time, z, y, x) float32 14MB dask.array<chunksize=(1, 39, 340, 270), meta=np.ndarray>
    wind_north                 (Time, z, y, x) float32 14MB dask.array<chunksize=(1, 39, 340, 270), meta=np.ndarray>
    wrf_projection             object 8B +proj=lcc +x_0=0 +y_0=0 +a=6370000 +...
Attributes: (12/149)
    TITLE:                            OUTPUT FROM WRF V4.1.3 MODEL
    START_DATE:                      2099-08-01_00:00:00
    SIMULATION_START_DATE:           2099-08-01_00:00:00
    WEST-EAST_GRID_DIMENSION:        271
    SOUTH-NORTH_GRID_DIMENSION:      341
    BOTTOM-TOP_GRID_DIMENSION:       40
    ...                              ...
    ISLAKE:                          21
    ISICE:                           15
    ISURBAN:                         13
    ISOILWATER:                      14
    HYBRID_OPT:                      0
    ETAC:                            0.0

Destaggering

If we naively try to calculate the wind speed from the U and V components, we get an error due to them having different shapes.

from metpy.calc import wind_speed

wind_speed(ds.U, ds.V)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[2], line 3
      1 from metpy.calc import wind_speed
----> 3 wind_speed(ds.U, ds.V)

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/metpy/xarray.py:1328, in preprocess_and_wrap.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
   1325     _mutate_arguments(bound_args, units.Quantity, lambda arg, _: arg.m)
   1327 # Evaluate inner calculation
-> 1328 result = func(*bound_args.args, **bound_args.kwargs)
   1330 # Wrap output based on match and match_unit
   1331 if match is None:

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/metpy/units.py:337, in check_units.<locals>.dec.<locals>.wrapper(*args, **kwargs)
    334 @functools.wraps(func)
    335 def wrapper(*args, **kwargs):
    336     _check_units_inner_helper(func, sig, defaults, dims, *args, **kwargs)
--> 337     return func(*args, **kwargs)

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/metpy/calc/basic.py:63, in wind_speed(u, v)
     33 @exporter.export
     34 @preprocess_and_wrap(wrap_like='u')
     35 @check_units('[speed]', '[speed]')
     36 def wind_speed(u, v):
     37     r"""Compute the wind speed from u and v-components.
     38 
     39     Parameters
   (...)
     61 
     62     """
---> 63     return np.hypot(u, v)

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/pint/facets/numpy/quantity.py:72, in NumpyQuantity.__array_ufunc__(self, ufunc, method, *inputs, **kwargs)
     65 # Replicate types from __array_function__
     66 types = {
     67     type(arg)
     68     for arg in list(inputs) + list(kwargs.values())
     69     if hasattr(arg, "__array_ufunc__")
     70 }
---> 72 return numpy_wrap("ufunc", ufunc, inputs, kwargs, types)

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/pint/facets/numpy/numpy_func.py:1071, in numpy_wrap(func_type, func, args, kwargs, types)
   1069 if name not in handled or any(is_upcast_type(t) for t in types):
   1070     return NotImplemented
-> 1071 return handled[name](*args, **kwargs)

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/pint/facets/numpy/numpy_func.py:322, in implement_func.<locals>.implementation(*args, **kwargs)
    317     stripped_args, stripped_kwargs = convert_to_consistent_units(
    318         *args, pre_calc_units=pre_calc_units, **kwargs
    319     )
    321 # Determine result through plain numpy function on stripped arguments
--> 322 result_magnitude = func(*stripped_args, **stripped_kwargs)
    324 if output_unit is None:
    325     # Short circuit and return magnitude alone
    326     return result_magnitude

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/dask/array/core.py:1605, in Array.__array_ufunc__(self, numpy_ufunc, method, *inputs, **kwargs)
   1603         return da_ufunc(*inputs, **kwargs)
   1604     else:
-> 1605         return elemwise(numpy_ufunc, *inputs, **kwargs)
   1606 elif method == "outer":
   1607     from dask.array import ufunc

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/dask/array/core.py:4957, in elemwise(op, out, where, dtype, name, *args, **kwargs)
   4953     shapes.append(out.shape)
   4955 shapes = [s if isinstance(s, Iterable) else () for s in shapes]
   4956 out_ndim = len(
-> 4957     broadcast_shapes(*shapes)
   4958 )  # Raises ValueError if dimensions mismatch
   4959 expr_inds = tuple(range(out_ndim))[::-1]
   4961 if dtype is not None:

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.10/site-packages/dask/array/core.py:4885, in broadcast_shapes(*shapes)
   4883         dim = 0 if 0 in sizes else np.max(sizes).item()
   4884     if any(i not in [-1, 0, 1, dim] and not np.isnan(i) for i in sizes):
-> 4885         raise ValueError(
   4886             "operands could not be broadcast together with "
   4887             "shapes {}".format(" ".join(map(str, shapes)))
   4888         )
   4889     out.append(dim)
   4890 return tuple(reversed(out))

ValueError: operands could not be broadcast together with shapes (1, 39, 340, 271) (1, 39, 341, 270)

Upon investigating the wind components, we can see that they are defined on the WRF-internal Arakawa-C grid, which causes the shapes to differ.

ds.U.sizes, ds.V.sizes

(Frozen({'Time': 1, 'z': 39, 'y': 340, 'x_stag': 271}),
 Frozen({'Time': 1, 'z': 39, 'y_stag': 341, 'x': 270}))

Destaggering is done in no time at all using the handy .xwrf accessor. We can now decide whether to destagger the whole Dataset…

destaggered = ds.xwrf.destagger().metpy.quantify()
destaggered['wind_speed'] = wind_speed(destaggered.U, destaggered.V)
destaggered.wind_speed

<xarray.DataArray 'wind_speed' (Time: 1, z: 39, y: 340, x: 270)> Size: 14MB
<Quantity(dask.array<hypot, shape=(1, 39, 340, 270), dtype=float32, chunksize=(1, 39, 340, 270), chunktype=numpy.ndarray>, 'meter / second')>
Coordinates:
    XTIME    (Time) datetime64[ns] 8B dask.array<chunksize=(1,), meta=np.ndarray>
  * Time     (Time) datetime64[ns] 8B 2099-10-01
    XLAT     (y, x) float32 367kB dask.array<chunksize=(340, 270), meta=np.ndarray>
    XLONG    (y, x) float32 367kB dask.array<chunksize=(340, 270), meta=np.ndarray>
  * z        (z) float32 156B 0.9969 0.9899 0.981 ... 0.0161 0.009174 0.002948
  * y        (y) float64 3kB -3.341e+05 -3.251e+05 ... 2.708e+06 2.717e+06
  * x        (x) float64 2kB -4.728e+06 -4.719e+06 ... -2.316e+06 -2.307e+06

… or whether we just want to destagger the two individual DataArrays.

ds = ds.metpy.quantify()
wind_speed(ds.U.xwrf.destagger(), ds.V.xwrf.destagger())

<xarray.DataArray 'hypot-0d45d0e070dec755a1195c5ff1fc2ee7' (Time: 1, z: 39,
                                                            y: 340, x: 270)> Size: 14MB
<Quantity(dask.array<hypot, shape=(1, 39, 340, 270), dtype=float32, chunksize=(1, 39, 340, 270), chunktype=numpy.ndarray>, 'meter / second')>
Coordinates:
    XTIME    (Time) datetime64[ns] 8B dask.array<chunksize=(1,), meta=np.ndarray>
  * Time     (Time) datetime64[ns] 8B 2099-10-01
    XLAT     (y, x) float32 367kB dask.array<chunksize=(340, 270), meta=np.ndarray>
    XLONG    (y, x) float32 367kB dask.array<chunksize=(340, 270), meta=np.ndarray>
  * z        (z) float32 156B 0.9969 0.9899 0.981 ... 0.0161 0.009174 0.002948
  * y        (y) float64 3kB -3.341e+05 -3.251e+05 ... 2.708e+06 2.717e+06
  * x        (x) float64 2kB -4.728e+06 -4.719e+06 ... -2.316e+06 -2.307e+06

Vertical interpolation using xgcm

We have now calculated the wind speed for the whole model domain. However, the z-layers are still in the native WRF sigma coordinate, which is of no practical use to us. So, in order to be able to analyze this data properly, we have to interpolate it onto pressure levels.

But, since xWRF prepared the dataset with the appropriate COMODO (and units) attributes, we can simply use xgcm with its Grid.transform function to solve this problem! However, since it doesn’t understand units yet, we have to work around it a bit:

import xgcm
import numpy as np
import pint_xarray

target_levels = np.array([250.]) # in hPa
air_pressure = destaggered.air_pressure.pint.to('hPa').metpy.dequantify()

grid = xgcm.Grid(destaggered, periodic=False)
_wind_speed = grid.transform(destaggered.wind_speed.metpy.dequantify(), 'Z', target_levels, target_data=air_pressure, method='log')
_wind_speed = _wind_speed.compute()

Finally, we can plot the result using hvplot.

import hvplot.xarray

_wind_speed.hvplot.quadmesh(
    x='XLONG',
    y='XLAT',
    title='Wind speed at 250 hPa',
    geo=True,
    project=True,
    alpha=0.9,
    cmap='inferno',
    clim=(_wind_speed.min().item(), _wind_speed.max().item()),
    clabel='wind speed [m/s]',
    tiles='OSM',
    dynamic=False
)