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/latest/lib/python3.9/site-packages/metpy/xarray.py:1329, in preprocess_and_wrap.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
1326 _mutate_arguments(bound_args, units.Quantity, lambda arg, _: arg.m)
1328 # Evaluate inner calculation
-> 1329 result = func(*bound_args.args, **bound_args.kwargs)
1331 # Wrap output based on match and match_unit
1332 if match is None:
File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/latest/lib/python3.9/site-packages/metpy/units.py:332, in check_units.<locals>.dec.<locals>.wrapper(*args, **kwargs)
329 @functools.wraps(func)
330 def wrapper(*args, **kwargs):
331 _check_units_inner_helper(func, sig, defaults, dims, *args, **kwargs)
--> 332 return func(*args, **kwargs)
File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/latest/lib/python3.9/site-packages/metpy/calc/basic.py:56, 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
(...)
54
55 """
---> 56 return np.hypot(u, v)
File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/latest/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/latest/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/latest/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/latest/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/latest/lib/python3.9/site-packages/dask/array/core.py:4792, in elemwise(op, out, where, dtype, name, *args, **kwargs)
4788 shapes.append(out.shape)
4790 shapes = [s if isinstance(s, Iterable) else () for s in shapes]
4791 out_ndim = len(
-> 4792 broadcast_shapes(*shapes)
4793 ) # Raises ValueError if dimensions mismatch
4794 expr_inds = tuple(range(out_ndim))[::-1]
4796 if dtype is not None:
File ~/checkouts/readthedocs.org/user_builds/xwrf/conda/latest/lib/python3.9/site-packages/dask/array/core.py:4720, in broadcast_shapes(*shapes)
4718 dim = 0 if 0 in sizes else np.max(sizes)
4719 if any(i not in [-1, 0, 1, dim] and not np.isnan(i) for i in sizes):
-> 4720 raise ValueError(
4721 "operands could not be broadcast together with "
4722 "shapes {}".format(" ".join(map(str, shapes)))
4723 )
4724 out.append(dim)
4725 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-2d38941a3edc0979aac9af862e720455' (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.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
)