Rolling window with `as_strided` #1837

fujiisoup · 2018-01-18T09:18:19Z

Closes Slow performance of rolling.reduce #1831, win_type for rolling() ? #1142, N-D rolling #819
Tests added
Tests passed
Passes git diff upstream/master **/*py | flake8 --diff
Fully documented, including whats-new.rst for all changes and api.rst for new API

I started to work for refactoring rollings.
As suggested in #1831 comment, I implemented rolling_window methods based on as_strided.

I got more than 1,000 times speed up! yey!

In [1]: import numpy as np
   ...: import xarray as xr
   ...: 
   ...: da = xr.DataArray(np.random.randn(10000, 3), dims=['x', 'y'])

with the master

%timeit da.rolling(x=5).reduce(np.mean)
1 loop, best of 3: 9.68 s per loop

with the current implementation

%timeit da.rolling(x=5).reduce(np.mean)
100 loops, best of 3: 5.29 ms per loop

and with the bottleneck

%timeit da.rolling(x=5).mean()
100 loops, best of 3: 2.62 ms per loop

My current concerns are

Can we expose the new rolling_window method of DataArray and Dataset to the public?
I think this method itself is useful for many usecases, such as short-term-FFT and convolution.
This also gives more flexible rolling operation, such as windowed moving average, strided rolling, and ND-rolling.
Is there any dask's equivalence to numpy's as_strided?
Currently, I just use a slice->concatenate path, but I don't think it is very efficient.
(Is it already efficient, as dask utilizes out-of-core computation?)

Any thoughts are welcome.

shoyer

Very nice! I have some minor suggestions for organization but this looks like a nice win for performance

shoyer · 2018-01-18T19:19:20Z

xarray/core/dataarray.py

@@ -2132,6 +2132,50 @@ def rank(self, dim, pct=False, keep_attrs=False):
        ds = self._to_temp_dataset().rank(dim, pct=pct, keep_attrs=keep_attrs)
        return self._from_temp_dataset(ds)

+    def rolling_window(self, dim, window, window_dim, center=True):


Can we make this a Rolling method instead? Maybe construct or build would be a good name for the method?

The API could look something like: da.rolling(b=3).construct('window_dim')

That helps keep the main DataArray/Dataset namespace a little more organized.

How about to_dataarray and to_dataset?
Personally, construct sounds something necessary to compute rolling method.

shoyer · 2018-01-18T19:20:20Z

xarray/core/variable.py

+        if isinstance(self.data, dask_array_type):
+            array = self.data
+
+            for d, pad in pad_widths.items():


add a comment noting dask/dask#1926 here

shoyer · 2018-01-18T19:21:13Z

xarray/core/variable.py

@@ -936,6 +936,43 @@ def shift(self, **shifts):
            result = result._shift_one_dim(dim, count)
        return result

+    def _pad(self, **pad_widths):


Call this _pad_with_fill_value?

shoyer · 2018-01-18T19:23:11Z

xarray/core/indexing.py

@@ -815,6 +820,17 @@ def __setitem__(self, key, value):
        array, key = self._indexing_array_and_key(key)
        array[key] = value

+    def rolling_window(self, axis, window):


Can we put this in duck_array_ops instead? I'd like to keep the indexing adapters simple.

shoyer · 2018-01-18T19:25:12Z

xarray/core/indexing.py

+        """
+
+        axis = nputils._validate_axis(self.array, axis)
+        rolling = nputils.rolling_window(np.swapaxes(self.array, axis, -1),


maybe add an axis argument directly to nputils.rolling_window?

shoyer · 2018-01-18T19:26:37Z

xarray/core/indexing.py

+        size = self.array.shape[axis] - window + 1
+        arrays = [self.array[(slice(None), ) * axis + (slice(w, size + w), )]
+                  for w in range(window)]
+        return da.stack(arrays, axis=-1)


Probably the most efficient way to do this would be to use dask's ghost cell support, but this is fine for now:
http://dask.pydata.org/en/latest/array-ghost.html

You would want to map the new efficient rolling window computation over each block.

see here as well:

xarray/xarray/core/dask_array_ops.py

Lines 11 to 25 in f3deb2f

def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):

'''wrapper to apply bottleneck moving window funcs on dask arrays'''

# inputs for ghost

if axis < 0:

axis = a.ndim + axis

depth = {d: 0 for d in range(a.ndim)}

depth[axis] = window - 1

boundary = {d: np.nan for d in range(a.ndim)}

# create ghosted arrays

ag = da.ghost.ghost(a, depth=depth, boundary=boundary)

# apply rolling func

out = ag.map_blocks(moving_func, window, min_count=min_count,

axis=axis, dtype=a.dtype)

# trim array

result = da.ghost.trim_internal(out, depth)

Thank you for the help.
But I don't yet come up with the solution for dask...
(Do you mean that we want to do as_strided-like operation for each ghosted-chunk?)

I think this could be in another PR, as it would take for along time for me to think of the correct path.

… public.

fujiisoup · 2018-01-19T14:35:31Z

During the work, I notice a somehow unexpected behavior of rolling.

I expected that with min_periods=1 option, we will get an array without nan,
It is true with center=False and also pd.rolling.

In [2]: da = xr.DataArray(np.arange(10), dims='x')
In [3]: da.rolling(x=3, min_periods=1, center=False).sum()
Out[3]: 
<xarray.DataArray (x: 10)>
array([  0.,   1.,   3.,   6.,   9.,  12.,  15.,  18.,  21.,  24.])
Dimensions without coordinates: x

In [5]: s = pd.Series(np.arange(10))
 s.rolling(3, min_periods=1, center=True).sum()
Out[7]: 
0     1.0
1     3.0
2     6.0
3     9.0
4    12.0
5    15.0
6    18.0
7    21.0
8    24.0
9    17.0
dtype: float64

But with center=True, we have a nan at the end.

In [4]: da.rolling(x=3, min_periods=1, center=True).sum()
Out[4]: 
<xarray.DataArray (x: 10)>
array([  1.,   3.,   6.,   9.,  12.,  15.,  18.,  21.,  24.,  nan])
Dimensions without coordinates: x

It is because we make shift operation after the rolling.

xarray/xarray/core/rolling.py

Lines 263 to 269 in 74d8318

    
               values = func(self.obj.data, window=self.window, 
        
                             min_count=min_count, axis=axis) 
        
           result = DataArray(values, self.obj.coords) 
        
           if self.center: 
        
               result = self._center_result(result)

If we pad the array before the rolling operation instead of shift, we will not get the last nan and the result would be the same to pandas.
(rolling.to_dataarray('window_dim') does this).

I think this path is more intuitive.
Any thoughts?

…_value to public.

fujiisoup · 2018-01-20T06:52:51Z

xarray/core/dtypes.py

@@ -34,6 +34,8 @@ def maybe_promote(dtype):
        fill_value = np.datetime64('NaT')
    elif np.issubdtype(dtype, np.timedelta64):
        fill_value = np.timedelta64('NaT')
+    elif dtype.kind == 'b':
+        fill_value = False


This is convenient for me, but it is not very clear whether False is equivalent to nan for boolean arrays.
If anyone has objections, I will consider different approach.

Indeed, let's consider other options here. This is used for the default value when reindexing/aligning.

fujiisoup · 2018-01-20T06:55:25Z

xarray/tests/test_dataarray.py

-                               da_rolling['index'])
+
+    np.testing.assert_allclose(s_rolling.values, da_rolling.values)
+    np.testing.assert_allclose(s_rolling.index, da_rolling['index'])


I updated the logic for the center=True case. Now our result is equivalent to pandas's rolling, including the last position.

Awesome, thanks!

fujiisoup · 2018-01-20T06:58:09Z

xarray/core/variable.py

@@ -936,6 +936,43 @@ def shift(self, **shifts):
            result = result._shift_one_dim(dim, count)
        return result

+    def pad_with_fill_value(self, **pad_widths):


I want to expose this to the public, which is used in my new logic in rolling.

shoyer

Looks great -- I have only a few smaller suggestions.

shoyer · 2018-01-20T19:43:20Z

xarray/core/dtypes.py

@@ -34,6 +34,8 @@ def maybe_promote(dtype):
        fill_value = np.datetime64('NaT')
    elif np.issubdtype(dtype, np.timedelta64):
        fill_value = np.timedelta64('NaT')
+    elif dtype.kind == 'b':
+        fill_value = False


Indeed, let's consider other options here. This is used for the default value when reindexing/aligning.

shoyer · 2018-01-20T19:47:55Z

xarray/core/rolling.py

+        # Find valid windows based on count.
+        # We do not use `reduced.count()` because it constructs a larger array
+        # (notice that `windows` is just a view)
+        counts = (~self.obj.isnull()).rolling(


For formatting long chains of method calls, I like to add extra parentheses and break every operation at the start of the line, e.g.,

counts = ((~self.obj.isnull()) .rolling(center=self.center, **{self.dim: self.window}) .to_dataarray('_rolling_window_dim') .sum(dim='_rolling_window_dim'))

I find this makes it easier to read

shoyer · 2018-01-20T19:48:51Z

xarray/core/rolling.py

-
-        return result
+        # restore dim order
+        return result.transpose(*self.obj.dims)


I don't think we need to restore dimension order any more. The result should already be calculated correctly.

shoyer · 2018-01-20T19:52:05Z

xarray/tests/test_dataarray.py

-                               da_rolling['index'])
+
+    np.testing.assert_allclose(s_rolling.values, da_rolling.values)
+    np.testing.assert_allclose(s_rolling.index, da_rolling['index'])


Awesome, thanks!

shoyer · 2018-01-20T19:53:19Z

doc/computation.rst


 .. ipython:: python

   @verbatim
   for label, arr_window in r:
      # arr_window is a view of x

+Finally, the rolling object has ``to_dataarray`` method, which gives a


Maybe add: (to_dataset for Rolling objects from Dataset)

shoyer · 2018-01-20T19:53:50Z

doc/whats-new.rst

+  dimension added to the last position. This enables more flexible operation,
+  such as strided rolling, windowed rolling, ND-rolling, and convolution.
+  (:issue:`1831`, :issue:`1142`, :issue:`819`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - Added nodatavals attribute to DataArray when using :py:func:`~xarray.open_rasterio`. (:issue:`1736`).


Add a bug fix note for the aggregations of the last element with center=True?

shoyer · 2018-01-20T19:55:45Z

xarray/core/rolling.py

+        # Find valid windows based on count.
+        # We do not use `reduced.count()` because it constructs a larger array
+        # (notice that `windows` is just a view)
+        counts = (~self.obj.isnull()).rolling(


Maybe we should add a short-cut here that doesn't bother to compute counts if the array's dtype cannot hold NaN? I think that would solve the issue with changing maybe_promote for booleans.

You could add a utility function to determine this based on whether the result of maybe_promote() has the same dtype as the input.

… review.

shoyer · 2018-01-21T00:47:45Z

xarray/core/nputils.py

@@ -133,3 +134,52 @@ def __setitem__(self, key, value):
        mixed_positions, vindex_positions = _advanced_indexer_subspaces(key)
        self._array[key] = np.moveaxis(value, vindex_positions,
                                       mixed_positions)
+
+
+def rolling_window(a, axis, window):


This is a small point, but can you swap the arguments for this function? That would let you set a default axis.

Bottleneck uses default arguments like move_sum(array, window, axis=-1) which I think is a nice convention:
https://kwgoodman.github.io/bottleneck-doc/reference.html#moving-window-functions

fujiisoup · 2018-01-21T07:01:38Z

xarray/core/variable.py

+        **pad_width: keyword arguments of the form {dim: (before, after)}
+            Number of values padded to the edges of each dimension.
+        """
+        if self.dtype.kind == 'b':


Is there a better way to get an appropriate fill_value for non-float arrays?

Maybe this function should take a fill value argument, which could default to dtypes.NA?

fujiisoup · 2018-01-21T07:03:11Z

xarray/tests/test_dataarray.py

@@ -3435,7 +3448,7 @@ def test_rolling_count_correct():
    result = da.rolling(time=11, min_periods=None).count()
    expected = DataArray(
        [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan,
-         np.nan, np.nan, np.nan, np.nan, 8], dims='time')
+         np.nan, np.nan, np.nan, np.nan, np.nan], dims='time')


I think the last element should be np.nan rather than 8, because 8 < min_periods=11.

…axis)

shoyer

Looks good to me -- thanks for all your work on this!

fujiisoup · 2018-02-19T02:14:35Z

@shoyer , thanks for the detailed review.

I noticed the benchmark test is still failing.
After fixing this, I will merge this.

Thanks :)

stickler-ci · 2018-02-24T15:15:42Z

xarray/tests/test_variable.py

+                              center=True)
+        # window/2 should be smaller than the smallest chunk size.
+        with pytest.raises(ValueError):
+            rw = v.rolling_window(dim='x', window=100, window_dim='x_w',


F841 local variable 'rw' is assigned to but never used

stickler-ci · 2018-02-24T15:19:14Z

xarray/tests/test_variable.py

+        import dask.array as da
+        v = Variable(['x'], da.arange(100, chunks=20))
+        # should not raise
+        rw = v.rolling_window(dim='x', window=10, window_dim='x_w',


F841 local variable 'rw' is assigned to but never used

shoyer · 2018-02-25T04:45:27Z

xarray/core/dask_array_ops.py

+    else:
+        start, end = window - 1, 0
+
+    drop_size = depth[axis] - offset - np.maximum(start, end)


Normally I think of size as a positive integer, but below you use -drop_size to make it positive. I think this would be clearer as drop_size = max(start, end) - offset - depth[axis] (use max() vs np.maximum as start and end are Python integers)

You are right. I thought it becomes sometimes negative.
Fixed.

shoyer · 2018-02-25T04:46:16Z

xarray/core/dask_array_ops.py

+            "more evenly divides the shape of your array." %
+            (window, depth[axis], min(a.chunks[axis])))
+
+    # We temporary use `reflect` boundary here, but the edge portion is


No longer correct?

shoyer · 2018-02-25T04:47:52Z

xarray/core/variable.py

+        """
+        if fill_value is dtypes.NA:  # np.nan is passed
+            dtype, fill_value = dtypes.maybe_promote(self.dtype)
+            array = self.astype(dtype).data


Use self.data.astype(dtype, copy=False) to avoid copying for numpy arrays unless necessary.

# Conflicts: # xarray/tests/test_duck_array_ops.py

stickler-ci · 2018-02-26T04:22:21Z

xarray/tests/test_duck_array_ops.py

+                       fill_value=np.nan)
+
+
+@pytest.mark.skipif(not has_dask, reason='This is for dask.')


F811 redefinition of unused 'test_dask_rolling' from line 308

# Conflicts: # xarray/core/duck_array_ops.py # xarray/core/missing.py # xarray/core/nputils.py # xarray/core/rolling.py # xarray/tests/test_duck_array_ops.py # xarray/tests/test_nputils.py

fujiisoup · 2018-03-01T00:59:37Z

@shyer, do you have further suggestions?
I think this is almost ready.

shoyer

This looks ready to me!

jhamman · 2018-03-01T03:41:48Z

nice work @fujiisoup!

jakirkham · 2018-06-13T05:24:30Z

xarray/core/variable.py

+            array = self.data
+
+            # Dask does not yet support pad. We manually implement it.
+            # https://github.com/dask/dask/issues/1926


Working on a pad function for Dask Arrays (akin to NumPy's pad) in PR ( dask/dask#3578 ). Would be curious to know if this will meet your needs. :)

Nice :)
Thanks for letting me know this.
I will update here after your PR is merged.

In the recent 0.18.1 release. Feedback welcome :)

fujiisoup added 5 commits January 17, 2018 00:27

Rolling_window for np.ndarray

789134c

Add pad method to Variable

fa4e857

Added rolling_window to DataArray and Dataset

52915f3

remove pad_value option. Support dask.rolling_window

b622007

Refactor rolling.reduce

36a1fe9

fujiisoup changed the title ~~Rolling window with as_strided~~ [WIP] Rolling window with as_strided Jan 18, 2018

fujiisoup added 2 commits January 18, 2018 22:54

add as_strided to npcompat. Tests added for reduce(np.nanmean)

71fed0f

Support boolean in maybe_promote

3960134

shoyer reviewed Jan 18, 2018

View reviewed changes

fujiisoup added 2 commits January 19, 2018 15:10

move rolling_window into duck_array_op. Make DataArray.rolling_window…

4bd38f3

… public.

Added to_dataarray and to_dataset to rolling object.

af8362e

fujiisoup added 3 commits January 20, 2018 10:06

Use pad in rolling to make compatible to pandas. Expose pad_with_fill…

76db6b5

…_value to public.

Refactor rolling

87f53af

flake8

c23cedb

fujiisoup commented Jan 20, 2018

View reviewed changes

fujiisoup added 5 commits January 20, 2018 16:06

Added a comment for dask's pad.

9547c57

Use fastpath in rolling.to_dataarray

1f71cff

Merge branch 'master' into rolling_window

724776f

Doc added.

73862eb

Revert not to use fastpath

859bb5c

shoyer reviewed Jan 20, 2018

View reviewed changes

fujiisoup added 5 commits January 21, 2018 10:20

Merge branch 'master' into rolling_window

d5fc24e

Remove maybe_prompt for Boolean. Some improvements based on @shoyer's…

05c72f0

… review.

Update test.

d55e498

Bug fix in test_rolling_count_correct

9393eb2

fill_value for boolean array

9c71a50

shoyer reviewed Jan 21, 2018

View reviewed changes

fujiisoup commented Jan 21, 2018

View reviewed changes

rolling_window(array, axis, window) -> rolling_window(array, window, …

54975b4

…axis)

fujiisoup added 2 commits February 18, 2018 16:34

Update doc

cc9c3d6

Merge branch 'master' into rolling_window

52cc48d

shoyer approved these changes Feb 18, 2018

View reviewed changes

Change boundary and add comments for dask_rolling_window.

2954cdf

Refactor dask_array_ops.rolling_window and np_utils.rolling_window

f19e531

stickler-ci reviewed Feb 24, 2018

View reviewed changes

flake8

a074df3

stickler-ci reviewed Feb 24, 2018

View reviewed changes

fujiisoup added 2 commits February 25, 2018 06:04

Simplify tests

f6f78a5

flake8 again.

0ec8aba

shoyer reviewed Feb 25, 2018

View reviewed changes

fujiisoup added 2 commits February 25, 2018 19:16

cleanup roling_window for dask.

0261cfe

Merge branch 'master' into rolling_window

a91c27f

# Conflicts: # xarray/tests/test_duck_array_ops.py

stickler-ci reviewed Feb 26, 2018

View reviewed changes

fujiisoup added 5 commits February 26, 2018 13:59

remove duplicates

c83d588

remvove duplicate

3bb4668

flake8

d0d89ce

delete unnecessary file.

eaba563

Merge branch 'master' into rolling_window

aeabdf5

# Conflicts: # xarray/core/duck_array_ops.py # xarray/core/missing.py # xarray/core/nputils.py # xarray/core/rolling.py # xarray/tests/test_duck_array_ops.py # xarray/tests/test_nputils.py

shoyer approved these changes Mar 1, 2018

View reviewed changes

fujiisoup merged commit dc3eebf into pydata:master Mar 1, 2018

fujiisoup deleted the rolling_window branch March 1, 2018 03:39

jhamman mentioned this pull request Mar 10, 2018

Efficient rolling 'trick' #1978

Closed

fujiisoup mentioned this pull request Mar 15, 2018

DataArray.rolling().mean() is way slower than it should be #1993

Closed

jakirkham reviewed Jun 13, 2018

View reviewed changes

fujiisoup mentioned this pull request Aug 17, 2018

Inconsistent results when calculating sums on float32 arrays w/ bottleneck installed #2370

Closed

kmsquire mentioned this pull request May 16, 2019

Fix rolling.constuct() example #2967

Merged

	def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):
	'''wrapper to apply bottleneck moving window funcs on dask arrays'''
	# inputs for ghost
	if axis < 0:
	axis = a.ndim + axis
	depth = {d: 0 for d in range(a.ndim)}
	depth[axis] = window - 1
	boundary = {d: np.nan for d in range(a.ndim)}
	# create ghosted arrays
	ag = da.ghost.ghost(a, depth=depth, boundary=boundary)
	# apply rolling func
	out = ag.map_blocks(moving_func, window, min_count=min_count,
	axis=axis, dtype=a.dtype)
	# trim array
	result = da.ghost.trim_internal(out, depth)

		fill_value=np.nan)


		@pytest.mark.skipif(not has_dask, reason='This is for dask.')

Rolling window with as_strided #1837

Rolling window with as_strided #1837

Conversation

fujiisoup commented Jan 18, 2018 • edited Loading

shoyer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fujiisoup commented Jan 19, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shoyer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shoyer left a comment

Choose a reason for hiding this comment

fujiisoup commented Feb 19, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fujiisoup commented Mar 1, 2018

shoyer left a comment

Choose a reason for hiding this comment

jhamman commented Mar 1, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Rolling window with `as_strided` #1837

Rolling window with `as_strided` #1837

fujiisoup commented Jan 18, 2018 •

edited

Loading