Guide and example on how to use nested grids with DataTrees#

This is an example on how to use PyDDA’s ability to handle nested grids using xarray DataTrees. In this example, we load radars with two pre-generated Cf/Radial grid. The fine grids are higher resolution grids that are contained within the coarser grid.

The DataTree structure that PyDDA follows is:

::

root
  |---nest_0/radar_1
  |---nest_0/radar_2
  |---nest_0/radar_n
  |---nest_1/radar_1
  |---nest_1/radar_2
  |---nest_1/radar_m

Each member of this tree is a DataTree itself. PyDDA will know if the DataTree contains data from a radar when the name of the node begins with radar_. The root node of each grid level, in this example, root and inner_nest will contain the keyword arguments that are inputs to :code:pydda.retrieval.get_dd_wind_field as attributes for the tree. PyDDA will use the attributes at each level as the arguments for the retrieval, allowing the user to vary the coefficients by grid level.

Using :code:pydda.retrieval.get_dd_wind_field_nested will allow PyDDA to perform the retrieval on the 0th grid first. It will then perform on the subsequent grid levels, using the previous nest as both the horizontal boundary conditions and initialization for the retrieval in the next nest. Finally, PyDDA will update the winds in the first grid by nearest- neighbor interpolation of the latter grid into the overlapping portion between the inner and outer grid level.

PyDDA will then return the retrieved wind fields as the “u”, “v”, and “w” DataArrays inside each of the root nodes for each level, in this case root and inner_nest.

## Do imports
import pydda
import matplotlib.pyplot as plt
import warnings
from xarray import DataTree

warnings.filterwarnings("ignore")

"""
We will load pregenerated grids for this case.
"""
test_coarse0 = pydda.io.read_grid(pydda.tests.get_sample_file("test_coarse0.nc"))
test_coarse1 = pydda.io.read_grid(pydda.tests.get_sample_file("test_coarse1.nc"))
test_fine0 = pydda.io.read_grid(pydda.tests.get_sample_file("test_fine0.nc"))
test_fine1 = pydda.io.read_grid(pydda.tests.get_sample_file("test_fine1.nc"))

"""
Initalize with a zero wind field. We have HRRR data already generated for this case inside
the example data files to provide a model constraint.
"""
test_coarse0 = pydda.initialization.make_constant_wind_field(
    test_coarse0, (0.0, 0.0, 0.0)
)

"""
Specify the retrieval parameters at each level
"""
kwargs_dict = dict(
    Cm=256.0,
    Co=1e-2,
    Cx=50.0,
    Cy=50.0,
    Cz=50.0,
    Cmod=1e-5,
    model_fields=["hrrr"],
    refl_field="DBZ",
    wind_tol=0.5,
    max_iterations=150,
    engine="scipy",
)

"""
Enforce equal times for each grid. This is required for the DataTree structure since time is an
inherited dimension.
"""
test_coarse1["time"] = test_coarse0["time"]
test_fine0["time"] = test_coarse0["time"]
test_fine1["time"] = test_coarse1["time"]
"""

Provide the overlying grid structure as specified above.
"""
tree_dict = {
    "/nest_0/radar_ktlx": test_coarse0,
    "/nest_0/radar_kict": test_coarse1,
    "/nest_1/radar_ktlx": test_fine0,
    "/nest_1/radar_kict": test_fine1,
}

tree = DataTree.from_dict(tree_dict)
tree["/nest_0/"].attrs = kwargs_dict
tree["/nest_1/"].attrs = kwargs_dict

"""
Perform the retrieval
"""

grid_tree = pydda.retrieval.get_dd_wind_field_nested(tree)

"""
Plot the coarse grid output and finer grid output
"""

fig, ax = plt.subplots(1, 2, figsize=(10, 5))
pydda.vis.plot_horiz_xsection_quiver(
    grid_tree["nest_0"],
    ax=ax[0],
    level=5,
    cmap="ChaseSpectral",
    vmin=-10,
    vmax=80,
    quiverkey_len=10.0,
    background_field="DBZ",
    bg_grid_no=1,
    w_vel_contours=[1, 2, 5, 10],
    quiver_spacing_x_km=50.0,
    quiver_spacing_y_km=50.0,
    quiverkey_loc="bottom_right",
)
pydda.vis.plot_horiz_xsection_quiver(
    grid_tree["nest_1"],
    ax=ax[1],
    level=5,
    cmap="ChaseSpectral",
    vmin=-10,
    vmax=80,
    quiverkey_len=10.0,
    background_field="DBZ",
    bg_grid_no=1,
    w_vel_contours=[1, 2, 5, 10],
    quiver_spacing_x_km=50.0,
    quiver_spacing_y_km=50.0,
    quiverkey_loc="bottom_right",
)

plt.show()

## You are using the Python ARM Radar Toolkit (Py-ART), an open source
## library for working with weather radar data. Py-ART is partly
## supported by the U.S. Department of Energy as part of the Atmospheric
## Radiation Measurement (ARM) Climate Research Facility, an Office of
## Science user facility.
##
## If you use this software to prepare a publication, please cite:
##
##     JJ Helmus and SM Collis, JORS 2016, doi: 10.5334/jors.119

Welcome to PyDDA 2.1.0
If you are using PyDDA in your publications, please cite:
Jackson et al. (2020) Journal of Open Research Science
Detecting Jax...
Jax/JaxOpt are not installed on your system, unable to use Jax engine.
Detecting TensorFlow...
TensorFlow detected. Checking for tensorflow-probability...
TensorFlow-probability detected. TensorFlow engine enabled!

False

Calculating weights for radars 0 and 1
Calculating weights for radars 1 and 0

Calculating weights for models...
Starting solver 
rmsVR = 16.21367982576569
Total points: 412470
The max of w_init is 0.0

Total number of model points: 2342081

Nfeval | Jvel    | Jmass   | Jsmooth |   Jbg   | Jvort   | Jmodel  | Jpoint  | Max w  
      0|4115.2916|   0.0000|   0.0000|   0.0000|   0.0000|11243.5276|   0.0000|   0.0000

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[1], line 68
tree["/nest_1/"].attrs = kwargs_dict
"""
Perform the retrieval
"""
---> 68 grid_tree = pydda.retrieval.get_dd_wind_field_nested(tree)
"""
Plot the coarse grid output and finer grid output
"""
fig, ax = plt.subplots(1, 2, figsize=(10, 5))

File ~/work/PyDDA/PyDDA/pydda/retrieval/nesting.py:56, in get_dd_wind_field_nested(grid_tree, **kwargs)
elif len(grid_list) > 0:
   my_kwargs = tree_attrs
---> 56     output_grids, output_parameters = get_dd_wind_field(grid_list, **my_kwargs)
   output_parameters = output_parameters.__dict__
   if in_parent is True:

File ~/work/PyDDA/PyDDA/pydda/retrieval/wind_retrieve.py:1474, in get_dd_wind_field(Grids, u_init, v_init, w_init, engine, **kwargs)
   w_init = new_grids[0]["w"].values.squeeze()
if (
   engine.lower() == "scipy"
   or engine.lower() == "jax"
   or engine.lower() == "auglag"
):
-> 1474     return _get_dd_wind_field_scipy(
       new_grids, u_init, v_init, w_init, engine, **kwargs
   )
elif engine.lower() == "tensorflow":
   return _get_dd_wind_field_tensorflow(
       new_grids, u_init, v_init, w_init, **kwargs
   )

File ~/work/PyDDA/PyDDA/pydda/retrieval/wind_retrieve.py:602, in _get_dd_wind_field_scipy(Grids, u_init, v_init, w_init, engine, points, vel_name, refl_field, u_back, v_back, z_back, frz, Co, Cm, Cx, Cy, Cz, Cb, Cv, Cmod, Cpoint, cvtol, gtol, Jveltol, Ut, Vt, low_pass_filter, mask_outside_opt, weights_obs, weights_model, weights_bg, max_iterations, mask_w_outside_opt, filter_window, filter_order, min_bca, max_bca, upper_bc, model_fields, output_cost_functions, roi, wind_tol, tolerance, const_boundary_cond, max_wind_mag)
parameters.print_out = False
if engine.lower() == "scipy":
--> 602     winds = fmin_l_bfgs_b(
       J_function,
       winds,
       args=(parameters,),
       maxiter=max_iterations,
       pgtol=tolerance,
       bounds=bounds,
       fprime=grad_J,
       disp=0,
       iprint=-1,
       callback=_vert_velocity_callback,
   )
else:
   def loss_and_gradient(x):

File /usr/share/miniconda3/envs/pydda-docs/lib/python3.10/site-packages/scipy/optimize/_lbfgsb_py.py:237, in fmin_l_bfgs_b(func, x0, fprime, args, approx_grad, bounds, m, factr, pgtol, epsilon, iprint, maxfun, maxiter, disp, callback, maxls)
callback = _wrap_callback(callback)
opts = {'disp': disp,
       'iprint': iprint,
       'maxcor': m,
   (...)
       'callback': callback,
       'maxls': maxls}
--> 237 res = _minimize_lbfgsb(fun, x0, args=args, jac=jac, bounds=bounds,
                      **opts)
d = {'grad': res['jac'],
    'task': res['message'],
    'funcalls': res['nfev'],
    'nit': res['nit'],
    'warnflag': res['status']}
f = res['fun']

File /usr/share/miniconda3/envs/pydda-docs/lib/python3.10/site-packages/scipy/optimize/_lbfgsb_py.py:369, in _minimize_lbfgsb(fun, x0, args, jac, bounds, disp, maxcor, ftol, gtol, eps, maxfun, maxiter, iprint, callback, maxls, finite_diff_rel_step, **unknown_options)
   low_bnd[i] = l
   l = 1
--> 369 if not np.isinf(u):
   upper_bnd[i] = u
   u = 1

KeyboardInterrupt: 