Memory efficient reduce functions

[mortonjt]

Hi, I'm trying to create memory efficient reduce functions, that can apply a map-reduce without blowing up memory

For instance, consider the following example, where f : N -> M could be a projection from a low dim space to a high dim space.

I'd like code to function something as follows

def _reduce_sum_f(fx, y, f):
    """Reduce the sum of f(x)"""
    return fx + f(y)

fg = lambda fx, y: _reduce_sum_f(fx, y, f)


# Create sample data
import xarray as xr
import numpy as np
x_values = np.arange(0, 5)
y_values = np.arange(10, 15)
data_values = np.random.rand(len(x_values), len(y_values))

# Create DataArray objects
x_da = xr.DataArray(x_values, dims='x', name='x')
y_da = xr.DataArray(y_values, dims='y', name='y')
data_da = xr.DataArray(data_values, dims=('x', 'y'), name='data')

# Create the Dataset
dataset = xr.Dataset({'data': data_da}, coords={'x': x_da, 'y': y_da})

# Apply toy reduce function
dataset.reduce(fg, axis=0)

Now, this code does not work (in this case, there is an error through for func doesn't have 'axis' specified). I'm not sure if this is currently possible in the xarray framework, but it would be strongly advantageous to have, since it could enable complex summaries of xarray objects using arbitrary functions.

[dcherian]

We usually use dask for this kind of thing.

[mortonjt]

Thanks. I'll need to tinker more with lazy evaluation with xarray / dask. If you have any favorite xarray+dask projects that leverage this type of approach, feel free to link.

[TomNicholas]

Hi @mortonjt !

If you want to do a common reduction (e.g. mean) without blowing up memory then xarray will do this automatically if dask is installed, as Deepak says. (You should read our docs page on dask though).

If you want to apply a custom function (including a custom reduction) then you need to look at xarray.apply_ufunc, which can apply custom dask-aware operations.