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
  * x                          (x) float64 2kB -4.728e+06 ... -2.307e+06
  * y_stag                     (y_stag) float64 3kB -3.386e+05 ... 2.721e+06
  * y                          (y) float64 3kB -3.341e+05 ... 2.717e+06
  * x_stag                     (x_stag) float64 2kB -4.733e+06 ... -2.303e+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/latest/lib/python3.10/site-packages/metpy/xarray.py:1330, in preprocess_and_wrap.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
   1327     _mutate_arguments(bound_args, units.Quantity, lambda arg, _: arg.m)
   1329 # Evaluate inner calculation
-> 1330 result = func(*bound_args.args, **bound_args.kwargs)
   1332 # Wrap output based on match and match_unit
   1333 if match is None:

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

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

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/latest/lib/python3.10/site-packages/dask/array/core.py:4950, in elemwise(op, out, where, dtype, name, *args, **kwargs)
   4946     shapes.append(out.shape)
   4948 shapes = [s if isinstance(s, Iterable) else () for s in shapes]
   4949 out_ndim = len(
-> 4950     broadcast_shapes(*shapes)
   4951 )  # Raises ValueError if dimensions mismatch
   4952 expr_inds = tuple(range(out_ndim))[::-1]
   4954 if dtype is not None:

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/latest/lib/python3.10/site-packages/dask/array/core.py:4878, in broadcast_shapes(*shapes)
   4876         dim = 0 if 0 in sizes else np.max(sizes).item()
   4877     if any(i not in [-1, 0, 1, dim] and not np.isnan(i) for i in sizes):
-> 4878         raise ValueError(
   4879             "operands could not be broadcast together with "
   4880             "shapes {}".format(" ".join(map(str, shapes)))
   4881         )
   4882     out.append(dim)
   4883 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
  * x        (x) float64 2kB -4.728e+06 -4.719e+06 ... -2.316e+06 -2.307e+06
  * y        (y) float64 3kB -3.341e+05 -3.251e+05 ... 2.708e+06 2.717e+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-04e556d1a8addf8c653a11e26a9c3a45' (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()

/home/docs/checkouts/readthedocs.org/user_builds/xwrf/conda/latest/lib/python3.10/site-packages/xgcm/grid.py:989: FutureWarning: From version 0.8.0 the Axis computation methods will be removed, in favour of using the Grid computation methods instead. i.e. use `Grid.transform` instead of `Axis.transform`
  warnings.warn(

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
)