Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Repalce xarray wrapper with new version that use apply_ufunc along ti…
…me dimension (#89) * Added completely new xarray wrapper * Update whatsnew * Example processing notebook is working again
- Loading branch information
Showing
8 changed files
with
598 additions
and
123 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,78 +1,77 @@ | ||
| from functools import wraps | ||
| from collections import OrderedDict | ||
| import inspect | ||
| import numpy as np | ||
| import xarray as xr | ||
|
|
||
|
|
||
| def xarray_loop_vars_over_dim(vars_to_loop, loop_dim): | ||
| def xarray_apply_along_time_dim(): | ||
| """ | ||
| A decorator to feed CML data with mutiple channels as xarray.DataArrays to CML processing functions | ||
| A decorator to apply CML processing function along the time dimension of DataArrays | ||
| Parameters | ||
| ---------- | ||
| vars_to_loop: list of strings | ||
| List of the names of the variables used as kwargs in the decorated function | ||
| which should have a dimension `loop_dim` for which the decorated function is | ||
| then applied individually to each item when looping over `loop_dim`. | ||
| loop_dim: basestring | ||
| Name of the dimension which all variables in `vars_to_loop` must have in common | ||
| and which will be looped over to apply the decorated function. | ||
| Examples | ||
| -------- | ||
| Here is an example for how this decorator is used for the WAA Schleiss function:: | ||
| @xarray_loop_vars_over_dim(vars_to_loop=["rsl", "baseline", "wet"], loop_dim="channel_id") | ||
| def waa_schleiss_2013(rsl, baseline, wet, waa_max, delta_t, tau): | ||
| # function body... | ||
| Here, `delta_t` and `tau` are not CML data xarray.DataArrays and hence do not | ||
| have to be looped over. | ||
| This will work if the decorated function takes 1D numpy arrays, which hold the | ||
| CML time series data, as arguments. Additional argument are also handled. | ||
| """ | ||
|
|
||
| def decorator(func): | ||
| @wraps(func) | ||
| def inner(*args, **kwargs): | ||
| # TODO: Maybe check if all args or kwargs are the same type | ||
| # The case with numpy array as arg | ||
| if len(args) > 0 and isinstance(args[0], np.ndarray): | ||
| return func(*args, **kwargs) | ||
| # The case with numpy array as kwarg | ||
| if len(kwargs.keys()) > 0 and isinstance( | ||
| kwargs[vars_to_loop[0]], np.ndarray | ||
| ): | ||
| return func(*args, **kwargs) | ||
| # The dummy case where nothing is passed. This is just to get the | ||
| # functions error message here instead of continuing to the loop below | ||
| if len(args) == 0 and len(kwargs) == 0: | ||
| return func(*args, **kwargs) | ||
| new_args_dict = _get_new_args_dict(func=func, args=args, kwargs=kwargs) | ||
|
|
||
| loop_dim_id_list = list( | ||
| np.atleast_1d(kwargs[vars_to_loop[0]][loop_dim].values) | ||
| ) | ||
| if len(loop_dim_id_list) > 1: | ||
| kwargs_vars_to_loop = {var: kwargs.pop(var) for var in vars_to_loop} | ||
| data_list = [] | ||
| for loop_dim_id in loop_dim_id_list: | ||
| for var in vars_to_loop: | ||
| kwargs[var] = kwargs_vars_to_loop[var].sel( | ||
| {loop_dim: loop_dim_id} | ||
| ).values | ||
| data_list.append(func(**kwargs)) | ||
| return xr.DataArray( | ||
| data=np.stack(data_list), | ||
| dims=(loop_dim, "time"), | ||
| coords={ | ||
| loop_dim: kwargs_vars_to_loop[vars_to_loop[0]][loop_dim].values, | ||
| "time": kwargs_vars_to_loop[vars_to_loop[0]].time, | ||
| }, | ||
| ) | ||
| # Check if any arg has a time dim. Also build the input_core_dims | ||
| # list which will be required below and in our case will contain either | ||
| # ['time'] or [] because we do only care about applying along the time | ||
| # dim if the arg has one. Note that input_core_dims has to have the same | ||
| # number of entries as func has args. | ||
| input_core_dims = [] | ||
| found_time_dim = False | ||
| for arg in new_args_dict.values(): | ||
| try: | ||
| if "time" in arg.dims: | ||
| found_time_dim = True | ||
| input_core_dims.append(["time"]) | ||
| else: | ||
| input_core_dims.append([]) | ||
| except AttributeError: | ||
| input_core_dims.append([]) | ||
|
|
||
| # If now arg has a `time` dim, then we just call the function as it would | ||
| # be called without the wrapper. | ||
| if not found_time_dim: | ||
| return func(*args, **kwargs) | ||
| else: | ||
| return xr.DataArray( | ||
| data=func(**kwargs), | ||
| dims=("time"), | ||
| coords={"time": kwargs[vars_to_loop[0]].time}, | ||
| return xr.apply_ufunc( | ||
| func, | ||
| *list(new_args_dict.values()), | ||
| input_core_dims=input_core_dims, | ||
| output_core_dims=[["time"]], | ||
| vectorize=True, | ||
| ) | ||
|
|
||
| return inner | ||
|
|
||
| return decorator | ||
|
|
||
|
|
||
| def _get_new_args_dict(func, args, kwargs): | ||
| """Build one dict from args, kwargs and function default args | ||
| The function signature is used to build one joint dict from args and kwargs and | ||
| additional from the default arguments found in the function signature. The order | ||
| of the args in this dict is the order of the args in the function signature and | ||
| hence the list of args can be used in cases where we can only supply *args, but | ||
| we have to work with a mixture of args, kwargs and default args as in | ||
| xarray.apply_ufunc in the xarray wrapper. | ||
| """ | ||
| new_args_dict = OrderedDict() | ||
| for i, (arg, parameter) in enumerate(inspect.signature(func).parameters.items()): | ||
| if i < len(args): | ||
| new_args_dict[arg] = args[i] | ||
| elif arg in kwargs.keys(): | ||
| new_args_dict[arg] = kwargs[arg] | ||
| else: | ||
| new_args_dict[arg] = parameter.default | ||
|
|
||
| return new_args_dict |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| from unittest import TestCase | ||
| from collections import OrderedDict | ||
| import pycomlink.processing.xarray_wrapper | ||
|
|
||
|
|
||
| def test_get_new_args_dict(): | ||
| def foo_func(a, b, c=None, d=1): | ||
| return None | ||
|
|
||
| expected = OrderedDict( | ||
| [ | ||
| ("a", 123), | ||
| ("b", "foo"), | ||
| ("c", "bar"), | ||
| ("d", 1), | ||
| ] | ||
| ) | ||
| result = pycomlink.processing.xarray_wrapper._get_new_args_dict( | ||
| func=foo_func, args=[123, "foo", "bar", 1], kwargs={} | ||
| ) | ||
| TestCase().assertDictEqual(expected, result) | ||
|
|
||
| result = pycomlink.processing.xarray_wrapper._get_new_args_dict( | ||
| func=foo_func, | ||
| args=[ | ||
| 123, | ||
| ], | ||
| kwargs={"b": "foo", "c": "bar"}, | ||
| ) | ||
| TestCase().assertDictEqual(expected, result) |