from __future__ import annotations import datetime import operator import warnings from unittest import mock import numpy as np import pandas as pd import pytest from packaging.version import Version import xarray as xr from xarray import DataArray, Dataset, Variable from xarray.core.groupby import _consolidate_slices from xarray.core.types import InterpOptions from xarray.tests import ( InaccessibleArray, assert_allclose, assert_equal, assert_identical, create_test_data, has_cftime, has_flox, requires_dask, requires_flox, requires_scipy, ) @pytest.fixture def dataset() -> xr.Dataset: ds = xr.Dataset( { "foo": (("x", "y", "z"), np.random.randn(3, 4, 2)), "baz": ("x", ["e", "f", "g"]), "cat": ("y", pd.Categorical(["cat1", "cat2", "cat2", "cat1"])), }, {"x": ("x", ["a", "b", "c"], {"name": "x"}), "y": [1, 2, 3, 4], "z": [1, 2]}, ) ds["boo"] = (("z", "y"), [["f", "g", "h", "j"]] * 2) return ds @pytest.fixture def array(dataset) -> xr.DataArray: return dataset["foo"] def test_consolidate_slices() -> None: assert _consolidate_slices([slice(3), slice(3, 5)]) == [slice(5)] assert _consolidate_slices([slice(2, 3), slice(3, 6)]) == [slice(2, 6)] assert _consolidate_slices([slice(2, 3, 1), slice(3, 6, 1)]) == [slice(2, 6, 1)] slices = [slice(2, 3), slice(5, 6)] assert _consolidate_slices(slices) == slices # ignore type because we're checking for an error anyway with pytest.raises(ValueError): _consolidate_slices([slice(3), 4]) # type: ignore[list-item] @pytest.mark.filterwarnings("ignore:return type") def test_groupby_dims_property(dataset, recwarn) -> None: # dims is sensitive to squeeze, always warn with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert dataset.groupby("x").dims == dataset.isel(x=1).dims assert dataset.groupby("y").dims == dataset.isel(y=1).dims # in pytest-8, pytest.warns() no longer clears all warnings recwarn.clear() # when squeeze=False, no warning should be raised assert tuple(dataset.groupby("x", squeeze=False).dims) == tuple( dataset.isel(x=slice(1, 2)).dims ) assert tuple(dataset.groupby("y", squeeze=False).dims) == tuple( dataset.isel(y=slice(1, 2)).dims ) assert len(recwarn) == 0 dataset = dataset.drop_vars(["cat"]) stacked = dataset.stack({"xy": ("x", "y")}) assert tuple(stacked.groupby("xy", squeeze=False).dims) == tuple( stacked.isel(xy=[0]).dims ) assert len(recwarn) == 0 def test_groupby_sizes_property(dataset) -> None: with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert dataset.groupby("x").sizes == dataset.isel(x=1).sizes with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert dataset.groupby("y").sizes == dataset.isel(y=1).sizes dataset = dataset.drop_vars("cat") stacked = dataset.stack({"xy": ("x", "y")}) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert stacked.groupby("xy").sizes == stacked.isel(xy=0).sizes def test_multi_index_groupby_map(dataset) -> None: # regression test for GH873 ds = dataset.isel(z=1, drop=True)[["foo"]] expected = 2 * ds # The function in `map` may be sensitive to squeeze, always warn with pytest.warns(UserWarning, match="The `squeeze` kwarg"): actual = ( ds.stack(space=["x", "y"]) .groupby("space") .map(lambda x: 2 * x) .unstack("space") ) assert_equal(expected, actual) def test_reduce_numeric_only(dataset) -> None: gb = dataset.groupby("x", squeeze=False) with xr.set_options(use_flox=False): expected = gb.sum() with xr.set_options(use_flox=True): actual = gb.sum() assert_identical(expected, actual) def test_multi_index_groupby_sum() -> None: # regression test for GH873 ds = xr.Dataset( {"foo": (("x", "y", "z"), np.ones((3, 4, 2)))}, {"x": ["a", "b", "c"], "y": [1, 2, 3, 4]}, ) expected = ds.sum("z") actual = ds.stack(space=["x", "y"]).groupby("space").sum("z").unstack("space") assert_equal(expected, actual) def test_groupby_da_datetime() -> None: # test groupby with a DataArray of dtype datetime for GH1132 # create test data times = pd.date_range("2000-01-01", periods=4) foo = xr.DataArray([1, 2, 3, 4], coords=dict(time=times), dims="time") # create test index reference_dates = [times[0], times[2]] labels = reference_dates[0:1] * 2 + reference_dates[1:2] * 2 ind = xr.DataArray( labels, coords=dict(time=times), dims="time", name="reference_date" ) g = foo.groupby(ind) actual = g.sum(dim="time") expected = xr.DataArray( [3, 7], coords=dict(reference_date=reference_dates), dims="reference_date" ) assert_equal(expected, actual) def test_groupby_duplicate_coordinate_labels() -> None: # fix for http://stackoverflow.com/questions/38065129 array = xr.DataArray([1, 2, 3], [("x", [1, 1, 2])]) expected = xr.DataArray([3, 3], [("x", [1, 2])]) actual = array.groupby("x").sum() assert_equal(expected, actual) def test_groupby_input_mutation() -> None: # regression test for GH2153 array = xr.DataArray([1, 2, 3], [("x", [2, 2, 1])]) array_copy = array.copy() expected = xr.DataArray([3, 3], [("x", [1, 2])]) actual = array.groupby("x").sum() assert_identical(expected, actual) assert_identical(array, array_copy) # should not modify inputs @pytest.mark.parametrize("use_flox", [True, False]) def test_groupby_indexvariable(use_flox: bool) -> None: # regression test for GH7919 array = xr.DataArray([1, 2, 3], [("x", [2, 2, 1])]) iv = xr.IndexVariable(dims="x", data=pd.Index(array.x.values)) with xr.set_options(use_flox=use_flox): actual = array.groupby(iv).sum() actual = array.groupby(iv).sum() expected = xr.DataArray([3, 3], [("x", [1, 2])]) assert_identical(expected, actual) @pytest.mark.parametrize( "obj", [ xr.DataArray([1, 2, 3, 4, 5, 6], [("x", [1, 1, 1, 2, 2, 2])]), xr.Dataset({"foo": ("x", [1, 2, 3, 4, 5, 6])}, {"x": [1, 1, 1, 2, 2, 2]}), ], ) def test_groupby_map_shrink_groups(obj) -> None: expected = obj.isel(x=[0, 1, 3, 4]) actual = obj.groupby("x").map(lambda f: f.isel(x=[0, 1])) assert_identical(expected, actual) @pytest.mark.parametrize( "obj", [ xr.DataArray([1, 2, 3], [("x", [1, 2, 2])]), xr.Dataset({"foo": ("x", [1, 2, 3])}, {"x": [1, 2, 2]}), ], ) def test_groupby_map_change_group_size(obj) -> None: def func(group): if group.sizes["x"] == 1: result = group.isel(x=[0, 0]) else: result = group.isel(x=[0]) return result expected = obj.isel(x=[0, 0, 1]) actual = obj.groupby("x").map(func) assert_identical(expected, actual) def test_da_groupby_map_func_args() -> None: def func(arg1, arg2, arg3=0): return arg1 + arg2 + arg3 array = xr.DataArray([1, 1, 1], [("x", [1, 2, 3])]) expected = xr.DataArray([3, 3, 3], [("x", [1, 2, 3])]) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): actual = array.groupby("x").map(func, args=(1,), arg3=1) assert_identical(expected, actual) def test_ds_groupby_map_func_args() -> None: def func(arg1, arg2, arg3=0): return arg1 + arg2 + arg3 dataset = xr.Dataset({"foo": ("x", [1, 1, 1])}, {"x": [1, 2, 3]}) expected = xr.Dataset({"foo": ("x", [3, 3, 3])}, {"x": [1, 2, 3]}) # The function in `map` may be sensitive to squeeze, always warn with pytest.warns(UserWarning, match="The `squeeze` kwarg"): actual = dataset.groupby("x").map(func, args=(1,), arg3=1) assert_identical(expected, actual) def test_da_groupby_empty() -> None: empty_array = xr.DataArray([], dims="dim") with pytest.raises(ValueError): empty_array.groupby("dim") @requires_dask def test_dask_da_groupby_quantile() -> None: # Only works when the grouped reduction can run blockwise # Scalar quantile expected = xr.DataArray( data=[2, 5], coords={"x": [1, 2], "quantile": 0.5}, dims="x" ) array = xr.DataArray( data=[1, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x" ) with pytest.raises(ValueError): array.chunk(x=1).groupby("x").quantile(0.5) # will work blockwise with flox actual = array.chunk(x=3).groupby("x").quantile(0.5) assert_identical(expected, actual) # will work blockwise with flox actual = array.chunk(x=-1).groupby("x").quantile(0.5) assert_identical(expected, actual) @requires_dask def test_dask_da_groupby_median() -> None: expected = xr.DataArray(data=[2, 5], coords={"x": [1, 2]}, dims="x") array = xr.DataArray( data=[1, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x" ) with xr.set_options(use_flox=False): actual = array.chunk(x=1).groupby("x").median() assert_identical(expected, actual) with xr.set_options(use_flox=True): actual = array.chunk(x=1).groupby("x").median() assert_identical(expected, actual) # will work blockwise with flox actual = array.chunk(x=3).groupby("x").median() assert_identical(expected, actual) # will work blockwise with flox actual = array.chunk(x=-1).groupby("x").median() assert_identical(expected, actual) def test_da_groupby_quantile() -> None: array = xr.DataArray( data=[1, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x" ) # Scalar quantile expected = xr.DataArray( data=[2, 5], coords={"x": [1, 2], "quantile": 0.5}, dims="x" ) actual = array.groupby("x").quantile(0.5) assert_identical(expected, actual) # Vector quantile expected = xr.DataArray( data=[[1, 3], [4, 6]], coords={"x": [1, 2], "quantile": [0, 1]}, dims=("x", "quantile"), ) actual = array.groupby("x").quantile([0, 1]) assert_identical(expected, actual) array = xr.DataArray( data=[np.nan, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x" ) for skipna in (True, False, None): e = [np.nan, 5] if skipna is False else [2.5, 5] expected = xr.DataArray(data=e, coords={"x": [1, 2], "quantile": 0.5}, dims="x") actual = array.groupby("x").quantile(0.5, skipna=skipna) assert_identical(expected, actual) # Multiple dimensions array = xr.DataArray( data=[[1, 11, 26], [2, 12, 22], [3, 13, 23], [4, 16, 24], [5, 15, 25]], coords={"x": [1, 1, 1, 2, 2], "y": [0, 0, 1]}, dims=("x", "y"), ) actual_x = array.groupby("x").quantile(0, dim=...) expected_x = xr.DataArray( data=[1, 4], coords={"x": [1, 2], "quantile": 0}, dims="x" ) assert_identical(expected_x, actual_x) actual_y = array.groupby("y").quantile(0, dim=...) expected_y = xr.DataArray( data=[1, 22], coords={"y": [0, 1], "quantile": 0}, dims="y" ) assert_identical(expected_y, actual_y) actual_xx = array.groupby("x").quantile(0) expected_xx = xr.DataArray( data=[[1, 11, 22], [4, 15, 24]], coords={"x": [1, 2], "y": [0, 0, 1], "quantile": 0}, dims=("x", "y"), ) assert_identical(expected_xx, actual_xx) actual_yy = array.groupby("y").quantile(0) expected_yy = xr.DataArray( data=[[1, 26], [2, 22], [3, 23], [4, 24], [5, 25]], coords={"x": [1, 1, 1, 2, 2], "y": [0, 1], "quantile": 0}, dims=("x", "y"), ) assert_identical(expected_yy, actual_yy) times = pd.date_range("2000-01-01", periods=365) x = [0, 1] foo = xr.DataArray( np.reshape(np.arange(365 * 2), (365, 2)), coords={"time": times, "x": x}, dims=("time", "x"), ) g = foo.groupby(foo.time.dt.month) actual = g.quantile(0, dim=...) expected = xr.DataArray( data=[ 0.0, 62.0, 120.0, 182.0, 242.0, 304.0, 364.0, 426.0, 488.0, 548.0, 610.0, 670.0, ], coords={"month": np.arange(1, 13), "quantile": 0}, dims="month", ) assert_identical(expected, actual) actual = g.quantile(0, dim="time")[:2] expected = xr.DataArray( data=[[0.0, 1], [62.0, 63]], coords={"month": [1, 2], "x": [0, 1], "quantile": 0}, dims=("month", "x"), ) assert_identical(expected, actual) # method keyword array = xr.DataArray(data=[1, 2, 3, 4], coords={"x": [1, 1, 2, 2]}, dims="x") expected = xr.DataArray( data=[1, 3], coords={"x": [1, 2], "quantile": 0.5}, dims="x" ) actual = array.groupby("x").quantile(0.5, method="lower") assert_identical(expected, actual) def test_ds_groupby_quantile() -> None: ds = xr.Dataset( data_vars={"a": ("x", [1, 2, 3, 4, 5, 6])}, coords={"x": [1, 1, 1, 2, 2, 2]} ) # Scalar quantile expected = xr.Dataset( data_vars={"a": ("x", [2, 5])}, coords={"quantile": 0.5, "x": [1, 2]} ) actual = ds.groupby("x").quantile(0.5) assert_identical(expected, actual) # Vector quantile expected = xr.Dataset( data_vars={"a": (("x", "quantile"), [[1, 3], [4, 6]])}, coords={"x": [1, 2], "quantile": [0, 1]}, ) actual = ds.groupby("x").quantile([0, 1]) assert_identical(expected, actual) ds = xr.Dataset( data_vars={"a": ("x", [np.nan, 2, 3, 4, 5, 6])}, coords={"x": [1, 1, 1, 2, 2, 2]}, ) for skipna in (True, False, None): e = [np.nan, 5] if skipna is False else [2.5, 5] expected = xr.Dataset( data_vars={"a": ("x", e)}, coords={"quantile": 0.5, "x": [1, 2]} ) actual = ds.groupby("x").quantile(0.5, skipna=skipna) assert_identical(expected, actual) # Multiple dimensions ds = xr.Dataset( data_vars={ "a": ( ("x", "y"), [[1, 11, 26], [2, 12, 22], [3, 13, 23], [4, 16, 24], [5, 15, 25]], ) }, coords={"x": [1, 1, 1, 2, 2], "y": [0, 0, 1]}, ) actual_x = ds.groupby("x").quantile(0, dim=...) expected_x = xr.Dataset({"a": ("x", [1, 4])}, coords={"x": [1, 2], "quantile": 0}) assert_identical(expected_x, actual_x) actual_y = ds.groupby("y").quantile(0, dim=...) expected_y = xr.Dataset({"a": ("y", [1, 22])}, coords={"y": [0, 1], "quantile": 0}) assert_identical(expected_y, actual_y) actual_xx = ds.groupby("x").quantile(0) expected_xx = xr.Dataset( {"a": (("x", "y"), [[1, 11, 22], [4, 15, 24]])}, coords={"x": [1, 2], "y": [0, 0, 1], "quantile": 0}, ) assert_identical(expected_xx, actual_xx) actual_yy = ds.groupby("y").quantile(0) expected_yy = xr.Dataset( {"a": (("x", "y"), [[1, 26], [2, 22], [3, 23], [4, 24], [5, 25]])}, coords={"x": [1, 1, 1, 2, 2], "y": [0, 1], "quantile": 0}, ).transpose() assert_identical(expected_yy, actual_yy) times = pd.date_range("2000-01-01", periods=365) x = [0, 1] foo = xr.Dataset( {"a": (("time", "x"), np.reshape(np.arange(365 * 2), (365, 2)))}, coords=dict(time=times, x=x), ) g = foo.groupby(foo.time.dt.month) actual = g.quantile(0, dim=...) expected = xr.Dataset( { "a": ( "month", [ 0.0, 62.0, 120.0, 182.0, 242.0, 304.0, 364.0, 426.0, 488.0, 548.0, 610.0, 670.0, ], ) }, coords={"month": np.arange(1, 13), "quantile": 0}, ) assert_identical(expected, actual) actual = g.quantile(0, dim="time").isel(month=slice(None, 2)) expected = xr.Dataset( data_vars={"a": (("month", "x"), [[0.0, 1], [62.0, 63]])}, coords={"month": [1, 2], "x": [0, 1], "quantile": 0}, ) assert_identical(expected, actual) ds = xr.Dataset(data_vars={"a": ("x", [1, 2, 3, 4])}, coords={"x": [1, 1, 2, 2]}) # method keyword expected = xr.Dataset( data_vars={"a": ("x", [1, 3])}, coords={"quantile": 0.5, "x": [1, 2]} ) actual = ds.groupby("x").quantile(0.5, method="lower") assert_identical(expected, actual) @pytest.mark.parametrize("as_dataset", [False, True]) def test_groupby_quantile_interpolation_deprecated(as_dataset: bool) -> None: array = xr.DataArray(data=[1, 2, 3, 4], coords={"x": [1, 1, 2, 2]}, dims="x") arr: xr.DataArray | xr.Dataset arr = array.to_dataset(name="name") if as_dataset else array with pytest.warns( FutureWarning, match="`interpolation` argument to quantile was renamed to `method`", ): actual = arr.quantile(0.5, interpolation="lower") expected = arr.quantile(0.5, method="lower") assert_identical(actual, expected) with warnings.catch_warnings(record=True): with pytest.raises(TypeError, match="interpolation and method keywords"): arr.quantile(0.5, method="lower", interpolation="lower") def test_da_groupby_assign_coords() -> None: actual = xr.DataArray( [[3, 4, 5], [6, 7, 8]], dims=["y", "x"], coords={"y": range(2), "x": range(3)} ) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): actual1 = actual.groupby("x").assign_coords({"y": [-1, -2]}) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): actual2 = actual.groupby("x").assign_coords(y=[-1, -2]) expected = xr.DataArray( [[3, 4, 5], [6, 7, 8]], dims=["y", "x"], coords={"y": [-1, -2], "x": range(3)} ) assert_identical(expected, actual1) assert_identical(expected, actual2) repr_da = xr.DataArray( np.random.randn(10, 20, 6, 24), dims=["x", "y", "z", "t"], coords={ "z": ["a", "b", "c", "a", "b", "c"], "x": [1, 1, 1, 2, 2, 3, 4, 5, 3, 4], "t": xr.date_range("2001-01-01", freq="ME", periods=24, use_cftime=False), "month": ("t", list(range(1, 13)) * 2), }, ) @pytest.mark.parametrize("dim", ["x", "y", "z", "month"]) @pytest.mark.parametrize("obj", [repr_da, repr_da.to_dataset(name="a")]) def test_groupby_repr(obj, dim) -> None: actual = repr(obj.groupby(dim)) expected = f"{obj.__class__.__name__}GroupBy" expected += f", grouped over {dim!r}" expected += f"\n{len(np.unique(obj[dim]))!r} groups with labels " if dim == "x": expected += "1, 2, 3, 4, 5." elif dim == "y": expected += "0, 1, 2, 3, 4, 5, ..., 15, 16, 17, 18, 19." elif dim == "z": expected += "'a', 'b', 'c'." elif dim == "month": expected += "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12." assert actual == expected @pytest.mark.parametrize("obj", [repr_da, repr_da.to_dataset(name="a")]) def test_groupby_repr_datetime(obj) -> None: actual = repr(obj.groupby("t.month")) expected = f"{obj.__class__.__name__}GroupBy" expected += ", grouped over 'month'" expected += f"\n{len(np.unique(obj.t.dt.month))!r} groups with labels " expected += "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12." assert actual == expected @pytest.mark.filterwarnings("ignore:Converting non-nanosecond") @pytest.mark.filterwarnings("ignore:invalid value encountered in divide:RuntimeWarning") def test_groupby_drops_nans() -> None: # GH2383 # nan in 2D data variable (requires stacking) ds = xr.Dataset( { "variable": (("lat", "lon", "time"), np.arange(60.0).reshape((4, 3, 5))), "id": (("lat", "lon"), np.arange(12.0).reshape((4, 3))), }, coords={"lat": np.arange(4), "lon": np.arange(3), "time": np.arange(5)}, ) ds["id"].values[0, 0] = np.nan ds["id"].values[3, 0] = np.nan ds["id"].values[-1, -1] = np.nan grouped = ds.groupby(ds.id) # non reduction operation expected1 = ds.copy() expected1.variable.values[0, 0, :] = np.nan expected1.variable.values[-1, -1, :] = np.nan expected1.variable.values[3, 0, :] = np.nan actual1 = grouped.map(lambda x: x).transpose(*ds.variable.dims) assert_identical(actual1, expected1) # reduction along grouped dimension actual2 = grouped.mean() stacked = ds.stack({"xy": ["lat", "lon"]}) expected2 = ( stacked.variable.where(stacked.id.notnull()) .rename({"xy": "id"}) .to_dataset() .reset_index("id", drop=True) .assign(id=stacked.id.values) .dropna("id") .transpose(*actual2.variable.dims) ) assert_identical(actual2, expected2) # reduction operation along a different dimension actual3 = grouped.mean("time") expected3 = ds.mean("time").where(ds.id.notnull()) assert_identical(actual3, expected3) # NaN in non-dimensional coordinate array = xr.DataArray([1, 2, 3], [("x", [1, 2, 3])]) array["x1"] = ("x", [1, 1, np.nan]) expected4 = xr.DataArray(3, [("x1", [1])]) actual4 = array.groupby("x1").sum() assert_equal(expected4, actual4) # NaT in non-dimensional coordinate array["t"] = ( "x", [ np.datetime64("2001-01-01"), np.datetime64("2001-01-01"), np.datetime64("NaT"), ], ) expected5 = xr.DataArray(3, [("t", [np.datetime64("2001-01-01")])]) actual5 = array.groupby("t").sum() assert_equal(expected5, actual5) # test for repeated coordinate labels array = xr.DataArray([0, 1, 2, 4, 3, 4], [("x", [np.nan, 1, 1, np.nan, 2, np.nan])]) expected6 = xr.DataArray([3, 3], [("x", [1, 2])]) actual6 = array.groupby("x").sum() assert_equal(expected6, actual6) def test_groupby_grouping_errors() -> None: dataset = xr.Dataset({"foo": ("x", [1, 1, 1])}, {"x": [1, 2, 3]}) with pytest.raises( ValueError, match=r"None of the data falls within bins with edges" ): dataset.groupby_bins("x", bins=[0.1, 0.2, 0.3]) with pytest.raises( ValueError, match=r"None of the data falls within bins with edges" ): dataset.to_dataarray().groupby_bins("x", bins=[0.1, 0.2, 0.3]) with pytest.raises(ValueError, match=r"All bin edges are NaN."): dataset.groupby_bins("x", bins=[np.nan, np.nan, np.nan]) with pytest.raises(ValueError, match=r"All bin edges are NaN."): dataset.to_dataarray().groupby_bins("x", bins=[np.nan, np.nan, np.nan]) with pytest.raises(ValueError, match=r"Failed to group data."): dataset.groupby(dataset.foo * np.nan) with pytest.raises(ValueError, match=r"Failed to group data."): dataset.to_dataarray().groupby(dataset.foo * np.nan) def test_groupby_reduce_dimension_error(array) -> None: grouped = array.groupby("y") # assert_identical(array, grouped.mean()) with pytest.raises(ValueError, match=r"cannot reduce over dimensions"): grouped.mean("huh") with pytest.raises(ValueError, match=r"cannot reduce over dimensions"): grouped.mean(("x", "y", "asd")) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert_identical(array.mean("x"), grouped.reduce(np.mean, "x")) assert_allclose(array.mean(["x", "z"]), grouped.reduce(np.mean, ["x", "z"])) grouped = array.groupby("y", squeeze=False) assert_identical(array, grouped.mean()) assert_identical(array.mean("x"), grouped.reduce(np.mean, "x")) assert_allclose(array.mean(["x", "z"]), grouped.reduce(np.mean, ["x", "z"])) def test_groupby_multiple_string_args(array) -> None: with pytest.raises(TypeError): array.groupby("x", "y") def test_groupby_bins_timeseries() -> None: ds = xr.Dataset() ds["time"] = xr.DataArray( pd.date_range("2010-08-01", "2010-08-15", freq="15min"), dims="time" ) ds["val"] = xr.DataArray(np.ones(ds["time"].shape), dims="time") time_bins = pd.date_range(start="2010-08-01", end="2010-08-15", freq="24h") actual = ds.groupby_bins("time", time_bins).sum() expected = xr.DataArray( 96 * np.ones((14,)), dims=["time_bins"], coords={"time_bins": pd.cut(time_bins, time_bins).categories}, ).to_dataset(name="val") assert_identical(actual, expected) def test_groupby_none_group_name() -> None: # GH158 # xarray should not fail if a DataArray's name attribute is None data = np.arange(10) + 10 da = xr.DataArray(data) # da.name = None key = xr.DataArray(np.floor_divide(data, 2)) mean = da.groupby(key).mean() assert "group" in mean.dims def test_groupby_getitem(dataset) -> None: with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert_identical(dataset.sel(x="a"), dataset.groupby("x")["a"]) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert_identical(dataset.sel(z=1), dataset.groupby("z")[1]) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert_identical(dataset.foo.sel(x="a"), dataset.foo.groupby("x")["a"]) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert_identical(dataset.foo.sel(z=1), dataset.foo.groupby("z")[1]) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert_identical(dataset.cat.sel(y=1), dataset.cat.groupby("y")[1]) assert_identical(dataset.sel(x=["a"]), dataset.groupby("x", squeeze=False)["a"]) assert_identical(dataset.sel(z=[1]), dataset.groupby("z", squeeze=False)[1]) assert_identical( dataset.foo.sel(x=["a"]), dataset.foo.groupby("x", squeeze=False)["a"] ) assert_identical(dataset.foo.sel(z=[1]), dataset.foo.groupby("z", squeeze=False)[1]) assert_identical(dataset.cat.sel(y=[1]), dataset.cat.groupby("y", squeeze=False)[1]) with pytest.raises( NotImplementedError, match="Cannot broadcast 1d-only pandas categorical array." ): dataset.groupby("boo", squeeze=False) dataset = dataset.drop_vars(["cat"]) actual = ( dataset.groupby("boo", squeeze=False)["f"].unstack().transpose("x", "y", "z") ) expected = dataset.sel(y=[1], z=[1, 2]).transpose("x", "y", "z") assert_identical(expected, actual) def test_groupby_dataset() -> None: data = Dataset( {"z": (["x", "y"], np.random.randn(3, 5))}, {"x": ("x", list("abc")), "c": ("x", [0, 1, 0]), "y": range(5)}, ) groupby = data.groupby("x", squeeze=False) assert len(groupby) == 3 expected_groups = {"a": slice(0, 1), "b": slice(1, 2), "c": slice(2, 3)} assert groupby.groups == expected_groups expected_items = [ ("a", data.isel(x=[0])), ("b", data.isel(x=[1])), ("c", data.isel(x=[2])), ] for actual1, expected1 in zip(groupby, expected_items): assert actual1[0] == expected1[0] assert_equal(actual1[1], expected1[1]) def identity(x): return x for k in ["x", "c", "y"]: actual2 = data.groupby(k, squeeze=False).map(identity) assert_equal(data, actual2) def test_groupby_dataset_squeeze_None() -> None: """Delete when removing squeeze.""" data = Dataset( {"z": (["x", "y"], np.random.randn(3, 5))}, {"x": ("x", list("abc")), "c": ("x", [0, 1, 0]), "y": range(5)}, ) groupby = data.groupby("x") assert len(groupby) == 3 expected_groups = {"a": 0, "b": 1, "c": 2} with pytest.warns(UserWarning, match="The `squeeze` kwarg"): assert groupby.groups == expected_groups expected_items = [ ("a", data.isel(x=0)), ("b", data.isel(x=1)), ("c", data.isel(x=2)), ] with pytest.warns(UserWarning, match="The `squeeze` kwarg"): for actual1, expected1 in zip(groupby, expected_items): assert actual1[0] == expected1[0] assert_equal(actual1[1], expected1[1]) def identity(x): return x with pytest.warns(UserWarning, match="The `squeeze` kwarg"): for k in ["x", "c"]: actual2 = data.groupby(k).map(identity) assert_equal(data, actual2) def test_groupby_dataset_returns_new_type() -> None: data = Dataset({"z": (["x", "y"], np.random.randn(3, 5))}) actual1 = data.groupby("x", squeeze=False).map(lambda ds: ds["z"]) expected1 = data["z"] assert_identical(expected1, actual1) actual2 = data["z"].groupby("x", squeeze=False).map(lambda x: x.to_dataset()) expected2 = data assert_identical(expected2, actual2) def test_groupby_dataset_iter() -> None: data = create_test_data() for n, (t, sub) in enumerate(list(data.groupby("dim1", squeeze=False))[:3]): assert data["dim1"][n] == t assert_equal(data["var1"][[n]], sub["var1"]) assert_equal(data["var2"][[n]], sub["var2"]) assert_equal(data["var3"][:, [n]], sub["var3"]) def test_groupby_dataset_errors() -> None: data = create_test_data() with pytest.raises(TypeError, match=r"`group` must be"): data.groupby(np.arange(10)) # type: ignore[arg-type,unused-ignore] with pytest.raises(ValueError, match=r"length does not match"): data.groupby(data["dim1"][:3]) with pytest.raises(TypeError, match=r"`group` must be"): data.groupby(data.coords["dim1"].to_index()) def test_groupby_dataset_reduce() -> None: data = Dataset( { "xy": (["x", "y"], np.random.randn(3, 4)), "xonly": ("x", np.random.randn(3)), "yonly": ("y", np.random.randn(4)), "letters": ("y", ["a", "a", "b", "b"]), } ) expected = data.mean("y") expected["yonly"] = expected["yonly"].variable.set_dims({"x": 3}) actual = data.groupby("x").mean(...) assert_allclose(expected, actual) actual = data.groupby("x").mean("y") assert_allclose(expected, actual) letters = data["letters"] expected = Dataset( { "xy": data["xy"].groupby(letters).mean(...), "xonly": (data["xonly"].mean().variable.set_dims({"letters": 2})), "yonly": data["yonly"].groupby(letters).mean(), } ) actual = data.groupby("letters").mean(...) assert_allclose(expected, actual) @pytest.mark.parametrize("squeeze", [True, False]) def test_groupby_dataset_math(squeeze: bool) -> None: def reorder_dims(x): return x.transpose("dim1", "dim2", "dim3", "time") ds = create_test_data() ds["dim1"] = ds["dim1"] grouped = ds.groupby("dim1", squeeze=squeeze) expected = reorder_dims(ds + ds.coords["dim1"]) actual = grouped + ds.coords["dim1"] assert_identical(expected, reorder_dims(actual)) actual = ds.coords["dim1"] + grouped assert_identical(expected, reorder_dims(actual)) ds2 = 2 * ds expected = reorder_dims(ds + ds2) actual = grouped + ds2 assert_identical(expected, reorder_dims(actual)) actual = ds2 + grouped assert_identical(expected, reorder_dims(actual)) def test_groupby_math_more() -> None: ds = create_test_data() grouped = ds.groupby("numbers") zeros = DataArray([0, 0, 0, 0], [("numbers", range(4))]) expected = (ds + Variable("dim3", np.zeros(10))).transpose( "dim3", "dim1", "dim2", "time" ) actual = grouped + zeros assert_equal(expected, actual) actual = zeros + grouped assert_equal(expected, actual) with pytest.raises(ValueError, match=r"incompat.* grouped binary"): grouped + ds with pytest.raises(ValueError, match=r"incompat.* grouped binary"): ds + grouped with pytest.raises(TypeError, match=r"only support binary ops"): grouped + 1 # type: ignore[operator] with pytest.raises(TypeError, match=r"only support binary ops"): grouped + grouped # type: ignore[operator] with pytest.raises(TypeError, match=r"in-place operations"): ds += grouped # type: ignore[arg-type] ds = Dataset( { "x": ("time", np.arange(100)), "time": pd.date_range("2000-01-01", periods=100), } ) with pytest.raises(ValueError, match=r"incompat.* grouped binary"): ds + ds.groupby("time.month") def test_groupby_math_bitshift() -> None: # create new dataset of int's only ds = Dataset( { "x": ("index", np.ones(4, dtype=int)), "y": ("index", np.ones(4, dtype=int) * -1), "level": ("index", [1, 1, 2, 2]), "index": [0, 1, 2, 3], } ) shift = DataArray([1, 2, 1], [("level", [1, 2, 8])]) left_expected = Dataset( { "x": ("index", [2, 2, 4, 4]), "y": ("index", [-2, -2, -4, -4]), "level": ("index", [2, 2, 8, 8]), "index": [0, 1, 2, 3], } ) left_manual = [] for lev, group in ds.groupby("level"): shifter = shift.sel(level=lev) left_manual.append(group << shifter) left_actual = xr.concat(left_manual, dim="index").reset_coords(names="level") assert_equal(left_expected, left_actual) left_actual = (ds.groupby("level") << shift).reset_coords(names="level") assert_equal(left_expected, left_actual) right_expected = Dataset( { "x": ("index", [0, 0, 2, 2]), "y": ("index", [-1, -1, -2, -2]), "level": ("index", [0, 0, 4, 4]), "index": [0, 1, 2, 3], } ) right_manual = [] for lev, group in left_expected.groupby("level"): shifter = shift.sel(level=lev) right_manual.append(group >> shifter) right_actual = xr.concat(right_manual, dim="index").reset_coords(names="level") assert_equal(right_expected, right_actual) right_actual = (left_expected.groupby("level") >> shift).reset_coords(names="level") assert_equal(right_expected, right_actual) @pytest.mark.parametrize("use_flox", [True, False]) def test_groupby_bins_cut_kwargs(use_flox: bool) -> None: da = xr.DataArray(np.arange(12).reshape(6, 2), dims=("x", "y")) x_bins = (0, 2, 4, 6) with xr.set_options(use_flox=use_flox): actual = da.groupby_bins( "x", bins=x_bins, include_lowest=True, right=False, squeeze=False ).mean() expected = xr.DataArray( np.array([[1.0, 2.0], [5.0, 6.0], [9.0, 10.0]]), dims=("x_bins", "y"), coords={ "x_bins": ("x_bins", pd.IntervalIndex.from_breaks(x_bins, closed="left")) }, ) assert_identical(expected, actual) @pytest.mark.parametrize("indexed_coord", [True, False]) def test_groupby_bins_math(indexed_coord) -> None: N = 7 da = DataArray(np.random.random((N, N)), dims=("x", "y")) if indexed_coord: da["x"] = np.arange(N) da["y"] = np.arange(N) g = da.groupby_bins("x", np.arange(0, N + 1, 3)) mean = g.mean() expected = da.isel(x=slice(1, None)) - mean.isel(x_bins=("x", [0, 0, 0, 1, 1, 1])) actual = g - mean assert_identical(expected, actual) def test_groupby_math_nD_group() -> None: N = 40 da = DataArray( np.random.random((N, N)), dims=("x", "y"), coords={ "labels": ( "x", np.repeat(["a", "b", "c", "d", "e", "f", "g", "h"], repeats=N // 8), ), }, ) da["labels2d"] = xr.broadcast(da.labels, da)[0] g = da.groupby("labels2d") mean = g.mean() expected = da - mean.sel(labels2d=da.labels2d) expected["labels"] = expected.labels.broadcast_like(expected.labels2d) actual = g - mean assert_identical(expected, actual) da["num"] = ( "x", np.repeat([1, 2, 3, 4, 5, 6, 7, 8], repeats=N // 8), ) da["num2d"] = xr.broadcast(da.num, da)[0] g = da.groupby_bins("num2d", bins=[0, 4, 6]) mean = g.mean() idxr = np.digitize(da.num2d, bins=(0, 4, 6), right=True)[:30, :] - 1 expanded_mean = mean.drop_vars("num2d_bins").isel(num2d_bins=(("x", "y"), idxr)) expected = da.isel(x=slice(30)) - expanded_mean expected["labels"] = expected.labels.broadcast_like(expected.labels2d) expected["num"] = expected.num.broadcast_like(expected.num2d) expected["num2d_bins"] = (("x", "y"), mean.num2d_bins.data[idxr]) actual = g - mean assert_identical(expected, actual) def test_groupby_dataset_math_virtual() -> None: ds = Dataset({"x": ("t", [1, 2, 3])}, {"t": pd.date_range("20100101", periods=3)}) grouped = ds.groupby("t.day") actual = grouped - grouped.mean(...) expected = Dataset({"x": ("t", [0, 0, 0])}, ds[["t", "t.day"]]) assert_identical(actual, expected) def test_groupby_math_dim_order() -> None: da = DataArray( np.ones((10, 10, 12)), dims=("x", "y", "time"), coords={"time": pd.date_range("2001-01-01", periods=12, freq="6h")}, ) grouped = da.groupby("time.day") result = grouped - grouped.mean() assert result.dims == da.dims def test_groupby_dataset_nan() -> None: # nan should be excluded from groupby ds = Dataset({"foo": ("x", [1, 2, 3, 4])}, {"bar": ("x", [1, 1, 2, np.nan])}) actual = ds.groupby("bar").mean(...) expected = Dataset({"foo": ("bar", [1.5, 3]), "bar": [1, 2]}) assert_identical(actual, expected) def test_groupby_dataset_order() -> None: # groupby should preserve variables order ds = Dataset() for vn in ["a", "b", "c"]: ds[vn] = DataArray(np.arange(10), dims=["t"]) data_vars_ref = list(ds.data_vars.keys()) ds = ds.groupby("t").mean(...) data_vars = list(ds.data_vars.keys()) assert data_vars == data_vars_ref # coords are now at the end of the list, so the test below fails # all_vars = list(ds.variables.keys()) # all_vars_ref = list(ds.variables.keys()) # .assertEqual(all_vars, all_vars_ref) def test_groupby_dataset_fillna() -> None: ds = Dataset({"a": ("x", [np.nan, 1, np.nan, 3])}, {"x": [0, 1, 2, 3]}) expected = Dataset({"a": ("x", range(4))}, {"x": [0, 1, 2, 3]}) for target in [ds, expected]: target.coords["b"] = ("x", [0, 0, 1, 1]) actual = ds.groupby("b").fillna(DataArray([0, 2], dims="b")) assert_identical(expected, actual) actual = ds.groupby("b").fillna(Dataset({"a": ("b", [0, 2])})) assert_identical(expected, actual) # attrs with groupby ds.attrs["attr"] = "ds" ds.a.attrs["attr"] = "da" actual = ds.groupby("b").fillna(Dataset({"a": ("b", [0, 2])})) assert actual.attrs == ds.attrs assert actual.a.name == "a" assert actual.a.attrs == ds.a.attrs def test_groupby_dataset_where() -> None: # groupby ds = Dataset({"a": ("x", range(5))}, {"c": ("x", [0, 0, 1, 1, 1])}) cond = Dataset({"a": ("c", [True, False])}) expected = ds.copy(deep=True) expected["a"].values = np.array([0, 1] + [np.nan] * 3) actual = ds.groupby("c").where(cond) assert_identical(expected, actual) # attrs with groupby ds.attrs["attr"] = "ds" ds.a.attrs["attr"] = "da" actual = ds.groupby("c").where(cond) assert actual.attrs == ds.attrs assert actual.a.name == "a" assert actual.a.attrs == ds.a.attrs def test_groupby_dataset_assign() -> None: ds = Dataset({"a": ("x", range(3))}, {"b": ("x", ["A"] * 2 + ["B"])}) actual = ds.groupby("b").assign(c=lambda ds: 2 * ds.a) expected = ds.merge({"c": ("x", [0, 2, 4])}) assert_identical(actual, expected) actual = ds.groupby("b").assign(c=lambda ds: ds.a.sum()) expected = ds.merge({"c": ("x", [1, 1, 2])}) assert_identical(actual, expected) actual = ds.groupby("b").assign_coords(c=lambda ds: ds.a.sum()) expected = expected.set_coords("c") assert_identical(actual, expected) def test_groupby_dataset_map_dataarray_func() -> None: # regression GH6379 ds = Dataset({"foo": ("x", [1, 2, 3, 4])}, coords={"x": [0, 0, 1, 1]}) actual = ds.groupby("x").map(lambda grp: grp.foo.mean()) expected = DataArray([1.5, 3.5], coords={"x": [0, 1]}, dims="x", name="foo") assert_identical(actual, expected) def test_groupby_dataarray_map_dataset_func() -> None: # regression GH6379 da = DataArray([1, 2, 3, 4], coords={"x": [0, 0, 1, 1]}, dims="x", name="foo") actual = da.groupby("x").map(lambda grp: grp.mean().to_dataset()) expected = xr.Dataset({"foo": ("x", [1.5, 3.5])}, coords={"x": [0, 1]}) assert_identical(actual, expected) @requires_flox @pytest.mark.parametrize("kwargs", [{"method": "map-reduce"}, {"engine": "numpy"}]) def test_groupby_flox_kwargs(kwargs) -> None: ds = Dataset({"a": ("x", range(5))}, {"c": ("x", [0, 0, 1, 1, 1])}) with xr.set_options(use_flox=False): expected = ds.groupby("c").mean() with xr.set_options(use_flox=True): actual = ds.groupby("c").mean(**kwargs) assert_identical(expected, actual) class TestDataArrayGroupBy: @pytest.fixture(autouse=True) def setup(self) -> None: self.attrs = {"attr1": "value1", "attr2": 2929} self.x = np.random.random((10, 20)) self.v = Variable(["x", "y"], self.x) self.va = Variable(["x", "y"], self.x, self.attrs) self.ds = Dataset({"foo": self.v}) self.dv = self.ds["foo"] self.mindex = pd.MultiIndex.from_product( [["a", "b"], [1, 2]], names=("level_1", "level_2") ) self.mda = DataArray([0, 1, 2, 3], coords={"x": self.mindex}, dims="x") self.da = self.dv.copy() self.da.coords["abc"] = ("y", np.array(["a"] * 9 + ["c"] + ["b"] * 10)) self.da.coords["y"] = 20 + 100 * self.da["y"] def test_stack_groupby_unsorted_coord(self) -> None: data = [[0, 1], [2, 3]] data_flat = [0, 1, 2, 3] dims = ["x", "y"] y_vals = [2, 3] arr = xr.DataArray(data, dims=dims, coords={"y": y_vals}) actual1 = arr.stack(z=dims).groupby("z").first() midx1 = pd.MultiIndex.from_product([[0, 1], [2, 3]], names=dims) expected1 = xr.DataArray(data_flat, dims=["z"], coords={"z": midx1}) assert_equal(actual1, expected1) # GH: 3287. Note that y coord values are not in sorted order. arr = xr.DataArray(data, dims=dims, coords={"y": y_vals[::-1]}) actual2 = arr.stack(z=dims).groupby("z").first() midx2 = pd.MultiIndex.from_product([[0, 1], [3, 2]], names=dims) expected2 = xr.DataArray(data_flat, dims=["z"], coords={"z": midx2}) assert_equal(actual2, expected2) def test_groupby_iter(self) -> None: for (act_x, act_dv), (exp_x, exp_ds) in zip( self.dv.groupby("y", squeeze=False), self.ds.groupby("y", squeeze=False) ): assert exp_x == act_x assert_identical(exp_ds["foo"], act_dv) with pytest.warns(UserWarning, match="The `squeeze` kwarg"): for (_, exp_dv), (_, act_dv) in zip( self.dv.groupby("x"), self.dv.groupby("x") ): assert_identical(exp_dv, act_dv) def test_groupby_properties(self) -> None: grouped = self.da.groupby("abc") expected_groups = {"a": range(0, 9), "c": [9], "b": range(10, 20)} assert expected_groups.keys() == grouped.groups.keys() for key in expected_groups: expected_group = expected_groups[key] actual_group = grouped.groups[key] # TODO: array_api doesn't allow slice: assert not isinstance(expected_group, slice) assert not isinstance(actual_group, slice) np.testing.assert_array_equal(expected_group, actual_group) assert 3 == len(grouped) @pytest.mark.parametrize( "by, use_da", [("x", False), ("y", False), ("y", True), ("abc", False)] ) @pytest.mark.parametrize("shortcut", [True, False]) @pytest.mark.parametrize("squeeze", [None, True, False]) def test_groupby_map_identity(self, by, use_da, shortcut, squeeze, recwarn) -> None: expected = self.da if use_da: by = expected.coords[by] def identity(x): return x grouped = expected.groupby(by, squeeze=squeeze) actual = grouped.map(identity, shortcut=shortcut) assert_identical(expected, actual) # abc is not a dim coordinate so no warnings expected! if (by.name if use_da else by) != "abc": assert len(recwarn) == (1 if squeeze in [None, True] else 0) def test_groupby_sum(self) -> None: array = self.da grouped = array.groupby("abc") expected_sum_all = Dataset( { "foo": Variable( ["abc"], np.array( [ self.x[:, :9].sum(), self.x[:, 10:].sum(), self.x[:, 9:10].sum(), ] ).T, ), "abc": Variable(["abc"], np.array(["a", "b", "c"])), } )["foo"] assert_allclose(expected_sum_all, grouped.reduce(np.sum, dim=...)) assert_allclose(expected_sum_all, grouped.sum(...)) expected = DataArray( [ array["y"].values[idx].sum() for idx in [slice(9), slice(10, None), slice(9, 10)] ], [["a", "b", "c"]], ["abc"], ) actual = array["y"].groupby("abc").map(np.sum) assert_allclose(expected, actual) actual = array["y"].groupby("abc").sum(...) assert_allclose(expected, actual) expected_sum_axis1 = Dataset( { "foo": ( ["x", "abc"], np.array( [ self.x[:, :9].sum(1), self.x[:, 10:].sum(1), self.x[:, 9:10].sum(1), ] ).T, ), "abc": Variable(["abc"], np.array(["a", "b", "c"])), } )["foo"] assert_allclose(expected_sum_axis1, grouped.reduce(np.sum, "y")) assert_allclose(expected_sum_axis1, grouped.sum("y")) @pytest.mark.parametrize("method", ["sum", "mean", "median"]) def test_groupby_reductions(self, method) -> None: array = self.da grouped = array.groupby("abc") reduction = getattr(np, method) expected = Dataset( { "foo": Variable( ["x", "abc"], np.array( [ reduction(self.x[:, :9], axis=-1), reduction(self.x[:, 10:], axis=-1), reduction(self.x[:, 9:10], axis=-1), ] ).T, ), "abc": Variable(["abc"], np.array(["a", "b", "c"])), } )["foo"] with xr.set_options(use_flox=False): actual_legacy = getattr(grouped, method)(dim="y") with xr.set_options(use_flox=True): actual_npg = getattr(grouped, method)(dim="y") assert_allclose(expected, actual_legacy) assert_allclose(expected, actual_npg) def test_groupby_count(self) -> None: array = DataArray( [0, 0, np.nan, np.nan, 0, 0], coords={"cat": ("x", ["a", "b", "b", "c", "c", "c"])}, dims="x", ) actual = array.groupby("cat").count() expected = DataArray([1, 1, 2], coords=[("cat", ["a", "b", "c"])]) assert_identical(actual, expected) @pytest.mark.parametrize("shortcut", [True, False]) @pytest.mark.parametrize("keep_attrs", [None, True, False]) def test_groupby_reduce_keep_attrs( self, shortcut: bool, keep_attrs: bool | None ) -> None: array = self.da array.attrs["foo"] = "bar" actual = array.groupby("abc").reduce( np.mean, keep_attrs=keep_attrs, shortcut=shortcut ) with xr.set_options(use_flox=False): expected = array.groupby("abc").mean(keep_attrs=keep_attrs) assert_identical(expected, actual) @pytest.mark.parametrize("keep_attrs", [None, True, False]) def test_groupby_keep_attrs(self, keep_attrs: bool | None) -> None: array = self.da array.attrs["foo"] = "bar" with xr.set_options(use_flox=False): expected = array.groupby("abc").mean(keep_attrs=keep_attrs) with xr.set_options(use_flox=True): actual = array.groupby("abc").mean(keep_attrs=keep_attrs) # values are tested elsewhere, here we just check data # TODO: add check_attrs kwarg to assert_allclose actual.data = expected.data assert_identical(expected, actual) def test_groupby_map_center(self) -> None: def center(x): return x - np.mean(x) array = self.da grouped = array.groupby("abc") expected_ds = array.to_dataset() exp_data = np.hstack( [center(self.x[:, :9]), center(self.x[:, 9:10]), center(self.x[:, 10:])] ) expected_ds["foo"] = (["x", "y"], exp_data) expected_centered = expected_ds["foo"] assert_allclose(expected_centered, grouped.map(center)) def test_groupby_map_ndarray(self) -> None: # regression test for #326 array = self.da grouped = array.groupby("abc") actual = grouped.map(np.asarray) # type: ignore[arg-type] # TODO: Not sure using np.asarray like this makes sense with array api assert_equal(array, actual) def test_groupby_map_changes_metadata(self) -> None: def change_metadata(x): x.coords["x"] = x.coords["x"] * 2 x.attrs["fruit"] = "lemon" return x array = self.da grouped = array.groupby("abc") actual = grouped.map(change_metadata) expected = array.copy() expected = change_metadata(expected) assert_equal(expected, actual) @pytest.mark.parametrize("squeeze", [True, False]) def test_groupby_math_squeeze(self, squeeze: bool) -> None: array = self.da grouped = array.groupby("x", squeeze=squeeze) expected = array + array.coords["x"] actual = grouped + array.coords["x"] assert_identical(expected, actual) actual = array.coords["x"] + grouped assert_identical(expected, actual) ds = array.coords["x"].to_dataset(name="X") expected = array + ds actual = grouped + ds assert_identical(expected, actual) actual = ds + grouped assert_identical(expected, actual) def test_groupby_math(self) -> None: array = self.da grouped = array.groupby("abc") expected_agg = (grouped.mean(...) - np.arange(3)).rename(None) actual = grouped - DataArray(range(3), [("abc", ["a", "b", "c"])]) actual_agg = actual.groupby("abc").mean(...) assert_allclose(expected_agg, actual_agg) with pytest.raises(TypeError, match=r"only support binary ops"): grouped + 1 # type: ignore[type-var] with pytest.raises(TypeError, match=r"only support binary ops"): grouped + grouped # type: ignore[type-var] with pytest.raises(TypeError, match=r"in-place operations"): array += grouped # type: ignore[arg-type] def test_groupby_math_not_aligned(self) -> None: array = DataArray( range(4), {"b": ("x", [0, 0, 1, 1]), "x": [0, 1, 2, 3]}, dims="x" ) other = DataArray([10], coords={"b": [0]}, dims="b") actual = array.groupby("b") + other expected = DataArray([10, 11, np.nan, np.nan], array.coords) assert_identical(expected, actual) # regression test for #7797 other = array.groupby("b").sum() actual = array.sel(x=[0, 1]).groupby("b") - other expected = DataArray([-1, 0], {"b": ("x", [0, 0]), "x": [0, 1]}, dims="x") assert_identical(expected, actual) other = DataArray([10], coords={"c": 123, "b": [0]}, dims="b") actual = array.groupby("b") + other expected = DataArray([10, 11, np.nan, np.nan], array.coords) expected.coords["c"] = (["x"], [123] * 2 + [np.nan] * 2) assert_identical(expected, actual) other_ds = Dataset({"a": ("b", [10])}, {"b": [0]}) actual_ds = array.groupby("b") + other_ds expected_ds = Dataset({"a": ("x", [10, 11, np.nan, np.nan])}, array.coords) assert_identical(expected_ds, actual_ds) def test_groupby_restore_dim_order(self) -> None: array = DataArray( np.random.randn(5, 3), coords={"a": ("x", range(5)), "b": ("y", range(3))}, dims=["x", "y"], ) for by, expected_dims in [ ("x", ("x", "y")), ("y", ("x", "y")), ("a", ("a", "y")), ("b", ("x", "b")), ]: result = array.groupby(by, squeeze=False).map(lambda x: x.squeeze()) assert result.dims == expected_dims def test_groupby_restore_coord_dims(self) -> None: array = DataArray( np.random.randn(5, 3), coords={ "a": ("x", range(5)), "b": ("y", range(3)), "c": (("x", "y"), np.random.randn(5, 3)), }, dims=["x", "y"], ) for by, expected_dims in [ ("x", ("x", "y")), ("y", ("x", "y")), ("a", ("a", "y")), ("b", ("x", "b")), ]: result = array.groupby(by, squeeze=False, restore_coord_dims=True).map( lambda x: x.squeeze() )["c"] assert result.dims == expected_dims def test_groupby_first_and_last(self) -> None: array = DataArray([1, 2, 3, 4, 5], dims="x") by = DataArray(["a"] * 2 + ["b"] * 3, dims="x", name="ab") expected = DataArray([1, 3], [("ab", ["a", "b"])]) actual = array.groupby(by).first() assert_identical(expected, actual) expected = DataArray([2, 5], [("ab", ["a", "b"])]) actual = array.groupby(by).last() assert_identical(expected, actual) array = DataArray(np.random.randn(5, 3), dims=["x", "y"]) expected = DataArray(array[[0, 2]], {"ab": ["a", "b"]}, ["ab", "y"]) actual = array.groupby(by).first() assert_identical(expected, actual) actual = array.groupby("x").first() expected = array # should be a no-op assert_identical(expected, actual) def make_groupby_multidim_example_array(self) -> DataArray: return DataArray( [[[0, 1], [2, 3]], [[5, 10], [15, 20]]], coords={ "lon": (["ny", "nx"], [[30, 40], [40, 50]]), "lat": (["ny", "nx"], [[10, 10], [20, 20]]), }, dims=["time", "ny", "nx"], ) def test_groupby_multidim(self) -> None: array = self.make_groupby_multidim_example_array() for dim, expected_sum in [ ("lon", DataArray([5, 28, 23], coords=[("lon", [30.0, 40.0, 50.0])])), ("lat", DataArray([16, 40], coords=[("lat", [10.0, 20.0])])), ]: actual_sum = array.groupby(dim).sum(...) assert_identical(expected_sum, actual_sum) def test_groupby_multidim_map(self) -> None: array = self.make_groupby_multidim_example_array() actual = array.groupby("lon").map(lambda x: x - x.mean()) expected = DataArray( [[[-2.5, -6.0], [-5.0, -8.5]], [[2.5, 3.0], [8.0, 8.5]]], coords=array.coords, dims=array.dims, ) assert_identical(expected, actual) @pytest.mark.parametrize("use_flox", [True, False]) @pytest.mark.parametrize("coords", [np.arange(4), np.arange(4)[::-1], [2, 0, 3, 1]]) @pytest.mark.parametrize( "cut_kwargs", ( {"labels": None, "include_lowest": True}, {"labels": None, "include_lowest": False}, {"labels": ["a", "b"]}, {"labels": [1.2, 3.5]}, {"labels": ["b", "a"]}, ), ) def test_groupby_bins( self, coords: np.typing.ArrayLike, use_flox: bool, cut_kwargs: dict, ) -> None: array = DataArray( np.arange(4), dims="dim_0", coords={"dim_0": coords}, name="a" ) # the first value should not be part of any group ("right" binning) array[0] = 99 # bins follow conventions for pandas.cut # http://pandas.pydata.org/pandas-docs/stable/generated/pandas.cut.html bins = [0, 1.5, 5] df = array.to_dataframe() df["dim_0_bins"] = pd.cut(array["dim_0"], bins, **cut_kwargs) expected_df = df.groupby("dim_0_bins", observed=True).sum() # TODO: can't convert df with IntervalIndex to Xarray expected = ( expected_df.reset_index(drop=True) .to_xarray() .assign_coords(index=np.array(expected_df.index)) .rename({"index": "dim_0_bins"})["a"] ) with xr.set_options(use_flox=use_flox): actual = array.groupby_bins("dim_0", bins=bins, **cut_kwargs).sum() assert_identical(expected, actual) actual = array.groupby_bins("dim_0", bins=bins, **cut_kwargs).map( lambda x: x.sum() ) assert_identical(expected, actual) # make sure original array dims are unchanged assert len(array.dim_0) == 4 def test_groupby_bins_ellipsis(self) -> None: da = xr.DataArray(np.ones((2, 3, 4))) bins = [-1, 0, 1, 2] with xr.set_options(use_flox=False): actual = da.groupby_bins("dim_0", bins).mean(...) with xr.set_options(use_flox=True): expected = da.groupby_bins("dim_0", bins).mean(...) assert_allclose(actual, expected) @pytest.mark.parametrize("use_flox", [True, False]) def test_groupby_bins_gives_correct_subset(self, use_flox: bool) -> None: # GH7766 rng = np.random.default_rng(42) coords = rng.normal(5, 5, 1000) bins = np.logspace(-4, 1, 10) labels = [ "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", ] # xArray # Make a mock dataarray darr = xr.DataArray(coords, coords=[coords], dims=["coords"]) expected = xr.DataArray( [np.nan, np.nan, 1, 1, 1, 8, 31, 104, 542], dims="coords_bins", coords={"coords_bins": labels}, ) gb = darr.groupby_bins("coords", bins, labels=labels) with xr.set_options(use_flox=use_flox): actual = gb.count() assert_identical(actual, expected) def test_groupby_bins_empty(self) -> None: array = DataArray(np.arange(4), [("x", range(4))]) # one of these bins will be empty bins = [0, 4, 5] bin_coords = pd.cut(array["x"], bins).categories actual = array.groupby_bins("x", bins).sum() expected = DataArray([6, np.nan], dims="x_bins", coords={"x_bins": bin_coords}) assert_identical(expected, actual) # make sure original array is unchanged # (was a problem in earlier versions) assert len(array.x) == 4 def test_groupby_bins_multidim(self) -> None: array = self.make_groupby_multidim_example_array() bins = [0, 15, 20] bin_coords = pd.cut(array["lat"].values.flat, bins).categories expected = DataArray([16, 40], dims="lat_bins", coords={"lat_bins": bin_coords}) actual = array.groupby_bins("lat", bins).map(lambda x: x.sum()) assert_identical(expected, actual) # modify the array coordinates to be non-monotonic after unstacking array["lat"].data = np.array([[10.0, 20.0], [20.0, 10.0]]) expected = DataArray([28, 28], dims="lat_bins", coords={"lat_bins": bin_coords}) actual = array.groupby_bins("lat", bins).map(lambda x: x.sum()) assert_identical(expected, actual) bins = [-2, -1, 0, 1, 2] field = DataArray(np.ones((5, 3)), dims=("x", "y")) by = DataArray( np.array([[-1.5, -1.5, 0.5, 1.5, 1.5] * 3]).reshape(5, 3), dims=("x", "y") ) actual = field.groupby_bins(by, bins=bins).count() bincoord = np.array( [ pd.Interval(left, right, closed="right") for left, right in zip(bins[:-1], bins[1:]) ], dtype=object, ) expected = DataArray( np.array([6, np.nan, 3, 6]), dims="group_bins", coords={"group_bins": bincoord}, ) assert_identical(actual, expected) def test_groupby_bins_sort(self) -> None: data = xr.DataArray( np.arange(100), dims="x", coords={"x": np.linspace(-100, 100, num=100)} ) binned_mean = data.groupby_bins("x", bins=11).mean() assert binned_mean.to_index().is_monotonic_increasing with xr.set_options(use_flox=True): actual = data.groupby_bins("x", bins=11).count() with xr.set_options(use_flox=False): expected = data.groupby_bins("x", bins=11).count() assert_identical(actual, expected) def test_groupby_assign_coords(self) -> None: array = DataArray([1, 2, 3, 4], {"c": ("x", [0, 0, 1, 1])}, dims="x") actual = array.groupby("c").assign_coords(d=lambda a: a.mean()) expected = array.copy() expected.coords["d"] = ("x", [1.5, 1.5, 3.5, 3.5]) assert_identical(actual, expected) def test_groupby_fillna(self) -> None: a = DataArray([np.nan, 1, np.nan, 3], coords={"x": range(4)}, dims="x") fill_value = DataArray([0, 1], dims="y") actual = a.fillna(fill_value) expected = DataArray( [[0, 1], [1, 1], [0, 1], [3, 3]], coords={"x": range(4)}, dims=("x", "y") ) assert_identical(expected, actual) b = DataArray(range(4), coords={"x": range(4)}, dims="x") expected = b.copy() for target in [a, expected]: target.coords["b"] = ("x", [0, 0, 1, 1]) actual = a.groupby("b").fillna(DataArray([0, 2], dims="b")) assert_identical(expected, actual) class TestDataArrayResample: @pytest.mark.parametrize("use_cftime", [True, False]) def test_resample(self, use_cftime: bool) -> None: if use_cftime and not has_cftime: pytest.skip() times = xr.date_range( "2000-01-01", freq="6h", periods=10, use_cftime=use_cftime ) def resample_as_pandas(array, *args, **kwargs): array_ = array.copy(deep=True) if use_cftime: array_["time"] = times.to_datetimeindex() result = DataArray.from_series( array_.to_series().resample(*args, **kwargs).mean() ) if use_cftime: result = result.convert_calendar( calendar="standard", use_cftime=use_cftime ) return result array = DataArray(np.arange(10), [("time", times)]) actual = array.resample(time="24h").mean() expected = resample_as_pandas(array, "24h") assert_identical(expected, actual) actual = array.resample(time="24h").reduce(np.mean) assert_identical(expected, actual) actual = array.resample(time="24h", closed="right").mean() expected = resample_as_pandas(array, "24h", closed="right") assert_identical(expected, actual) with pytest.raises(ValueError, match=r"index must be monotonic"): array[[2, 0, 1]].resample(time="1D") @pytest.mark.parametrize("use_cftime", [True, False]) def test_resample_doctest(self, use_cftime: bool) -> None: # run the doctest example here so we are not surprised if use_cftime and not has_cftime: pytest.skip() da = xr.DataArray( np.array([1, 2, 3, 1, 2, np.nan]), dims="time", coords=dict( time=( "time", xr.date_range( "2001-01-01", freq="ME", periods=6, use_cftime=use_cftime ), ), labels=("time", np.array(["a", "b", "c", "c", "b", "a"])), ), ) actual = da.resample(time="3ME").count() expected = DataArray( [1, 3, 1], dims="time", coords={ "time": xr.date_range( "2001-01-01", freq="3ME", periods=3, use_cftime=use_cftime ) }, ) assert_identical(actual, expected) def test_da_resample_func_args(self) -> None: def func(arg1, arg2, arg3=0.0): return arg1.mean("time") + arg2 + arg3 times = pd.date_range("2000", periods=3, freq="D") da = xr.DataArray([1.0, 1.0, 1.0], coords=[times], dims=["time"]) expected = xr.DataArray([3.0, 3.0, 3.0], coords=[times], dims=["time"]) actual = da.resample(time="D").map(func, args=(1.0,), arg3=1.0) assert_identical(actual, expected) def test_resample_first(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.arange(10), [("time", times)]) # resample to same frequency actual = array.resample(time="6h").first() assert_identical(array, actual) actual = array.resample(time="1D").first() expected = DataArray([0, 4, 8], [("time", times[::4])]) assert_identical(expected, actual) # verify that labels don't use the first value actual = array.resample(time="24h").first() expected = DataArray(array.to_series().resample("24h").first()) assert_identical(expected, actual) # missing values array = array.astype(float) array[:2] = np.nan actual = array.resample(time="1D").first() expected = DataArray([2, 4, 8], [("time", times[::4])]) assert_identical(expected, actual) actual = array.resample(time="1D").first(skipna=False) expected = DataArray([np.nan, 4, 8], [("time", times[::4])]) assert_identical(expected, actual) # regression test for http://stackoverflow.com/questions/33158558/ array = Dataset({"time": times})["time"] actual = array.resample(time="1D").last() expected_times = pd.to_datetime( ["2000-01-01T18", "2000-01-02T18", "2000-01-03T06"], unit="ns" ) expected = DataArray(expected_times, [("time", times[::4])], name="time") assert_identical(expected, actual) def test_resample_bad_resample_dim(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.arange(10), [("__resample_dim__", times)]) with pytest.raises(ValueError, match=r"Proxy resampling dimension"): array.resample(**{"__resample_dim__": "1D"}).first() # type: ignore[arg-type] @requires_scipy def test_resample_drop_nondim_coords(self) -> None: xs = np.arange(6) ys = np.arange(3) times = pd.date_range("2000-01-01", freq="6h", periods=5) data = np.tile(np.arange(5), (6, 3, 1)) xx, yy = np.meshgrid(xs * 5, ys * 2.5) tt = np.arange(len(times), dtype=int) array = DataArray(data, {"time": times, "x": xs, "y": ys}, ("x", "y", "time")) xcoord = DataArray(xx.T, {"x": xs, "y": ys}, ("x", "y")) ycoord = DataArray(yy.T, {"x": xs, "y": ys}, ("x", "y")) tcoord = DataArray(tt, {"time": times}, ("time",)) ds = Dataset({"data": array, "xc": xcoord, "yc": ycoord, "tc": tcoord}) ds = ds.set_coords(["xc", "yc", "tc"]) # Select the data now, with the auxiliary coordinates in place array = ds["data"] # Re-sample actual = array.resample(time="12h", restore_coord_dims=True).mean("time") assert "tc" not in actual.coords # Up-sample - filling actual = array.resample(time="1h", restore_coord_dims=True).ffill() assert "tc" not in actual.coords # Up-sample - interpolation actual = array.resample(time="1h", restore_coord_dims=True).interpolate( "linear" ) assert "tc" not in actual.coords def test_resample_keep_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.ones(10), [("time", times)]) array.attrs["meta"] = "data" result = array.resample(time="1D").mean(keep_attrs=True) expected = DataArray([1, 1, 1], [("time", times[::4])], attrs=array.attrs) assert_identical(result, expected) def test_resample_skipna(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.ones(10), [("time", times)]) array[1] = np.nan result = array.resample(time="1D").mean(skipna=False) expected = DataArray([np.nan, 1, 1], [("time", times[::4])]) assert_identical(result, expected) def test_upsample(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=5) array = DataArray(np.arange(5), [("time", times)]) # Forward-fill actual = array.resample(time="3h").ffill() expected = DataArray(array.to_series().resample("3h").ffill()) assert_identical(expected, actual) # Backward-fill actual = array.resample(time="3h").bfill() expected = DataArray(array.to_series().resample("3h").bfill()) assert_identical(expected, actual) # As frequency actual = array.resample(time="3h").asfreq() expected = DataArray(array.to_series().resample("3h").asfreq()) assert_identical(expected, actual) # Pad actual = array.resample(time="3h").pad() expected = DataArray(array.to_series().resample("3h").ffill()) assert_identical(expected, actual) # Nearest rs = array.resample(time="3h") actual = rs.nearest() new_times = rs.groupers[0].full_index expected = DataArray(array.reindex(time=new_times, method="nearest")) assert_identical(expected, actual) def test_upsample_nd(self) -> None: # Same as before, but now we try on multi-dimensional DataArrays. xs = np.arange(6) ys = np.arange(3) times = pd.date_range("2000-01-01", freq="6h", periods=5) data = np.tile(np.arange(5), (6, 3, 1)) array = DataArray(data, {"time": times, "x": xs, "y": ys}, ("x", "y", "time")) # Forward-fill actual = array.resample(time="3h").ffill() expected_data = np.repeat(data, 2, axis=-1) expected_times = times.to_series().resample("3h").asfreq().index expected_data = expected_data[..., : len(expected_times)] expected = DataArray( expected_data, {"time": expected_times, "x": xs, "y": ys}, ("x", "y", "time"), ) assert_identical(expected, actual) # Backward-fill actual = array.resample(time="3h").ffill() expected_data = np.repeat(np.flipud(data.T).T, 2, axis=-1) expected_data = np.flipud(expected_data.T).T expected_times = times.to_series().resample("3h").asfreq().index expected_data = expected_data[..., : len(expected_times)] expected = DataArray( expected_data, {"time": expected_times, "x": xs, "y": ys}, ("x", "y", "time"), ) assert_identical(expected, actual) # As frequency actual = array.resample(time="3h").asfreq() expected_data = np.repeat(data, 2, axis=-1).astype(float)[..., :-1] expected_data[..., 1::2] = np.nan expected_times = times.to_series().resample("3h").asfreq().index expected = DataArray( expected_data, {"time": expected_times, "x": xs, "y": ys}, ("x", "y", "time"), ) assert_identical(expected, actual) # Pad actual = array.resample(time="3h").pad() expected_data = np.repeat(data, 2, axis=-1) expected_data[..., 1::2] = expected_data[..., ::2] expected_data = expected_data[..., :-1] expected_times = times.to_series().resample("3h").asfreq().index expected = DataArray( expected_data, {"time": expected_times, "x": xs, "y": ys}, ("x", "y", "time"), ) assert_identical(expected, actual) def test_upsample_tolerance(self) -> None: # Test tolerance keyword for upsample methods bfill, pad, nearest times = pd.date_range("2000-01-01", freq="1D", periods=2) times_upsampled = pd.date_range("2000-01-01", freq="6h", periods=5) array = DataArray(np.arange(2), [("time", times)]) # Forward fill actual = array.resample(time="6h").ffill(tolerance="12h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. expected = DataArray([0.0, 0.0, 0.0, np.nan, 1.0], [("time", times_upsampled)]) assert_identical(expected, actual) # Backward fill actual = array.resample(time="6h").bfill(tolerance="12h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. expected = DataArray([0.0, np.nan, 1.0, 1.0, 1.0], [("time", times_upsampled)]) assert_identical(expected, actual) # Nearest actual = array.resample(time="6h").nearest(tolerance="6h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. expected = DataArray([0, 0, np.nan, 1, 1], [("time", times_upsampled)]) assert_identical(expected, actual) @requires_scipy def test_upsample_interpolate(self) -> None: from scipy.interpolate import interp1d xs = np.arange(6) ys = np.arange(3) times = pd.date_range("2000-01-01", freq="6h", periods=5) z = np.arange(5) ** 2 data = np.tile(z, (6, 3, 1)) array = DataArray(data, {"time": times, "x": xs, "y": ys}, ("x", "y", "time")) expected_times = times.to_series().resample("1h").asfreq().index # Split the times into equal sub-intervals to simulate the 6 hour # to 1 hour up-sampling new_times_idx = np.linspace(0, len(times) - 1, len(times) * 5) kinds: list[InterpOptions] = [ "linear", "nearest", "zero", "slinear", "quadratic", "cubic", "polynomial", ] for kind in kinds: kwargs = {} if kind == "polynomial": kwargs["order"] = 1 actual = array.resample(time="1h").interpolate(kind, **kwargs) # using interp1d, polynomial order is to set directly in kind using int f = interp1d( np.arange(len(times)), data, kind=kwargs["order"] if kind == "polynomial" else kind, axis=-1, bounds_error=True, assume_sorted=True, ) expected_data = f(new_times_idx) expected = DataArray( expected_data, {"time": expected_times, "x": xs, "y": ys}, ("x", "y", "time"), ) # Use AllClose because there are some small differences in how # we upsample timeseries versus the integer indexing as I've # done here due to floating point arithmetic assert_allclose(expected, actual, rtol=1e-16) @requires_scipy @pytest.mark.filterwarnings("ignore:Converting non-nanosecond") def test_upsample_interpolate_bug_2197(self) -> None: dates = pd.date_range("2007-02-01", "2007-03-01", freq="D") da = xr.DataArray(np.arange(len(dates)), [("time", dates)]) result = da.resample(time="ME").interpolate("linear") expected_times = np.array( [np.datetime64("2007-02-28"), np.datetime64("2007-03-31")] ) expected = xr.DataArray([27.0, np.nan], [("time", expected_times)]) assert_equal(result, expected) @requires_scipy def test_upsample_interpolate_regression_1605(self) -> None: dates = pd.date_range("2016-01-01", "2016-03-31", freq="1D") expected = xr.DataArray( np.random.random((len(dates), 2, 3)), dims=("time", "x", "y"), coords={"time": dates}, ) actual = expected.resample(time="1D").interpolate("linear") assert_allclose(actual, expected, rtol=1e-16) @requires_dask @requires_scipy @pytest.mark.parametrize("chunked_time", [True, False]) def test_upsample_interpolate_dask(self, chunked_time: bool) -> None: from scipy.interpolate import interp1d xs = np.arange(6) ys = np.arange(3) times = pd.date_range("2000-01-01", freq="6h", periods=5) z = np.arange(5) ** 2 data = np.tile(z, (6, 3, 1)) array = DataArray(data, {"time": times, "x": xs, "y": ys}, ("x", "y", "time")) chunks = {"x": 2, "y": 1} if chunked_time: chunks["time"] = 3 expected_times = times.to_series().resample("1h").asfreq().index # Split the times into equal sub-intervals to simulate the 6 hour # to 1 hour up-sampling new_times_idx = np.linspace(0, len(times) - 1, len(times) * 5) kinds: list[InterpOptions] = [ "linear", "nearest", "zero", "slinear", "quadratic", "cubic", "polynomial", ] for kind in kinds: kwargs = {} if kind == "polynomial": kwargs["order"] = 1 actual = array.chunk(chunks).resample(time="1h").interpolate(kind, **kwargs) actual = actual.compute() # using interp1d, polynomial order is to set directly in kind using int f = interp1d( np.arange(len(times)), data, kind=kwargs["order"] if kind == "polynomial" else kind, axis=-1, bounds_error=True, assume_sorted=True, ) expected_data = f(new_times_idx) expected = DataArray( expected_data, {"time": expected_times, "x": xs, "y": ys}, ("x", "y", "time"), ) # Use AllClose because there are some small differences in how # we upsample timeseries versus the integer indexing as I've # done here due to floating point arithmetic assert_allclose(expected, actual, rtol=1e-16) def test_resample_base(self) -> None: times = pd.date_range("2000-01-01T02:03:01", freq="6h", periods=10) array = DataArray(np.arange(10), [("time", times)]) base = 11 with pytest.warns(FutureWarning, match="the `base` parameter to resample"): actual = array.resample(time="24h", base=base).mean() expected = DataArray( array.to_series().resample("24h", offset=f"{base}h").mean() ) assert_identical(expected, actual) def test_resample_offset(self) -> None: times = pd.date_range("2000-01-01T02:03:01", freq="6h", periods=10) array = DataArray(np.arange(10), [("time", times)]) offset = pd.Timedelta("11h") actual = array.resample(time="24h", offset=offset).mean() expected = DataArray(array.to_series().resample("24h", offset=offset).mean()) assert_identical(expected, actual) def test_resample_origin(self) -> None: times = pd.date_range("2000-01-01T02:03:01", freq="6h", periods=10) array = DataArray(np.arange(10), [("time", times)]) origin = "start" actual = array.resample(time="24h", origin=origin).mean() expected = DataArray(array.to_series().resample("24h", origin=origin).mean()) assert_identical(expected, actual) @pytest.mark.parametrize( "loffset", [ "-12h", datetime.timedelta(hours=-12), pd.Timedelta(hours=-12), pd.DateOffset(hours=-12), ], ) def test_resample_loffset(self, loffset) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.arange(10), [("time", times)]) with pytest.warns(FutureWarning, match="`loffset` parameter"): actual = array.resample(time="24h", loffset=loffset).mean() series = array.to_series().resample("24h").mean() if not isinstance(loffset, pd.DateOffset): loffset = pd.Timedelta(loffset) series.index = series.index + loffset expected = DataArray(series) assert_identical(actual, expected) def test_resample_invalid_loffset(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.arange(10), [("time", times)]) with pytest.warns( FutureWarning, match="Following pandas, the `loffset` parameter" ): with pytest.raises(ValueError, match="`loffset` must be"): array.resample(time="24h", loffset=1).mean() # type: ignore class TestDatasetResample: def test_resample_and_first(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)), "bar": ("time", np.random.randn(10), {"meta": "data"}), "time": times, } ) actual = ds.resample(time="1D").first(keep_attrs=True) expected = ds.isel(time=[0, 4, 8]) assert_identical(expected, actual) # upsampling expected_time = pd.date_range("2000-01-01", freq="3h", periods=19) expected = ds.reindex(time=expected_time) actual = ds.resample(time="3h") for how in ["mean", "sum", "first", "last"]: method = getattr(actual, how) result = method() assert_equal(expected, result) for method in [np.mean]: result = actual.reduce(method) assert_equal(expected, result) def test_resample_min_count(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)), "bar": ("time", np.random.randn(10), {"meta": "data"}), "time": times, } ) # inject nan ds["foo"] = xr.where(ds["foo"] > 2.0, np.nan, ds["foo"]) actual = ds.resample(time="1D").sum(min_count=1) expected = xr.concat( [ ds.isel(time=slice(i * 4, (i + 1) * 4)).sum("time", min_count=1) for i in range(3) ], dim=actual["time"], ) assert_allclose(expected, actual) def test_resample_by_mean_with_keep_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)), "bar": ("time", np.random.randn(10), {"meta": "data"}), "time": times, } ) ds.attrs["dsmeta"] = "dsdata" resampled_ds = ds.resample(time="1D").mean(keep_attrs=True) actual = resampled_ds["bar"].attrs expected = ds["bar"].attrs assert expected == actual actual = resampled_ds.attrs expected = ds.attrs assert expected == actual def test_resample_loffset(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)), "bar": ("time", np.random.randn(10), {"meta": "data"}), "time": times, } ) ds.attrs["dsmeta"] = "dsdata" def test_resample_by_mean_discarding_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)), "bar": ("time", np.random.randn(10), {"meta": "data"}), "time": times, } ) ds.attrs["dsmeta"] = "dsdata" resampled_ds = ds.resample(time="1D").mean(keep_attrs=False) assert resampled_ds["bar"].attrs == {} assert resampled_ds.attrs == {} def test_resample_by_last_discarding_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)), "bar": ("time", np.random.randn(10), {"meta": "data"}), "time": times, } ) ds.attrs["dsmeta"] = "dsdata" resampled_ds = ds.resample(time="1D").last(keep_attrs=False) assert resampled_ds["bar"].attrs == {} assert resampled_ds.attrs == {} @requires_scipy def test_resample_drop_nondim_coords(self) -> None: xs = np.arange(6) ys = np.arange(3) times = pd.date_range("2000-01-01", freq="6h", periods=5) data = np.tile(np.arange(5), (6, 3, 1)) xx, yy = np.meshgrid(xs * 5, ys * 2.5) tt = np.arange(len(times), dtype=int) array = DataArray(data, {"time": times, "x": xs, "y": ys}, ("x", "y", "time")) xcoord = DataArray(xx.T, {"x": xs, "y": ys}, ("x", "y")) ycoord = DataArray(yy.T, {"x": xs, "y": ys}, ("x", "y")) tcoord = DataArray(tt, {"time": times}, ("time",)) ds = Dataset({"data": array, "xc": xcoord, "yc": ycoord, "tc": tcoord}) ds = ds.set_coords(["xc", "yc", "tc"]) # Re-sample actual = ds.resample(time="12h").mean("time") assert "tc" not in actual.coords # Up-sample - filling actual = ds.resample(time="1h").ffill() assert "tc" not in actual.coords # Up-sample - interpolation actual = ds.resample(time="1h").interpolate("linear") assert "tc" not in actual.coords def test_resample_old_api(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { "foo": (["time", "x", "y"], np.random.randn(10, 5, 3)), "bar": ("time", np.random.randn(10), {"meta": "data"}), "time": times, } ) with pytest.raises(TypeError, match=r"resample\(\) no longer supports"): ds.resample("1D", "time") # type: ignore[arg-type] with pytest.raises(TypeError, match=r"resample\(\) no longer supports"): ds.resample("1D", dim="time", how="mean") # type: ignore[arg-type] with pytest.raises(TypeError, match=r"resample\(\) no longer supports"): ds.resample("1D", dim="time") # type: ignore[arg-type] def test_resample_ds_da_are_the_same(self) -> None: time = pd.date_range("2000-01-01", freq="6h", periods=365 * 4) ds = xr.Dataset( { "foo": (("time", "x"), np.random.randn(365 * 4, 5)), "time": time, "x": np.arange(5), } ) assert_allclose( ds.resample(time="ME").mean()["foo"], ds.foo.resample(time="ME").mean() ) def test_ds_resample_apply_func_args(self) -> None: def func(arg1, arg2, arg3=0.0): return arg1.mean("time") + arg2 + arg3 times = pd.date_range("2000", freq="D", periods=3) ds = xr.Dataset({"foo": ("time", [1.0, 1.0, 1.0]), "time": times}) expected = xr.Dataset({"foo": ("time", [3.0, 3.0, 3.0]), "time": times}) actual = ds.resample(time="D").map(func, args=(1.0,), arg3=1.0) assert_identical(expected, actual) def test_groupby_cumsum() -> None: ds = xr.Dataset( {"foo": (("x",), [7, 3, 1, 1, 1, 1, 1])}, coords={"x": [0, 1, 2, 3, 4, 5, 6], "group_id": ("x", [0, 0, 1, 1, 2, 2, 2])}, ) actual = ds.groupby("group_id").cumsum(dim="x") expected = xr.Dataset( { "foo": (("x",), [7, 10, 1, 2, 1, 2, 3]), }, coords={ "x": [0, 1, 2, 3, 4, 5, 6], "group_id": ds.group_id, }, ) # TODO: Remove drop_vars when GH6528 is fixed # when Dataset.cumsum propagates indexes, and the group variable? assert_identical(expected.drop_vars(["x", "group_id"]), actual) actual = ds.foo.groupby("group_id").cumsum(dim="x") expected.coords["group_id"] = ds.group_id expected.coords["x"] = np.arange(7) assert_identical(expected.foo, actual) def test_groupby_cumprod() -> None: ds = xr.Dataset( {"foo": (("x",), [7, 3, 0, 1, 1, 2, 1])}, coords={"x": [0, 1, 2, 3, 4, 5, 6], "group_id": ("x", [0, 0, 1, 1, 2, 2, 2])}, ) actual = ds.groupby("group_id").cumprod(dim="x") expected = xr.Dataset( { "foo": (("x",), [7, 21, 0, 0, 1, 2, 2]), }, coords={ "x": [0, 1, 2, 3, 4, 5, 6], "group_id": ds.group_id, }, ) # TODO: Remove drop_vars when GH6528 is fixed # when Dataset.cumsum propagates indexes, and the group variable? assert_identical(expected.drop_vars(["x", "group_id"]), actual) actual = ds.foo.groupby("group_id").cumprod(dim="x") expected.coords["group_id"] = ds.group_id expected.coords["x"] = np.arange(7) assert_identical(expected.foo, actual) @pytest.mark.parametrize( "method, expected_array", [ ("cumsum", [1.0, 2.0, 5.0, 6.0, 2.0, 2.0]), ("cumprod", [1.0, 2.0, 6.0, 6.0, 2.0, 2.0]), ], ) def test_resample_cumsum(method: str, expected_array: list[float]) -> None: ds = xr.Dataset( {"foo": ("time", [1, 2, 3, 1, 2, np.nan])}, coords={ "time": xr.date_range("01-01-2001", freq="ME", periods=6, use_cftime=False), }, ) actual = getattr(ds.resample(time="3ME"), method)(dim="time") expected = xr.Dataset( {"foo": (("time",), expected_array)}, coords={ "time": xr.date_range("01-01-2001", freq="ME", periods=6, use_cftime=False), }, ) # TODO: Remove drop_vars when GH6528 is fixed # when Dataset.cumsum propagates indexes, and the group variable? assert_identical(expected.drop_vars(["time"]), actual) actual = getattr(ds.foo.resample(time="3ME"), method)(dim="time") expected.coords["time"] = ds.time assert_identical(expected.drop_vars(["time"]).foo, actual) def test_groupby_binary_op_regression() -> None: # regression test for #7797 # monthly timeseries that should return "zero anomalies" everywhere time = xr.date_range("2023-01-01", "2023-12-31", freq="MS") data = np.linspace(-1, 1, 12) x = xr.DataArray(data, coords={"time": time}) clim = xr.DataArray(data, coords={"month": np.arange(1, 13, 1)}) # seems to give the correct result if we use the full x, but not with a slice x_slice = x.sel(time=["2023-04-01"]) # two typical ways of computing anomalies anom_gb = x_slice.groupby("time.month") - clim assert_identical(xr.zeros_like(anom_gb), anom_gb) def test_groupby_multiindex_level() -> None: # GH6836 midx = pd.MultiIndex.from_product([list("abc"), [0, 1]], names=("one", "two")) mda = xr.DataArray(np.random.rand(6, 3), [("x", midx), ("y", range(3))]) groups = mda.groupby("one").groups assert groups == {"a": [0, 1], "b": [2, 3], "c": [4, 5]} @requires_flox @pytest.mark.parametrize("func", ["sum", "prod"]) @pytest.mark.parametrize("skipna", [True, False]) @pytest.mark.parametrize("min_count", [None, 1]) def test_min_count_vs_flox(func: str, min_count: int | None, skipna: bool) -> None: da = DataArray( data=np.array([np.nan, 1, 1, np.nan, 1, 1]), dims="x", coords={"labels": ("x", np.array([1, 2, 3, 1, 2, 3]))}, ) gb = da.groupby("labels") method = operator.methodcaller(func, min_count=min_count, skipna=skipna) with xr.set_options(use_flox=True): actual = method(gb) with xr.set_options(use_flox=False): expected = method(gb) assert_identical(actual, expected) @pytest.mark.parametrize("use_flox", [True, False]) def test_min_count_error(use_flox: bool) -> None: if use_flox and not has_flox: pytest.skip() da = DataArray( data=np.array([np.nan, 1, 1, np.nan, 1, 1]), dims="x", coords={"labels": ("x", np.array([1, 2, 3, 1, 2, 3]))}, ) with xr.set_options(use_flox=use_flox): with pytest.raises(TypeError): da.groupby("labels").mean(min_count=1) @requires_dask def test_groupby_math_auto_chunk() -> None: da = xr.DataArray( [[1, 2, 3], [1, 2, 3], [1, 2, 3]], dims=("y", "x"), coords={"label": ("x", [2, 2, 1])}, ) sub = xr.DataArray( InaccessibleArray(np.array([1, 2])), dims="label", coords={"label": [1, 2]} ) actual = da.chunk(x=1, y=2).groupby("label") - sub assert actual.chunksizes == {"x": (1, 1, 1), "y": (2, 1)} @pytest.mark.parametrize("use_flox", [True, False]) def test_groupby_dim_no_dim_equal(use_flox: bool) -> None: # https://github.com/pydata/xarray/issues/8263 da = DataArray( data=[1, 2, 3, 4], dims="lat", coords={"lat": np.linspace(0, 1.01, 4)} ) with xr.set_options(use_flox=use_flox): actual1 = da.drop_vars("lat").groupby("lat", squeeze=False).sum() actual2 = da.groupby("lat", squeeze=False).sum() assert_identical(actual1, actual2.drop_vars("lat")) @requires_flox def test_default_flox_method() -> None: import flox.xarray da = xr.DataArray([1, 2, 3], dims="x", coords={"label": ("x", [2, 2, 1])}) result = xr.DataArray([3, 3], dims="label", coords={"label": [1, 2]}) with mock.patch("flox.xarray.xarray_reduce", return_value=result) as mocked_reduce: da.groupby("label").sum() kwargs = mocked_reduce.call_args.kwargs if Version(flox.__version__) < Version("0.9.0"): assert kwargs["method"] == "cohorts" else: assert "method" not in kwargs