Applying custom functions to rolling

[rhyswhitley]

Hello,

What is the correct way of implementing a custom function to the rolling method in xarray. For instance, with a pandas DataFrame if I wanted to do something like a cumulative sum with a decay constant over a 7-day window, this operation is quite trivial:

def decayed_sum(x, k=0.5):
    sum_ = 0
    for  i in range(1, x.size+1):
        sum_ += x[-1]*k**i
    return sum_

df = pd.DataFrame({'value':np.random.negative_binomial(5, 0.9, size=60)})
df.rolling(window=7).apply(decayed_sum)

However, with an xarray DataArray I'm at a loss as to how this works or translates to the same operation that would be applied over space and time. So how would the decayed_sum function be applied to a DataArray that looks like this:

da = xr.DataArray(
    np.random.negative_binomial(5, 0.9, size=(60, 10, 10)),
    coords=[range(60), range(10), range(10)],
    dims=["time", "x", "y"],
)

Looking for examples has proved difficult, just wondering if anyone can help me out here? Thanks in advance.

[TomNicholas]

HI @rhyswhitley , sorry you found this tricky. If you have any ideas as to how our documentation could be clearer we are all ears!

That said, do you not just need DataArrayRolling.reduce(func)? I.e. da.rolling(window=7).reduce(decayed_sum, k=0.5)

the same operation that would be applied over space and time.

Can you explain a bit more about what you want to do here? You want to use a 1D function to reduce over a 2D rolling window?

[rhyswhitley]

Thanks for getting back to me @TomNicholas.

You want to use a 1D function to reduce over a 2D rolling window?

Yes, so for some context we could imagine the data as gridded rainfall over time and the decayed_sum function represents a simple 'bucket' being filled and emptied over that period, where the 7-day moving window captures the antecedent conditions. So I need to calculate the bucket 'fullness' per day for each grid point.

I had tried as you suggested, da.rolling(window=7).reduce(decayed_sum, k=0.5), but I get the following error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-8-b29d26464029> in <module>
      5     return sum_
      6 
----> 7 da.rolling(time=7).reduce(decayed_sum, k=0.5)

~/.conda/envs/xarray/lib/python3.7/site-packages/xarray/core/rolling.py in reduce(self, func, **kwargs)
    234         rolling_dim = _get_new_dimname(self.obj.dims, '_rolling_dim')
    235         windows = self.construct(rolling_dim)
--> 236         result = windows.reduce(func, dim=rolling_dim, **kwargs)
    237 
    238         # Find valid windows based on count.

~/.conda/envs/xarray/lib/python3.7/site-packages/xarray/core/dataarray.py in reduce(self, func, dim, axis, keep_attrs, **kwargs)
   1589             summarized data and the indicated dimension(s) removed.
   1590         """
-> 1591         var = self.variable.reduce(func, dim, axis, keep_attrs, **kwargs)
   1592         return self._replace_maybe_drop_dims(var)
   1593 

~/.conda/envs/xarray/lib/python3.7/site-packages/xarray/core/variable.py in reduce(self, func, dim, axis, keep_attrs, allow_lazy, **kwargs)
   1340             axis = self.get_axis_num(dim)
   1341         data = func(self.data if allow_lazy else self.values,
-> 1342                     axis=axis, **kwargs)
   1343 
   1344         if getattr(data, 'shape', ()) == self.shape:

TypeError: decayed_sum() got an unexpected keyword argument 'axis'

My guess is that the rolling method expands the DataArray to 4-dimensions adding the window, but then doesn't know how to reduce this back to the original 3-dimensions? I feel like this should be quite simple, but I have no idea how why it's not working.

If you have any ideas as to how our documentation could be clearer we are all ears!

Don't get me wrong, I think you guys do a great job. But I think this is a critical thing missing in the documentation, as the current examples only demonstrate using numpy functions. Happy to help add this as an example to the docs if we can get it working.

[mathause]

There are several ways to get this to work & I agree it's not trivial. It sounds like reduce expects a function that can be applied over all array elements at once. So I tried to wrap the function in np.vectorize - but that did not work. What did work was to use np.apply_along_axis:

def decayed_sum2(data, axis, k):
    return np.apply_along_axis(decayed_sum, axis[0], data)

da.rolling(time=7).reduce(decayed_sum2, k=0.5)

but then I thought that we can probably write decayed_sum as 'proper' ufunc that makes use of numpy functions (slightly longer due to error checking):

def decayed_sum3(x, axis, k=0.5):

    if isinstance(axis, tuple):
        if len(axis) == 1:
            axis = axis[0]
        else:
            raise ValueError("Can only reduce over one axis!")

    weights = k ** np.arange(x.shape[axis])
    return (x * weights).sum(axis=axis)

da.rolling(time=7).reduce(decayed_sum3, k=0.5)

(this works because the new axis is added at the end, so numpy broadcasting is happy).

Finally, you could also use da.rolling(...).construct(...) to calculate this in xarray directly (but this will certainly materialize the array while the other versions should use less RAM):

weights = xr.DataArray(0.5 ** np.arange(7), dims="time_window")
(da.rolling(time=7).construct("time_window") * weights).sum("time_window")

[TomNicholas]

Using np.apply_along_axis to create a version of reduced_func that accepts an axis kwarg is what I was also going to suggest. I use that trick a lot with xarray.apply_ufunc too.

But I think this is a critical thing missing in the documentation, as the current examples only demonstrate using numpy functions. Happy to help add this as an example to the docs if we can get it working.

As opposed to an example using xarray functions? Or as opposed to an example using a function that doesn't accept an axis kwarg?

[rhyswhitley]

Fantastic! So I can confirm both solutions work for me, and my preference is for following the ufunc style (decayed_sum3). Thanks guys, I'll accept this as answered.

As opposed to an example using xarray functions? Or as opposed to an example using a function that doesn't accept an axis kwarg?

So the latter one (in bold), but this may actually be stated in the documentation and I didn't pick it up. Just wondering if adding an example of the above would help some in the community.

[nepomnyi]

Hello dear Xarray team.

I have the same problem, which I've been failing to solve. Lest to deviate attention from the original question, I will, simply, ask clarifying questions about the provided solutions if you don't mind.

The first question is about the function decayed_sum2(data, axis, k). What is data there? Isn't it the data within the boundaries of the rolling window? If so, how do I obtain it and why didn't you use the data argument in the subsequent line da.rolling(time=7).reduce(decayed_sum2, k=0.5)?

The second question is about the function decayed_sum3(x, axis, k=0.5). I don't understand anything about it. But for the sake of asking something specific, let me formulate the following question. What do you mean by "proper ufunc"? Also, why are you introducing np.arange(x.shape[axis])?

And let me formulate my wish to the team please. But don't take it as a sign of impudence; we, users, tend to have wishes. You, maintainers, can tell me to get lost with my wish. That's perfectly fine.
You see, rolling windows - which are also called moving windows or filters - are of great importance, at least, in engineering data analysis. One would want to "impose" a rolling window on one's data, execute some function within the window, move the window, execute the function again and so on. The crucial point here is that the function can be whatever one wants. But - as far as I understand the discussion above - the Xarray team assumed only one type of functions: the functions that can be applied over all elements within the rolling window at once. That reduces the flexibility of Xarray. If one wants another type of functions, one has to adhere to the headache inducing tricks like the ones given in the solutions. So, my wish is to being able to execute any type of function I want in .reduce(func) easily without any additional tricks, just like the author of the original post exemplified with pandas.

Thank you in advance.
Ivan

[TomNicholas]

Thanks for raising your concerns @nepomnyi ! Feedback is always appreciated :)

The first question is about the function decayed_sum2(data, axis, k). What is data there? Isn't it the data within the boundaries of the rolling window?

I don't think data is special here, I think it's just a numpy array, same as x is in the next example. The examples could have been clearer in this respect.

Isn't it the data within the boundaries of the rolling window?

Yes.

If so, how do I obtain it and why didn't you use the data argument in the subsequent line da.rolling(time=7).reduce(decayed_sum2, k=0.5)?

data (or x) is coming from xarray extracting the da.data attribute of the xr.DataArray object you called .rolling on, and using numpy's sliding_window_view to index it before passing it into the function you pass to .reduce. Therefore you don't need to provide an argument for data explicitly.

What do you mean by "proper ufunc"?

Numpy "universal functions" have an explicit but complicated definition. For the purposes of applying them in xarray the main thing to understand is that they need to accept an axis argument. These docs on xr.apply_ufunc might be helpful.

about the function decayed_sum3(x, axis, k=0.5). I don't understand anything about it.

🥲 @mathause has rewritten the original loop-based cumulative sum by instead creating an array of weights and then performing a weighted sum inside the decayed_sum3 function.

Also, why are you introducing np.arange(x.shape[axis])?

The line weights = k ** np.arange(x.shape[axis]) is essentially a vectorized way of writing the x[-1]*k**i in the first decayed_sum example given. @mathause just used np.arange to create an array without an explicit loop (i.e. in a "vectorized" manner).

But don't take it as a sign of impudence; we, users, tend to have wishes. You, maintainers, can tell me to get lost with my wish. That's perfectly fine.

Remember that this a permeable boundary - users can become contributors and maintainers!

assumed only one type of functions: the functions that can be applied over all elements within the rolling window at once.

Can you describe the other types of functions you are thinking of? All I know of are

1. "Simple rolling window reductions" (e.g. rolling average)
2. "Weighted rolling window reductions", of which rolling exponential windows are one example.

Xarray currently supports (1) for any reduction function, and (2) for any reduction function but only for exponentially decaying weights (see https://docs.xarray.dev/en/latest/generated/xarray.DataArray.rolling_exp.html). We might imagine supporting a general weighted rolling reduction similar to how weighted groupby has been suggested.

So, my wish is to being able to execute any type of function I want in .reduce(func) easily without any additional tricks, just like the author of the original post exemplified with pandas.

In the end the only additional trick here is that the function you pass needs to understand which axis you want to reduce along. This is an unavoidable addition that reflects xarray's generalisation of pandas from 2D to ND.

[nepomnyi]

@TomNicholas thank you for taking time to explain all that to me!

With regards to the data, I had understood it myself, eventually; your explanation confirms my view of what's going on there, thank you.

With regards to my second question, I still don't have a clear understanding of it. But you pointed out the directions I should look at, which is helpful indeed, thank you. You mentioned that np.arange(x.shape[axis]) is a way to vectorize the function which what I have been trying to do in my case.

To answer your last questions I will have to use my case as an example.

I think your very last sentence downplays the complexity associated with the trick. Here's why. .reduce(func) needs an axis. Thus, a user has to incorporate axis into his function. The important point here, as I see it, is that the way axis is incorporated into the user's function depends on a case, it is not a universal procedure. That's what make .reduce(func) difficult to use: every time one has to think how to incorporate axis into one's function.

My understanding of the statement "It sounds like .reduce() expects a function that can be applied over all array elements at once" is the following. Assume a rolling window is a (2,3) numpy array. There is a function add that adds 5 to every element of the array. When I do .reduce(add), I obtain array of the same size (2,3). I.e. I don't actually reduce the dimensions of my rolling window. And I don't have to bother with the axis trick. That's how I understand that a function can be applied over all array elements at once. That is what I meant saying that Xarray's .reduce() works easily only for functions that can be applied over all array elements at once.

Once one wants to actually reduce the dimensions of the rolling window, one has to bother with the axis trick in order to make .reduce() work.

From that perspective, my answer to your question what other kinds of functions I have in mind is all the functions that imply reduction of the rolling window dimensions because all of them will require additional work from the user associated with axis.

You say that kind of work is unavoidable.

I say ok, but there are cases when that kind of work becomes increasingly complicated and my wish is to make it less complicated.

Here's an example of a more complicated case (my case).

I have a 2D vector field (representing a 2D slice of a water flow). Each vector is a velocity. I store my field in an xr.Dataset. The coordinates of my dataset are x and y (i.e. horizontal and vertical dimensions of my vector field in mm). My dataset has 2 data arrays. The first one stores horizontal components of velocity vectors at every point. The second one stores vertical components. I want to take a 9x9 moving window, impose it on every single point in my vector field. Within every moving window, I want to perform a complex set of calculations involving coordinates and velocity components at every point within the moving window and resulting in one number. I want to create a new DataArray in my Dataset and store that number for every point in the field in that new DataArray. I.e. I want to reduce my 2D moving window to a number. Moreover, since the set of calculations I perform in every moving window is complex, I don't want to use for loops. Instead, I want to vectorize my calculations to reduce the computation time. And the cherry on top of all it is the following. I'm dealing with both horizontal and vertical components of velocities. They are stored in different data arrays. The data arrays are named. It would have been easy for me to use the names of my data arrays to write the code for my calculations. You can't imagine how "suprised" I was to learn that xarray.core.rolling.DatasetRolling returns a numpy array! A numpy array where all my named dimensions are not named anymore and I have no idea how exactly my velocity components are stored in this numpy array!

I understand I can find a solution to my case. But I've been looking for it for 2 days already to no avail. I think I would spend at least a week on it didn't I give up and decide to try ds.rolling().construct() instead (which is still ahead of me). From my perspective, .reduce() takes too much effort to implement.

[dcherian]

You can iterate over the rolling object. Does that help?

We could add rolling.map to copy the pandas syntax.

[nepomnyi]

Hello @dcherian, I think rolling.map would be the best.

I could iterate, of course, if I had all the time for me, till the Sun explodes.

.reduce(func) supplies a numpy array to func. If .reduce(func) is called on a Dataset then it supplies - one by one - each of the dataset's DataArrays as numpy arrays to func. All that means .reduce(func) has a pretty niche applicability: it is more or less convenient to use when it is more or less easy to work with the DataArrays as numpy arrays and only one by one.

It means that - for more general cases - either .reduce() has to be rewritten or another solution has to be found.

Searching the Internet, I have discovered another solution which worked in my case. The solution consists of several bits and, man, was it a pain in the rear side to piece those bits together!
Here's the solution.
First, create a rolling window object: rollingW = ds.rolling({'x':3, 'y':3}, center=True)
Second, construct the dataset: dsRoll = rollingW.construct(x='rollWx', y='rollWy')
Third, stack the rolling windows: dsStacked = dsRoll.stack(gridcell=['x','y','t'])
Fourth, group the dataset: dsGroupped = dsStacked.groupby('gridcell')
Fifth, apply my custom function to the groups: newArray = dsGroupped.map(myCustomFunction)
Sixth, unstack the result in order to add it to the original data set: ds['newArray'] = newArray.unstack('gridcell')

As you see, this solution takes 6 steps. And one has to figure out those steps. From this perspective, your rolling.map idea is going be a life savior.

There is one issue I would like to bring up though.

I had solved my problem with mere for loops initially. And it would take approximately 4 minutes to process one velocity field. My entire data set consists of 500*14 velocity fields. It was taking me 2 weeks to process it. That is why I started looking into rolling windows: to speed up the things. The solution that I showed above, takes 1 minute to process one velocity field. Still, too long.

And here we are: I think that if you just copy pandas syntax, you won't make rolling.map fast. If rolling.map was implemented, I wish it was really fast.

Thank you in advance.

[singledoggy]

As I understand it, reduce does not fully utilize the multi-core processing and performs calculations sequentially. Is this why it is slow in similar situations? So, if I want to speed it up simply, perhaps I need to add some parallel operations?

[singledoggy]

This is a straightforward example of parallel computing a function. I hope it can assist those who encounter similar situations.

import dask

roll = Mydata.rolling(time=24)
delayed = [
    dask.delayed(my_function)(arr_window) for label, arr_window in roll
]

result = dask.compute(*delayed)