forked from pydata/xarray
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdask_array_ops.py
More file actions
26 lines (23 loc) · 783 Bytes
/
dask_array_ops.py
File metadata and controls
26 lines (23 loc) · 783 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
"""Define core operations for xarray objects.
"""
import numpy as np
try:
import dask.array as da
except ImportError:
pass
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)
return result