[1]:

import earthkit.hydro as ekh
import numpy as np
import xarray as xr

Handling xarray and array inputs and outputs¶

For every operation, earthkit-hydro has two versions:

an xarray version, called at ekh.submodule.operation
an array version, called at ekh.submodule.array.operation

This notebook demonstrates this for some simple operations.

[2]:

network = ekh.river_network.load("efas", "5", use_cache=False)

Cache disabled.

xarray versions¶

xarray is the high-level recommended approach for conducting operations on river networks.

For this tutorial, example xarray data is created, but in practice this can be loaded from any supported file with ekd.from_source(...).to_xarray() (or xr.open_dataset(...)).

[3]:

example_arr = np.random.rand(2, *network.shape)

lat = network.coords['lat']
lon = network.coords['lon']
step = np.arange(2)

example_da = xr.DataArray(
    example_arr,
    dims = ["step", "lat", "lon"],
    coords = {"step": step, "lat": lat, "lon": lon},
    name = "precip",
    attrs={"units": "m", "description": "Sample precipitation data"}
)
example_da

All operations can then be directly conducted on this xarray dataarray, and will return back an xarray dataarray.

For example, we can do an upstream.sum, or a catchments.sum.

[4]:

ekh.upstream.sum(network, example_da)

[5]:

ekh.catchments.sum(network, example_da, locations = {
        "gauge_1": (70.475, 28.32),
        "gauge_2": (42.225, 50.24)
    }
)

We can also pass a dataset, in which case we will be returned back a corresponding dataset.

[6]:

example_ds = xr.Dataset(
    data_vars={"var1": example_da, "var2": example_da+1}
)
example_ds

[7]:

ekh.upstream.sum(network, example_ds)

Note that the coordinates were automatically detected in the above example. In some cases, it will not be possible infer the spatial coordinates in the xarray object. See for example the following:

[8]:

example_da_uninferrable = xr.DataArray(
    example_arr,
    dims = ["step", "uninferrable_name_1", "uninferrable_name_2"],
    coords = {"step": step, "uninferrable_name_1": lat, "uninferrable_name_2": lon},
    name = "precip",
    attrs={"units": "m", "description": "Sample precipitation data"}
)

This will raise a ValueError: Could not autodetect xarray core dims. However, we can manually specify the input core dimensions for each of the xarray inputs via the argument input_core_dims.

[9]:

ekh.upstream.sum(network, example_da_uninferrable, input_core_dims=[["uninferrable_name_1", "uninferrable_name_2"]])

We need to specify this for each xarray argument. In the below example, we have two xarray arguments so we specify the input_core_dims for both, in the order defined by the function signature.

[10]:

ekh.upstream.sum(network, example_da, node_weights=example_da_uninferrable, input_core_dims=[["lat", "lon"], ["uninferrable_name_1", "uninferrable_name_2"]])

It is also possible to give purely array inputs, in which case an xarray dataarray is returned with the coordinate information of the river network. For example,

[11]:

ekh.upstream.sum(network, example_arr)

Note: Array inputs must follow the convention that the last dimensions are the river network related dimensions, in this case lat and lon. Therefore, for any vectorised operations, the extra dimensions must be leading.

array versions¶

A lower level API is also available for users interested in working directly with arrays. This API does not accept any xarray inputs, but works analagously to the xarray API. Taking the two above examples, the array forms are:

[12]:

ekh.upstream.array.sum(network, example_arr)

[12]:

array([[[0.15711655, 0.52936434, 0.56102303, ...,        nan,
                nan,        nan],
        [0.91444499, 1.23290502, 1.86772309, ...,        nan,
                nan,        nan],
        [0.38624433, 0.84895319, 0.16244652, ...,        nan,
                nan,        nan],
        ...,
        [       nan,        nan,        nan, ..., 0.75432537,
         0.71705345, 0.09432709],
        [       nan,        nan,        nan, ..., 0.47122936,
         2.46722595, 3.7863624 ],
        [       nan,        nan,        nan, ..., 1.28929405,
         0.00772117, 0.84704023]],

       [[0.57485927, 0.04379043, 0.1061605 , ...,        nan,
                nan,        nan],
        [0.59805427, 0.65665193, 0.96566858, ...,        nan,
                nan,        nan],
        [0.68846108, 1.54718075, 0.7100234 , ...,        nan,
                nan,        nan],
        ...,
        [       nan,        nan,        nan, ..., 0.03131001,
         0.65784284, 0.57760851],
        [       nan,        nan,        nan, ..., 0.49580478,
         3.73241039, 4.60202261],
        [       nan,        nan,        nan, ..., 1.65342032,
         0.71465062, 0.44616404]]])

[13]:

ekh.catchments.array.sum(network, example_arr, locations = {
        "gauge_1": (70.475, 28.32),
        "gauge_2": (42.225, 50.24)
    }
)

[13]:

array([[  6710.97380367, 272520.24025306],
       [  6716.52716166, 272441.37894161]])