map_over_sweeps - Githubissues

xradar version: latest
Python version: any
Operating System: any

Description

There is ongoing discussion how to implement a map_over_sweeps functionality, which applies a function which works on sweep datasets to a DataTree's sweep nodes.

@syedhamidali did a great job in implementing such functionality for the old datatree.datatree in #203 based on accessors. Since DataTree will be a first class xarray citizen in the next xarray version (have a look for the updated xradar in feature branch https://github.com/openradar/xradar/tree/xarray-nightly) we should think about using the additional functionality (xarray has an map_over_datasets-decorator in combination with a DataTree-method of the same name.

During the discussion in #203 (a must read :grinning:) I've outlined a possible implementation which I paste here for verbosity:


# define the xradar decorator to just work on sweep nodes 
# and let other nodes out unchanged
# thats quick'ndirty, so the code might not be up to standards
import functools
import xarray as xr

def map_over_sweeps(func):

    @functools.wraps(func)
    def _func(*args, **kwargs):
        # we do not have access to any names or other DataTree specifics
        # so we need to rely on the layout of the given dataset
        # maybe we can add something like `is_sweep` to check for that?
        if "range" in args[0].dims:
            return func(*args, **kwargs)
        else:
            return args

    # we need another wrap here for the xarray decorator
    @xr.map_over_datasets
    def _map_over_sweeps(*args, **kwargs):
        return _func(*args, **kwargs)

    return _map_over_sweeps

# defined inside XarrayAccessor
def map_over_sweeps(self, func, *args, **kwargs):
    """    """
    # add our decorator for input function
    @map_over_sweeps
    def _func(*args, **kwargs):
        return func(*args, **kwargs)

    return self.xarray_obj.pipe(_func, *args, **kwargs)

The user would then have two options, one via accessor and one via decorator. There are surely use cases for both. The good thing is, that both ways are implemented solely by our and the xarray decorator and we only use xarray machinery (beside our little "range"-hack).

import xradar as xd
from open_radar_data import DATASETS

# Fetch the sample radar file
filename = DATASETS.fetch("sample_sgp_data.nc")

# Open the radar file into a DataTree object
dtree = xd.io.open_cfradial1_datatree(filename)

# without decorator
def calculate_rain_rate(ds, ref_field='DBZH'):
    def _rain_rate(dbz, a=200.0, b=1.6):
        Z = 10.0 ** (dbz / 10.0)
        return (Z / a) ** (1.0 / b)

    ds = ds.assign(RAIN_RATE=_rain_rate(ds[ref_field]))
    ds.RAIN_RATE.attrs = {'units': 'mm/h', 'long_name': 'Rain Rate'}
    return ds

# with decorator
@map_over_sweeps
def calculate_rain_rate2(ds, ref_field='DBZH'):
    def _rain_rate(dbz, a=200.0, b=1.6):
        Z = 10.0 ** (dbz / 10.0)
        return (Z / a) ** (1.0 / b)

    ds = ds.assign(RAIN_RATE=_rain_rate(ds[ref_field]))
    ds.RAIN_RATE.attrs = {'units': 'mm/h', 'long_name': 'Rain Rate'}
    return ds

# invocation via accessor
dtree2 = dtree.xradar.map_over_sweeps(calculate_rain_rate, ref_field='corrected_reflectivity_horizontal')
display(dtree2["sweep_0"])

# invocation via decorator + pipe
dtree3 = dtree.pipe(calculate_rain_rate2, ref_field='corrected_reflectivity_horizontal')
display(dtree3["sweep_0"])

The nice thing about that is, that users can re-use their functions which works on Datasets (eg sweeps) without any additional effort (beside adding the decorator).

Even in our codebase we can switch to this approach for our accessor based DataTree functionality.

@syedhamidali Are you willing to further lead the effort in implementing that? I'd happily assist you in that endeavour. We can use the xarray-nightly feature branch to already start the implementation process.

openradar / xradar

map_over_sweeps #216

Description