def median(array, dim, keep_attrs=False, skipna=False, **kwargs):
""" Runs a median on an dask-backed xarray.
This function does not scale!
It will rechunk along the given dimension, so make sure
your other chunk sizes are small enough that it
will fit into memory.
:param array: An xarray.DataArray or xarray.Dataset wrapping one or more dask arrays
:type array: xarray.DataArray or xarray.Dataset
:param dim: The name of the dim in array to calculate the median
:type dim: str
"""
if type(array) is xr.Dataset:
return array.apply(median,
dim=dim,
keep_attrs=keep_attrs,
skipna=skipna,
**kwargs)
if not hasattr(array.data, 'dask'):
return array.median(dim,
keep_attrs=keep_attrs,
skipna=skipna,
**kwargs)
array = array.chunk({dim: -1})
axis = array.dims.index(dim)
median_func = np.nanmedian if skipna else np.median
blocks = dask.array.map_blocks(median_func,
array.data,
dtype=array.dtype,
drop_axis=axis,
axis=axis,
**kwargs)
new_coords = {
k: v
for k, v in array.coords.items() if k != dim and dim not in v.dims
}
new_dims = tuple(d for d in array.dims if d != dim)
new_attrs = array.attrs if keep_attrs else None
return xr.DataArray(blocks,
coords=new_coords,
dims=new_dims,
attrs=new_attrs)
def test_xarray_reduce_multiple_groupers():
arr = np.ones((4, 12))
labels = np.array(
["a", "a", "c", "c", "c", "b", "b", "c", "c", "b", "b", "f"])
labels = np.array(labels)
labels2 = np.array([1, 2, 2, 1])
da = xr.DataArray(arr,
dims=("x", "y"),
coords={
"labels2": ("x", labels2),
"labels": ("y", labels)
}).expand_dims(z=4)
expected = xr.DataArray(
[[4, 4], [10, 10], [8, 8], [2, 2]],
dims=("labels", "labels2"),
coords={
"labels": ["a", "c", "b", "f"],
"labels2": [1, 2]
},
).expand_dims(z=4)
actual = xarray_reduce(da, da.labels, da.labels2, func="count")
xr.testing.assert_identical(expected, actual)
actual = xarray_reduce(da, "labels", "labels2", func="count", fill_value=0)
xr.testing.assert_identical(expected, actual)
with raise_if_dask_computes():
actual = xarray_reduce(da.chunk({
"x": 2,
"z": 1
}),
da.labels,
da.labels2,
func="count")
xr.testing.assert_identical(expected, actual)
def test_min_count_specific(dask, func, dim):
if dask and not has_dask:
pytest.skip("requires dask")
# Simple array with four non-NaN values.
da = DataArray(np.ones((6, 6), dtype=np.float64) * np.nan, dims=("a", "b"))
da[0][0] = 2
da[0][3] = 2
da[3][0] = 2
da[3][3] = 2
if dask:
da = da.chunk({"a": 3, "b": 3})
# Expected result if we set min_count to the number of non-NaNs in a
# row/column/the entire array.
if dim:
min_count = 2
expected = DataArray(
[4.0, np.nan, np.nan] * 2, dims=("a" if dim == "b" else "b",)
)
else:
min_count = 4
expected = DataArray(8.0 if func == "sum" else 16.0)
# Check for that min_count.
with raise_if_dask_computes():
actual = getattr(da, func)(dim, skipna=True, min_count=min_count)
assert_dask_array(actual, dask)
assert_allclose(actual, expected)
# With min_count being one higher, should get all NaN.
min_count += 1
expected *= np.nan
with raise_if_dask_computes():
actual = getattr(da, func)(dim, skipna=True, min_count=min_count)
assert_dask_array(actual, dask)
assert_allclose(actual, expected)