From 9796fe377f48a373d27eec6f7ee72c50e905e670 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Fri, 20 May 2022 14:44:56 +0200 Subject: [PATCH 01/17] add cftime support to polyval --- .gitignore | 2 ++ xarray/core/computation.py | 13 +++++++++---- xarray/tests/test_computation.py | 12 ++++++++++++ 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 686c7efa701..293e79fe806 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ *.py[cod] __pycache__ +.env +.venv # example caches from Hypothesis .hypothesis/ diff --git a/xarray/core/computation.py b/xarray/core/computation.py index da2db39525a..b90421bfaeb 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1950,13 +1950,18 @@ def _ensure_numeric(data: Dataset | DataArray) -> Dataset | DataArray: from .dataset import Dataset def to_floatable(x: DataArray) -> DataArray: - if x.dtype.kind == "M": - # datetimes + if x.dtype.kind in "MO": + # datetimes (CFIndexes are object type) + offset = ( + np.datetime64("1970-01-01") + if x.dtype.kind == "M" + else type(x.data[0])(1970, 1, 1) + ) return x.copy( data=datetime_to_numeric( x.data, - offset=np.datetime64("1970-01-01"), - datetime_unit="ns", + offset=offset, + datetime_unit="ns" ), ) elif x.dtype.kind == "m": diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index ec8b5a5bc7c..512e2673ea6 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -2018,6 +2018,18 @@ def test_where_attrs() -> None: xr.DataArray([1000.0, 2000.0, 3000.0], dims="x"), id="timedelta", ), + pytest.param( + xr.DataArray( + xr.date_range('1970-01-01', freq='1S', periods=3, use_cftime=True), dims="x" + ), + xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}), + xr.DataArray( + [0, 1e9, 2e9], + dims="x", + coords={"x": xr.date_range('1970-01-01', freq='1S', periods=3, use_cftime=True)}, + ), + id="cftime", + ), ], ) def test_polyval( From ff5ad1e0307125a97b5b045bcf941fd491173c2c Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Fri, 20 May 2022 14:52:34 +0200 Subject: [PATCH 02/17] format changes --- xarray/core/computation.py | 6 +----- xarray/tests/test_computation.py | 9 +++++++-- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index b90421bfaeb..7099d85fe1d 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1958,11 +1958,7 @@ def to_floatable(x: DataArray) -> DataArray: else type(x.data[0])(1970, 1, 1) ) return x.copy( - data=datetime_to_numeric( - x.data, - offset=offset, - datetime_unit="ns" - ), + data=datetime_to_numeric(x.data, offset=offset, datetime_unit="ns"), ) elif x.dtype.kind == "m": # timedeltas diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 512e2673ea6..76304913d5d 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -2020,13 +2020,18 @@ def test_where_attrs() -> None: ), pytest.param( xr.DataArray( - xr.date_range('1970-01-01', freq='1S', periods=3, use_cftime=True), dims="x" + xr.date_range("1970-01-01", freq="1S", periods=3, use_cftime=True), + dims="x", ), xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}), xr.DataArray( [0, 1e9, 2e9], dims="x", - coords={"x": xr.date_range('1970-01-01', freq='1S', periods=3, use_cftime=True)}, + coords={ + "x": xr.date_range( + "1970-01-01", freq="1S", periods=3, use_cftime=True + ) + }, ), id="cftime", ), From e5a9b60c6053fbfef449efce96f694cd04826113 Mon Sep 17 00:00:00 2001 From: dcherian Date: Fri, 20 May 2022 09:59:41 -0600 Subject: [PATCH 03/17] Squashed commit of the following: commit 398f1b660233402fdf7e631f0ee5f863cae564a2 Author: dcherian Date: Fri May 20 08:47:56 2022 -0600 Backward compatibility dask commit bde40e471ed7f9e5ba3a9b4066528195fb467931 Merge: 0783df364 4cae8d0ec Author: dcherian Date: Fri May 20 07:54:48 2022 -0600 Merge branch 'main' into dask-datetime-to-numeric * main: concatenate docs style (#6621) Typing for open_dataset/array/mfdataset and to_netcdf/zarr (#6612) {full,zeros,ones}_like typing (#6611) commit 0783df364716b934ad4ab0ac23570fe3e8f7b942 Merge: 5cff4f1bd 8de706151 Author: dcherian Date: Sun May 15 21:03:50 2022 -0600 Merge branch 'main' into dask-datetime-to-numeric * main: (24 commits) Fix overflow issue in decode_cf_datetime for dtypes <= np.uint32 (#6598) Enable flox in GroupBy and resample (#5734) Add setuptools as dependency in ASV benchmark CI (#6609) change polyval dim ordering (#6601) re-add timedelta support for polyval (#6599) Minor Dataset.map docstr clarification (#6595) New inline_array kwarg for open_dataset (#6566) Fix polyval overloads (#6593) Restore old MultiIndex dropping behaviour (#6592) [docs] add Dataset.assign_coords example (#6336) (#6558) Fix zarr append dtype checks (#6476) Add missing space in exception message (#6590) Doc Link to accessors list in extending-xarray.rst (#6587) Fix Dataset/DataArray.isel with drop=True and scalar DataArray indexes (#6579) Add some warnings about rechunking to the docs (#6569) [pre-commit.ci] pre-commit autoupdate (#6584) terminology.rst: fix link to Unidata's "netcdf_dataset_components" (#6583) Allow string formatting of scalar DataArrays (#5981) Fix mypy issues & reenable in tests (#6581) polyval: Use Horner's algorithm + support chunked inputs (#6548) ... commit 5cff4f1bd40cd107c76331c0bd54708ae538f2b0 Merge: dfe200d15 6144c61a7 Author: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> Date: Sun May 1 15:16:33 2022 -0700 Merge branch 'main' into dask-datetime-to-numeric commit dfe200d1502a62ca105b301053a4254c9286f20e Author: dcherian Date: Sun May 1 11:04:03 2022 -0600 Minor cleanup commit 35ed378f201ecf7c8cd80a7a6180583296781ed7 Author: dcherian Date: Sun May 1 10:57:36 2022 -0600 Support dask arrays in datetime_to_numeric --- xarray/core/duck_array_ops.py | 22 +++++++++++-- xarray/tests/test_duck_array_ops.py | 49 +++++++++++++++++++++++------ 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index 253a68b7205..033c238b959 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -431,7 +431,14 @@ def datetime_to_numeric(array, offset=None, datetime_unit=None, dtype=float): # Compute timedelta object. # For np.datetime64, this can silently yield garbage due to overflow. # One option is to enforce 1970-01-01 as the universal offset. - array = array - offset + + # This map_blocks call is for backwards compatibility. + # dask == 2021.04.1 does not support subtracting object arrays + # which is required for cftime + if is_duck_dask_array(array): + array = array.map_blocks(lambda a, b: a - b, offset) + else: + array = array - offset # Scalar is converted to 0d-array if not hasattr(array, "dtype"): @@ -517,10 +524,19 @@ def pd_timedelta_to_float(value, datetime_unit): return np_timedelta64_to_float(value, datetime_unit) +def _timedelta_to_seconds(array): + return np.reshape([a.total_seconds() for a in array.ravel()], array.shape) * 1e6 + + def py_timedelta_to_float(array, datetime_unit): """Convert a timedelta object to a float, possibly at a loss of resolution.""" - array = np.asarray(array) - array = np.reshape([a.total_seconds() for a in array.ravel()], array.shape) * 1e6 + array = asarray(array) + if is_duck_dask_array(array): + array = array.map_blocks( + _timedelta_to_seconds, meta=np.array([], dtype=np.float64) + ) + else: + array = _timedelta_to_seconds(array) conversion_factor = np.timedelta64(1, "us") / np.timedelta64(1, datetime_unit) return conversion_factor * array diff --git a/xarray/tests/test_duck_array_ops.py b/xarray/tests/test_duck_array_ops.py index c329bc50c56..392f1b91914 100644 --- a/xarray/tests/test_duck_array_ops.py +++ b/xarray/tests/test_duck_array_ops.py @@ -675,39 +675,68 @@ def test_multiple_dims(dtype, dask, skipna, func): assert_allclose(actual, expected) -def test_datetime_to_numeric_datetime64(): +@pytest.mark.parametrize("dask", [True, False]) +def test_datetime_to_numeric_datetime64(dask): + if dask and not has_dask: + pytest.skip("requires dask") + times = pd.date_range("2000", periods=5, freq="7D").values - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h") + if dask: + import dask.array + + times = dask.array.from_array(times, chunks=-1) + + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h") expected = 24 * np.arange(0, 35, 7) np.testing.assert_array_equal(result, expected) offset = times[1] - result = duck_array_ops.datetime_to_numeric(times, offset=offset, datetime_unit="h") + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, offset=offset, datetime_unit="h" + ) expected = 24 * np.arange(-7, 28, 7) np.testing.assert_array_equal(result, expected) dtype = np.float32 - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=dtype) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, datetime_unit="h", dtype=dtype + ) expected = 24 * np.arange(0, 35, 7).astype(dtype) np.testing.assert_array_equal(result, expected) @requires_cftime -def test_datetime_to_numeric_cftime(): +@pytest.mark.parametrize("dask", [True, False]) +def test_datetime_to_numeric_cftime(dask): + if dask and not has_dask: + pytest.skip("requires dask") + times = cftime_range("2000", periods=5, freq="7D", calendar="standard").values - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=int) + if dask: + import dask.array + + times = dask.array.from_array(times, chunks=-1) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=int) expected = 24 * np.arange(0, 35, 7) np.testing.assert_array_equal(result, expected) offset = times[1] - result = duck_array_ops.datetime_to_numeric( - times, offset=offset, datetime_unit="h", dtype=int - ) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, offset=offset, datetime_unit="h", dtype=int + ) expected = 24 * np.arange(-7, 28, 7) np.testing.assert_array_equal(result, expected) dtype = np.float32 - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=dtype) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, datetime_unit="h", dtype=dtype + ) expected = 24 * np.arange(0, 35, 7).astype(dtype) np.testing.assert_array_equal(result, expected) From 9c8f6e286c53279d630be28fd5507f356109af07 Mon Sep 17 00:00:00 2001 From: dcherian Date: Fri, 20 May 2022 10:01:15 -0600 Subject: [PATCH 04/17] Fix dask? --- xarray/core/computation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 7099d85fe1d..cb638d3da3b 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1955,7 +1955,7 @@ def to_floatable(x: DataArray) -> DataArray: offset = ( np.datetime64("1970-01-01") if x.dtype.kind == "M" - else type(x.data[0])(1970, 1, 1) + else type(x.data[0].compute())(1970, 1, 1) ) return x.copy( data=datetime_to_numeric(x.data, offset=offset, datetime_unit="ns"), From de8e16293db853ce605fcb53b14f1dacd63fb439 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Sat, 21 May 2022 22:31:50 +0200 Subject: [PATCH 05/17] attempt to fix cftime in polyval --- xarray/core/computation.py | 22 +++++++++++++++++++--- xarray/tests/test_computation.py | 26 ++++++++++++++++++++------ 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index cb638d3da3b..4e30e930be7 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1883,6 +1883,15 @@ def polyval(coord: Dataset, coeffs: Dataset, degree_dim: Hashable) -> Dataset: ... +@overload +def polyval( + coord: Dataset | DataArray, + coeffs: Dataset | DataArray, + degree_dim: Hashable = "degree", +) -> Dataset | DataArray: + ... + + def polyval( coord: Dataset | DataArray, coeffs: Dataset | DataArray, @@ -1949,13 +1958,20 @@ def _ensure_numeric(data: Dataset | DataArray) -> Dataset | DataArray: """ from .dataset import Dataset + def _cfoffset(x: DataArray) -> Any: + try: + # dask arrays require compute + cls = type(x.data[0].compute()) + except Exception: + cls = type(x.data[0]) + + return cls(1970, 1, 1) + def to_floatable(x: DataArray) -> DataArray: if x.dtype.kind in "MO": # datetimes (CFIndexes are object type) offset = ( - np.datetime64("1970-01-01") - if x.dtype.kind == "M" - else type(x.data[0].compute())(1970, 1, 1) + np.datetime64("1970-01-01") if x.dtype.kind == "M" else _cfoffset(x) ) return x.copy( data=datetime_to_numeric(x.data, offset=offset, datetime_unit="ns"), diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 76304913d5d..2baf62b1509 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -1936,14 +1936,17 @@ def test_where_attrs() -> None: assert actual.attrs == {} -@pytest.mark.parametrize("use_dask", [False, True]) @pytest.mark.parametrize( - ["x", "coeffs", "expected"], + "use_dask", [pytest.param(False, "nodask"), pytest.param(True, "dask")] +) +@pytest.mark.parametrize( + ["x", "coeffs", "expected", "max_computes"], [ pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([2, 3, 4], dims="degree", coords={"degree": [0, 1, 2]}), xr.DataArray([9, 2 + 6 + 16, 2 + 9 + 36], dims="x"), + 0, id="simple", ), pytest.param( @@ -1952,6 +1955,7 @@ def test_where_attrs() -> None: [[0, 1], [0, 1]], dims=("y", "degree"), coords={"degree": [0, 1]} ), xr.DataArray([[1, 1], [2, 2], [3, 3]], dims=("x", "y")), + 0, id="broadcast-x", ), pytest.param( @@ -1962,18 +1966,21 @@ def test_where_attrs() -> None: coords={"degree": [0, 1]}, ), xr.DataArray([1, 1, 1 + 3], dims="x"), + 0, id="shared-dim", ), pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([1, 0, 0], dims="degree", coords={"degree": [2, 1, 0]}), - xr.DataArray([1, 2**2, 3**2], dims="x"), + xr.DataArray([1, 2 ** 2, 3 ** 2], dims="x"), + 0, id="reordered-index", ), pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([5], dims="degree", coords={"degree": [3]}), - xr.DataArray([5, 5 * 2**3, 5 * 3**3], dims="x"), + xr.DataArray([5, 5 * 2 ** 3, 5 * 3 ** 3], dims="x"), + 0, id="sparse-index", ), pytest.param( @@ -1983,12 +1990,14 @@ def test_where_attrs() -> None: coords={"degree": [0, 1]}, ), xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [1, 1, 1])}), + 0, id="array-dataset", ), pytest.param( xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [2, 3, 4])}), xr.DataArray([1, 1], dims="degree", coords={"degree": [0, 1]}), xr.Dataset({"a": ("x", [2, 3, 4]), "b": ("x", [3, 4, 5])}), + 0, id="dataset-array", ), pytest.param( @@ -1998,6 +2007,7 @@ def test_where_attrs() -> None: coords={"degree": [0, 1]}, ), xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("y", [3, 4, 5])}), + 0, id="dataset-dataset", ), pytest.param( @@ -2008,6 +2018,7 @@ def test_where_attrs() -> None: dims="x", coords={"x": pd.date_range("1970-01-01", freq="s", periods=3)}, ), + 0, id="datetime", ), pytest.param( @@ -2016,6 +2027,7 @@ def test_where_attrs() -> None: ), xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}), xr.DataArray([1000.0, 2000.0, 3000.0], dims="x"), + 0, id="timedelta", ), pytest.param( @@ -2033,6 +2045,7 @@ def test_where_attrs() -> None: ) }, ), + 1, id="cftime", ), ], @@ -2042,14 +2055,15 @@ def test_polyval( x: xr.DataArray | xr.Dataset, coeffs: xr.DataArray | xr.Dataset, expected: xr.DataArray | xr.Dataset, + max_computes: int, ) -> None: if use_dask: if not has_dask: pytest.skip("requires dask") coeffs = coeffs.chunk({"degree": 2}) x = x.chunk({"x": 2}) - with raise_if_dask_computes(): - actual = xr.polyval(coord=x, coeffs=coeffs) # type: ignore + with raise_if_dask_computes(max_computes): + actual = xr.polyval(coord=x, coeffs=coeffs) xr.testing.assert_allclose(actual, expected) From 48b395604299d68df3f7d7461f07cded927e8437 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Sat, 21 May 2022 22:32:48 +0200 Subject: [PATCH 06/17] fix overload --- xarray/core/computation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 4e30e930be7..ded3abe7dd5 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1887,7 +1887,7 @@ def polyval(coord: Dataset, coeffs: Dataset, degree_dim: Hashable) -> Dataset: def polyval( coord: Dataset | DataArray, coeffs: Dataset | DataArray, - degree_dim: Hashable = "degree", + degree_dim: Hashable, ) -> Dataset | DataArray: ... From 9b3828bf627fdfda3abff20f5ffd6829861df3d3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 21 May 2022 20:33:48 +0000 Subject: [PATCH 07/17] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_computation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 2baf62b1509..693f8c249d6 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -1972,14 +1972,14 @@ def test_where_attrs() -> None: pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([1, 0, 0], dims="degree", coords={"degree": [2, 1, 0]}), - xr.DataArray([1, 2 ** 2, 3 ** 2], dims="x"), + xr.DataArray([1, 2**2, 3**2], dims="x"), 0, id="reordered-index", ), pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([5], dims="degree", coords={"degree": [3]}), - xr.DataArray([5, 5 * 2 ** 3, 5 * 3 ** 3], dims="x"), + xr.DataArray([5, 5 * 2**3, 5 * 3**3], dims="x"), 0, id="sparse-index", ), From 212e43160a50b5b1cd5c6d80df786817aeeab124 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Sat, 21 May 2022 22:46:47 +0200 Subject: [PATCH 08/17] fix test discovery typo --- xarray/tests/test_computation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 693f8c249d6..7a2cf38cbbe 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -1937,7 +1937,7 @@ def test_where_attrs() -> None: @pytest.mark.parametrize( - "use_dask", [pytest.param(False, "nodask"), pytest.param(True, "dask")] + "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] ) @pytest.mark.parametrize( ["x", "coeffs", "expected", "max_computes"], From 0e23685016ea26328f91518168a6599d50f3bbdc Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Sat, 21 May 2022 22:57:56 +0200 Subject: [PATCH 09/17] specify exception --- xarray/core/computation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index ded3abe7dd5..3b1492ecdd1 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1962,7 +1962,7 @@ def _cfoffset(x: DataArray) -> Any: try: # dask arrays require compute cls = type(x.data[0].compute()) - except Exception: + except AttributeError: cls = type(x.data[0]) return cls(1970, 1, 1) From c75d3f57b3092720ccab58cb91ed9d34f0d95f77 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Sat, 21 May 2022 22:57:59 +0200 Subject: [PATCH 10/17] Revert "fix overload" This reverts commit 48b395604299d68df3f7d7461f07cded927e8437. --- xarray/core/computation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 3b1492ecdd1..e0204c33cfa 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1887,7 +1887,7 @@ def polyval(coord: Dataset, coeffs: Dataset, degree_dim: Hashable) -> Dataset: def polyval( coord: Dataset | DataArray, coeffs: Dataset | DataArray, - degree_dim: Hashable, + degree_dim: Hashable = "degree", ) -> Dataset | DataArray: ... From d4aec0e2aa3cbe898332ada473b2aedaf1be526b Mon Sep 17 00:00:00 2001 From: dcherian Date: Sun, 22 May 2022 07:58:46 -0600 Subject: [PATCH 11/17] dask+cftime fixes --- xarray/core/computation.py | 10 ++++++---- xarray/core/duck_array_ops.py | 4 ++-- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index e0204c33cfa..041da188bbe 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1959,12 +1959,14 @@ def _ensure_numeric(data: Dataset | DataArray) -> Dataset | DataArray: from .dataset import Dataset def _cfoffset(x: DataArray) -> Any: - try: - # dask arrays require compute + if is_duck_dask_array(x.data): + # cftime dask arrays require compute cls = type(x.data[0].compute()) - except AttributeError: + if isinstance(cls, np.ndarray): + # we do not get a scalar back on dask == 2021.04.1 + cls = cls.item() + else: cls = type(x.data[0]) - return cls(1970, 1, 1) def to_floatable(x: DataArray) -> DataArray: diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index 033c238b959..e5a659a9ec9 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -435,8 +435,8 @@ def datetime_to_numeric(array, offset=None, datetime_unit=None, dtype=float): # This map_blocks call is for backwards compatibility. # dask == 2021.04.1 does not support subtracting object arrays # which is required for cftime - if is_duck_dask_array(array): - array = array.map_blocks(lambda a, b: a - b, offset) + if is_duck_dask_array(array) and np.issubdtype(array.dtype, np.object): + array = array.map_blocks(lambda a, b: a - b, offset, meta=array._meta) else: array = array - offset From fb29d819bc44a139e7c99abf768136f6892625ea Mon Sep 17 00:00:00 2001 From: dcherian Date: Sun, 22 May 2022 16:11:14 -0600 Subject: [PATCH 12/17] moar fix --- xarray/core/computation.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 041da188bbe..3ac13aa0020 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -30,7 +30,7 @@ from .merge import merge_attrs, merge_coordinates_without_align from .options import OPTIONS, _get_keep_attrs from .pycompat import is_duck_dask_array -from .utils import is_dict_like +from .utils import is_dict_like, is_scalar from .variable import Variable if TYPE_CHECKING: @@ -1959,15 +1959,11 @@ def _ensure_numeric(data: Dataset | DataArray) -> Dataset | DataArray: from .dataset import Dataset def _cfoffset(x: DataArray) -> Any: - if is_duck_dask_array(x.data): - # cftime dask arrays require compute - cls = type(x.data[0].compute()) - if isinstance(cls, np.ndarray): - # we do not get a scalar back on dask == 2021.04.1 - cls = cls.item() - else: - cls = type(x.data[0]) - return cls(1970, 1, 1) + scalar = x.compute().data[0] + if not is_scalar(scalar): + # we do not get a scalar back on dask == 2021.04.1 + scalar = scalar.item() + return type(scalar)(1970, 1, 1) def to_floatable(x: DataArray) -> DataArray: if x.dtype.kind in "MO": From 8c7b3f0ff1e004632ce224097f4c8b27f217895f Mon Sep 17 00:00:00 2001 From: dcherian Date: Fri, 27 May 2022 19:58:30 -0600 Subject: [PATCH 13/17] SKip cftime tests if not present --- xarray/tests/test_computation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 7a2cf38cbbe..cc27771e521 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -26,7 +26,7 @@ from xarray.core.pycompat import dask_version from xarray.core.types import T_Xarray -from . import has_dask, raise_if_dask_computes, requires_dask +from . import has_cftime, has_dask, raise_if_dask_computes, requires_dask def assert_identical(a, b): @@ -2047,6 +2047,7 @@ def test_where_attrs() -> None: ), 1, id="cftime", + marks=pytest.mark.skipif(not has_cftime, reason="raises"), ), ], ) From 76263687e95207e6c26192dca93dde89f6d2d25f Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Sat, 28 May 2022 09:15:48 +0200 Subject: [PATCH 14/17] move polyval_cftime test --- xarray/tests/test_computation.py | 74 ++++++++++++++++++-------------- 1 file changed, 41 insertions(+), 33 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index cc27771e521..8aa6c0ecaa0 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -26,7 +26,13 @@ from xarray.core.pycompat import dask_version from xarray.core.types import T_Xarray -from . import has_cftime, has_dask, raise_if_dask_computes, requires_dask +from . import ( + has_cftime, + has_dask, + raise_if_dask_computes, + requires_cftime, + requires_dask, +) def assert_identical(a, b): @@ -1940,13 +1946,12 @@ def test_where_attrs() -> None: "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] ) @pytest.mark.parametrize( - ["x", "coeffs", "expected", "max_computes"], + ["x", "coeffs", "expected"], [ pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([2, 3, 4], dims="degree", coords={"degree": [0, 1, 2]}), xr.DataArray([9, 2 + 6 + 16, 2 + 9 + 36], dims="x"), - 0, id="simple", ), pytest.param( @@ -1955,7 +1960,6 @@ def test_where_attrs() -> None: [[0, 1], [0, 1]], dims=("y", "degree"), coords={"degree": [0, 1]} ), xr.DataArray([[1, 1], [2, 2], [3, 3]], dims=("x", "y")), - 0, id="broadcast-x", ), pytest.param( @@ -1966,21 +1970,18 @@ def test_where_attrs() -> None: coords={"degree": [0, 1]}, ), xr.DataArray([1, 1, 1 + 3], dims="x"), - 0, id="shared-dim", ), pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([1, 0, 0], dims="degree", coords={"degree": [2, 1, 0]}), xr.DataArray([1, 2**2, 3**2], dims="x"), - 0, id="reordered-index", ), pytest.param( xr.DataArray([1, 2, 3], dims="x"), xr.DataArray([5], dims="degree", coords={"degree": [3]}), xr.DataArray([5, 5 * 2**3, 5 * 3**3], dims="x"), - 0, id="sparse-index", ), pytest.param( @@ -1990,14 +1991,12 @@ def test_where_attrs() -> None: coords={"degree": [0, 1]}, ), xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [1, 1, 1])}), - 0, id="array-dataset", ), pytest.param( xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [2, 3, 4])}), xr.DataArray([1, 1], dims="degree", coords={"degree": [0, 1]}), xr.Dataset({"a": ("x", [2, 3, 4]), "b": ("x", [3, 4, 5])}), - 0, id="dataset-array", ), pytest.param( @@ -2007,7 +2006,6 @@ def test_where_attrs() -> None: coords={"degree": [0, 1]}, ), xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("y", [3, 4, 5])}), - 0, id="dataset-dataset", ), pytest.param( @@ -2018,7 +2016,6 @@ def test_where_attrs() -> None: dims="x", coords={"x": pd.date_range("1970-01-01", freq="s", periods=3)}, ), - 0, id="datetime", ), pytest.param( @@ -2027,28 +2024,8 @@ def test_where_attrs() -> None: ), xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}), xr.DataArray([1000.0, 2000.0, 3000.0], dims="x"), - 0, id="timedelta", ), - pytest.param( - xr.DataArray( - xr.date_range("1970-01-01", freq="1S", periods=3, use_cftime=True), - dims="x", - ), - xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}), - xr.DataArray( - [0, 1e9, 2e9], - dims="x", - coords={ - "x": xr.date_range( - "1970-01-01", freq="1S", periods=3, use_cftime=True - ) - }, - ), - 1, - id="cftime", - marks=pytest.mark.skipif(not has_cftime, reason="raises"), - ), ], ) def test_polyval( @@ -2056,15 +2033,46 @@ def test_polyval( x: xr.DataArray | xr.Dataset, coeffs: xr.DataArray | xr.Dataset, expected: xr.DataArray | xr.Dataset, - max_computes: int, ) -> None: if use_dask: if not has_dask: pytest.skip("requires dask") coeffs = coeffs.chunk({"degree": 2}) x = x.chunk({"x": 2}) - with raise_if_dask_computes(max_computes): + + with raise_if_dask_computes(): actual = xr.polyval(coord=x, coeffs=coeffs) + + xr.testing.assert_allclose(actual, expected) + + +@requires_cftime +@pytest.mark.parametrize( + "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] +) +def test_polyval_cftime(use_dask: bool) -> None: + x = xr.DataArray( + xr.date_range("1970-01-01", freq="1S", periods=3, use_cftime=True), + dims="x", + ) + coeffs = xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}) + + if use_dask: + if not has_dask: + pytest.skip("requires dask") + coeffs = coeffs.chunk({"degree": 2}) + x = x.chunk({"x": 2}) + + with raise_if_dask_computes(max_computes=1): + actual = xr.polyval(coord=x, coeffs=coeffs) + + expected = xr.DataArray( + [0, 1e9, 2e9], + dims="x", + coords={ + "x": xr.date_range("1970-01-01", freq="1S", periods=3, use_cftime=True) + }, + ) xr.testing.assert_allclose(actual, expected) From f7314efe137a875879fcac2a0c5bbd0324f1340c Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Mon, 30 May 2022 17:33:42 +0200 Subject: [PATCH 15/17] add test with different date --- xarray/tests/test_computation.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 8aa6c0ecaa0..6ef87ba6cb1 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -2050,9 +2050,12 @@ def test_polyval( @pytest.mark.parametrize( "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] ) -def test_polyval_cftime(use_dask: bool) -> None: +@pytest.mark.parametrize("date", ["1970-01-01", "0753-04-21"]) +def test_polyval_cftime(use_dask: bool, date: str) -> None: + import cftime + x = xr.DataArray( - xr.date_range("1970-01-01", freq="1S", periods=3, use_cftime=True), + xr.date_range(date, freq="1S", periods=3, use_cftime=True), dims="x", ) coeffs = xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}) @@ -2066,12 +2069,15 @@ def test_polyval_cftime(use_dask: bool) -> None: with raise_if_dask_computes(max_computes=1): actual = xr.polyval(coord=x, coeffs=coeffs) - expected = xr.DataArray( - [0, 1e9, 2e9], - dims="x", - coords={ - "x": xr.date_range("1970-01-01", freq="1S", periods=3, use_cftime=True) - }, + t0 = xr.date_range(date, periods=1)[0] + offset = (t0 - cftime.DatetimeGregorian(1970, 1, 1)).total_seconds() * 1e9 + expected = ( + xr.DataArray( + [0, 1e9, 2e9], + dims="x", + coords={"x": xr.date_range(date, freq="1S", periods=3, use_cftime=True)}, + ) + + offset ) xr.testing.assert_allclose(actual, expected) From 4ac5327bb487a9e451a6e9c0a51b312e41b03804 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Mon, 30 May 2022 17:35:17 +0200 Subject: [PATCH 16/17] remove deprecation warning --- xarray/core/duck_array_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index e5a659a9ec9..5f09abdbcfd 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -435,7 +435,7 @@ def datetime_to_numeric(array, offset=None, datetime_unit=None, dtype=float): # This map_blocks call is for backwards compatibility. # dask == 2021.04.1 does not support subtracting object arrays # which is required for cftime - if is_duck_dask_array(array) and np.issubdtype(array.dtype, np.object): + if is_duck_dask_array(array) and np.issubdtype(array.dtype, object): array = array.map_blocks(lambda a, b: a - b, offset, meta=array._meta) else: array = array - offset From ae035923270f6ee76eb678156d16cfab922f1d6e Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Mon, 30 May 2022 18:34:20 +0200 Subject: [PATCH 17/17] add polyval integration test --- xarray/tests/test_computation.py | 45 ++++++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 6ef87ba6cb1..1eaa772206e 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -2082,8 +2082,8 @@ def test_polyval_cftime(use_dask: bool, date: str) -> None: xr.testing.assert_allclose(actual, expected) -def test_polyval_degree_dim_checks(): - x = (xr.DataArray([1, 2, 3], dims="x"),) +def test_polyval_degree_dim_checks() -> None: + x = xr.DataArray([1, 2, 3], dims="x") coeffs = xr.DataArray([2, 3, 4], dims="degree", coords={"degree": [0, 1, 2]}) with pytest.raises(ValueError): xr.polyval(x, coeffs.drop_vars("degree")) @@ -2091,6 +2091,47 @@ def test_polyval_degree_dim_checks(): xr.polyval(x, coeffs.assign_coords(degree=coeffs.degree.astype(float))) +@pytest.mark.parametrize( + "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] +) +@pytest.mark.parametrize( + "x", + [ + pytest.param(xr.DataArray([0, 1, 2], dims="x"), id="simple"), + pytest.param( + xr.DataArray(pd.date_range("1970-01-01", freq="ns", periods=3), dims="x"), + id="datetime", + ), + pytest.param( + xr.DataArray(np.array([0, 1, 2], dtype="timedelta64[ns]"), dims="x"), + id="timedelta", + ), + ], +) +@pytest.mark.parametrize( + "y", + [ + pytest.param(xr.DataArray([1, 6, 17], dims="x"), id="1D"), + pytest.param( + xr.DataArray([[1, 6, 17], [34, 57, 86]], dims=("y", "x")), id="2D" + ), + ], +) +def test_polyfit_polyval_integration( + use_dask: bool, x: xr.DataArray, y: xr.DataArray +) -> None: + y.coords["x"] = x + if use_dask: + if not has_dask: + pytest.skip("requires dask") + y = y.chunk({"x": 2}) + + fit = y.polyfit(dim="x", deg=2) + evaluated = xr.polyval(y.x, fit.polyfit_coefficients) + expected = y.transpose(*evaluated.dims) + xr.testing.assert_allclose(evaluated.variable, expected.variable) + + @pytest.mark.parametrize("use_dask", [False, True]) @pytest.mark.parametrize( "a, b, ae, be, dim, axis",