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_stag                     (x_stag) float64 2kB -4.733e+06 ... -2.303e+06
  * y                          (y) float64 3kB -3.341e+05 ... 2.717e+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.9/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/stable/lib/python3.9/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/stable/lib/python3.9/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.9/site-packages/pint/facets/numpy/quantity.py:73, in NumpyQuantity.__array_ufunc__(self, ufunc, method, *inputs, **kwargs)
     66 # Replicate types from __array_function__
     67 types = {
     68     type(arg)
     69     for arg in list(inputs) + list(kwargs.values())
     70     if hasattr(arg, "__array_ufunc__")
     71 }
---> 73 return numpy_wrap("ufunc", ufunc, inputs, kwargs, types)

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

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.9/site-packages/pint/facets/numpy/numpy_func.py:307, in implement_func.<locals>.implementation(*args, **kwargs)
    302     stripped_args, stripped_kwargs = convert_to_consistent_units(
    303         *args, pre_calc_units=pre_calc_units, **kwargs
    304     )
    306 # Determine result through plain numpy function on stripped arguments
--> 307 result_magnitude = func(*stripped_args, **stripped_kwargs)
    309 if output_unit is None:
    310     # Short circuit and return magnitude alone
    311     return result_magnitude

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.9/site-packages/dask/array/core.py:1594, in Array.__array_ufunc__(self, numpy_ufunc, method, *inputs, **kwargs)
   1592         return da_ufunc(*inputs, **kwargs)
   1593     else:
-> 1594         return elemwise(numpy_ufunc, *inputs, **kwargs)
   1595 elif method == "outer":
   1596     from dask.array import ufunc

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.9/site-packages/dask/array/core.py:4795, in elemwise(op, out, where, dtype, name, *args, **kwargs)
   4791     shapes.append(out.shape)
   4793 shapes = [s if isinstance(s, Iterable) else () for s in shapes]
   4794 out_ndim = len(
-> 4795     broadcast_shapes(*shapes)
   4796 )  # Raises ValueError if dimensions mismatch
   4797 expr_inds = tuple(range(out_ndim))[::-1]
   4799 if dtype is not None:

File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/stable/lib/python3.9/site-packages/dask/array/core.py:4723, in broadcast_shapes(*shapes)
   4721         dim = 0 if 0 in sizes else np.max(sizes)
   4722     if any(i not in [-1, 0, 1, dim] and not np.isnan(i) for i in sizes):
-> 4723         raise ValueError(
   4724             "operands could not be broadcast together with "
   4725             "shapes {}".format(" ".join(map(str, shapes)))
   4726         )
   4727     out.append(dim)
   4728 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-b399d2896e0546bebb6d8e7f14d37fb3' (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

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/stable/lib/python3.9/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
)