diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 5ef8ed119ec..cf81671c333 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -3358,9 +3358,12 @@ def map_blocks( ... clim = gb.mean(dim="time") ... return gb - clim >>> time = xr.cftime_range("1990-01", "1992-01", freq="M") + >>> month = xr.DataArray(time.month, coords={"time": time}, dims=["time"]) >>> np.random.seed(123) >>> array = xr.DataArray( - ... np.random.rand(len(time)), dims="time", coords=[time] + ... np.random.rand(len(time)), + ... dims=["time"], + ... coords={"time": time, "month": month}, ... ).chunk() >>> array.map_blocks(calculate_anomaly, template=array).compute() @@ -3371,21 +3374,19 @@ def map_blocks( 0.07673453, 0.22865714, 0.19063865, -0.0590131 ]) Coordinates: * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + month (time) int64 1 2 3 4 5 6 7 8 9 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 Note that one must explicitly use ``args=[]`` and ``kwargs={}`` to pass arguments to the function being applied in ``xr.map_blocks()``: >>> array.map_blocks( ... calculate_anomaly, kwargs={"groupby_type": "time.year"}, template=array, - ... ) + ... ) # doctest: +ELLIPSIS - array([ 0.15361741, -0.25671244, -0.31600032, 0.008463 , 0.1766172 , - -0.11974531, 0.43791243, 0.14197797, -0.06191987, -0.15073425, - -0.19967375, 0.18619794, -0.05100474, -0.42989909, -0.09153273, - 0.24841842, -0.30708526, -0.31412523, 0.04197439, 0.0422506 , - 0.14482397, 0.35985481, 0.23487834, 0.12144652]) + dask.array Coordinates: - * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + month (time) int64 dask.array """ from .parallel import map_blocks @@ -3875,9 +3876,10 @@ def argmin( >>> array.isel(array.argmin(...)) array(-1) - >>> array = xr.DataArray([[[3, 2, 1], [3, 1, 2], [2, 1, 3]], - ... [[1, 3, 2], [2, -5, 1], [2, 3, 1]]], - ... dims=("x", "y", "z")) + >>> array = xr.DataArray( + ... [[[3, 2, 1], [3, 1, 2], [2, 1, 3]], [[1, 3, 2], [2, -5, 1], [2, 3, 1]]], + ... dims=("x", "y", "z"), + ... ) >>> array.min(dim="x") array([[ 1, 2, 1], @@ -3977,9 +3979,10 @@ def argmax( array(3) - >>> array = xr.DataArray([[[3, 2, 1], [3, 1, 2], [2, 1, 3]], - ... [[1, 3, 2], [2, 5, 1], [2, 3, 1]]], - ... dims=("x", "y", "z")) + >>> array = xr.DataArray( + ... [[[3, 2, 1], [3, 1, 2], [2, 1, 3]], [[1, 3, 2], [2, 5, 1], [2, 3, 1]]], + ... dims=("x", "y", "z"), + ... ) >>> array.max(dim="x") array([[3, 3, 2], diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 69805f086f6..9fbaf7479db 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -5817,20 +5817,22 @@ def map_blocks( ... clim = gb.mean(dim="time") ... return gb - clim >>> time = xr.cftime_range("1990-01", "1992-01", freq="M") + >>> month = xr.DataArray(time.month, coords={"time": time}, dims=["time"]) >>> np.random.seed(123) >>> array = xr.DataArray( - ... np.random.rand(len(time)), dims="time", coords=[time] + ... np.random.rand(len(time)), + ... dims=["time"], + ... coords={"time": time, "month": month}, ... ).chunk() >>> ds = xr.Dataset({"a": array}) >>> ds.map_blocks(calculate_anomaly, template=ds).compute() - - array([ 0.12894847, 0.11323072, -0.0855964 , -0.09334032, 0.26848862, - 0.12382735, 0.22460641, 0.07650108, -0.07673453, -0.22865714, - -0.19063865, 0.0590131 , -0.12894847, -0.11323072, 0.0855964 , - 0.09334032, -0.26848862, -0.12382735, -0.22460641, -0.07650108, - 0.07673453, 0.22865714, 0.19063865, -0.0590131 ]) + + Dimensions: (time: 24) Coordinates: * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + month (time) int64 1 2 3 4 5 6 7 8 9 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 + Data variables: + a (time) float64 0.1289 0.1132 -0.0856 ... 0.2287 0.1906 -0.05901 Note that one must explicitly use ``args=[]`` and ``kwargs={}`` to pass arguments to the function being applied in ``xr.map_blocks()``: @@ -5838,14 +5840,13 @@ def map_blocks( >>> ds.map_blocks( ... calculate_anomaly, kwargs={"groupby_type": "time.year"}, template=ds, ... ) - - array([ 0.15361741, -0.25671244, -0.31600032, 0.008463 , 0.1766172 , - -0.11974531, 0.43791243, 0.14197797, -0.06191987, -0.15073425, - -0.19967375, 0.18619794, -0.05100474, -0.42989909, -0.09153273, - 0.24841842, -0.30708526, -0.31412523, 0.04197439, 0.0422506 , - 0.14482397, 0.35985481, 0.23487834, 0.12144652]) + + Dimensions: (time: 24) Coordinates: - * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + month (time) int64 dask.array + Data variables: + a (time) float64 dask.array """ from .parallel import map_blocks diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py index 07d61e595c9..6d5456f77f7 100644 --- a/xarray/core/parallel.py +++ b/xarray/core/parallel.py @@ -235,11 +235,14 @@ def map_blocks( ... clim = gb.mean(dim="time") ... return gb - clim >>> time = xr.cftime_range("1990-01", "1992-01", freq="M") + >>> month = xr.DataArray(time.month, coords={"time": time}, dims=["time"]) >>> np.random.seed(123) >>> array = xr.DataArray( - ... np.random.rand(len(time)), dims="time", coords=[time] + ... np.random.rand(len(time)), + ... dims=["time"], + ... coords={"time": time, "month": month}, ... ).chunk() - >>> xr.map_blocks(calculate_anomaly, array, template=array).compute() + >>> array.map_blocks(calculate_anomaly, template=array).compute() array([ 0.12894847, 0.11323072, -0.0855964 , -0.09334032, 0.26848862, 0.12382735, 0.22460641, 0.07650108, -0.07673453, -0.22865714, @@ -248,25 +251,20 @@ def map_blocks( 0.07673453, 0.22865714, 0.19063865, -0.0590131 ]) Coordinates: * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + month (time) int64 1 2 3 4 5 6 7 8 9 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 Note that one must explicitly use ``args=[]`` and ``kwargs={}`` to pass arguments to the function being applied in ``xr.map_blocks()``: - >>> xr.map_blocks( - ... calculate_anomaly, - ... array, - ... kwargs={"groupby_type": "time.year"}, - ... template=array, - ... ) + >>> array.map_blocks( + ... calculate_anomaly, kwargs={"groupby_type": "time.year"}, template=array, + ... ) # doctest: +ELLIPSIS - array([ 0.15361741, -0.25671244, -0.31600032, 0.008463 , 0.1766172 , - -0.11974531, 0.43791243, 0.14197797, -0.06191987, -0.15073425, - -0.19967375, 0.18619794, -0.05100474, -0.42989909, -0.09153273, - 0.24841842, -0.30708526, -0.31412523, 0.04197439, 0.0422506 , - 0.14482397, 0.35985481, 0.23487834, 0.12144652]) + dask.array Coordinates: - * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 - """ + * time (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00 + month (time) int64 dask.array + """ def _wrapper( func: Callable,