From 0c41bb4c255531144908504e7119a3ec855936bb Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Mon, 28 Oct 2024 17:38:14 -0400 Subject: [PATCH 1/5] use same function to floatize coords in polyfit and polyval --- xarray/core/dataset.py | 13 ++----------- xarray/tests/test_dataset.py | 26 ++++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 66ceea17b91..abe8a636b46 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -60,7 +60,7 @@ _contains_datetime_like_objects, get_chunksizes, ) -from xarray.core.computation import unify_chunks +from xarray.core.computation import _ensure_numeric, unify_chunks from xarray.core.coordinates import ( Coordinates, DatasetCoordinates, @@ -87,7 +87,6 @@ merge_coordinates_without_align, merge_core, ) -from xarray.core.missing import _floatize_x from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.core.types import ( Bins, @@ -9066,15 +9065,7 @@ def polyfit( variables = {} skipna_da = skipna - x: Any = self.coords[dim].variable - x = _floatize_x((x,), (x,))[0][0] - - try: - x = x.values.astype(np.float64) - except TypeError as e: - raise TypeError( - f"Dim {dim!r} must be castable to float64, got {type(x).__name__}." - ) from e + x: Any = _ensure_numeric(self.coords[dim]).astype(np.float64) xname = f"{self[dim].name}_" order = int(deg) + 1 diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index c89cfa85622..923e7420bf2 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6720,6 +6720,32 @@ def test_polyfit_warnings(self) -> None: ds.var1.polyfit("dim2", 10, full=True) assert len(ws) == 1 + def test_polyfit_polyval(self) -> None: + da = xr.DataArray([1.0, 2.0, 3.0], dims=["x"], coords=dict(x=[0, 1, 2])) + + out = da.polyfit("x", 3, full=False) + da_fitval = xr.polyval(da.x, out.polyfit_coefficients) + # polyval introduces very small errors (1e-16 here) + np.testing.assert_allclose(da_fitval, da) + + da = da.assign_coords(x=xr.date_range("2001-01-01", periods=3, freq="YS")) + out = da.polyfit("x", 3, full=False) + da_fitval = xr.polyval(da.x, out.polyfit_coefficients) + np.testing.assert_allclose(da_fitval, da) + + @pytest.mark.skipif(not has_cftime, reason="Test requires cftime.") + def test_polyfit_polyval_cftime(self) -> None: + da = xr.DataArray( + [1.0, 2.0, 3.0], + dims=["x"], + coords=dict( + x=xr.date_range("2001-01-01", periods=3, freq="YS", calendar="noleap") + ), + ) + out = da.polyfit("x", 3, full=False) + da_fitval = xr.polyval(da.x, out.polyfit_coefficients) + np.testing.assert_allclose(da_fitval, da) + @staticmethod def _test_data_var_interior( original_data_var, padded_data_var, padded_dim_name, expected_pad_values From 9b3e00fbf99303e32c9b6f7dacff97d1c6d95b9f Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Tue, 29 Oct 2024 11:34:25 -0400 Subject: [PATCH 2/5] Add whats new - fit typing - avoid warnings in tests --- doc/whats-new.rst | 2 ++ xarray/core/dataset.py | 2 +- xarray/tests/test_dataset.py | 12 +++++++----- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 52f75c1a6f1..659063f2cbf 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -37,6 +37,8 @@ Bug fixes - Fix inadvertent deep-copying of child data in DataTree. By `Stephan Hoyer `_. +- Fix regression in the interoperability of :py:meth:`DataArray.polyfit` and :py:meth:`xr.polyval` for date-time coordinates. (:pull:`9691`). + By `Pascal Bourgault `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index b0f94e6f1df..25ba5ffe67e 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -9065,7 +9065,7 @@ def polyfit( variables = {} skipna_da = skipna - x: Any = _ensure_numeric(self.coords[dim]).astype(np.float64) + x = np.asarray(_ensure_numeric(self.coords[dim]).astype(np.float64)) xname = f"{self[dim].name}_" order = int(deg) + 1 diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 923e7420bf2..f3994c15d56 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6721,25 +6721,27 @@ def test_polyfit_warnings(self) -> None: assert len(ws) == 1 def test_polyfit_polyval(self) -> None: - da = xr.DataArray([1.0, 2.0, 3.0], dims=["x"], coords=dict(x=[0, 1, 2])) + da = xr.DataArray( + np.arange(1, 10).astype(np.float64), dims=["x"], coords=dict(x=np.arange(9)) + ) out = da.polyfit("x", 3, full=False) da_fitval = xr.polyval(da.x, out.polyfit_coefficients) # polyval introduces very small errors (1e-16 here) np.testing.assert_allclose(da_fitval, da) - da = da.assign_coords(x=xr.date_range("2001-01-01", periods=3, freq="YS")) + da = da.assign_coords(x=xr.date_range("2001-01-01", periods=9, freq="YS")) out = da.polyfit("x", 3, full=False) da_fitval = xr.polyval(da.x, out.polyfit_coefficients) - np.testing.assert_allclose(da_fitval, da) + np.testing.assert_allclose(da_fitval, da, rtol=1e-3) @pytest.mark.skipif(not has_cftime, reason="Test requires cftime.") def test_polyfit_polyval_cftime(self) -> None: da = xr.DataArray( - [1.0, 2.0, 3.0], + np.arange(1, 10).astype(np.float64), dims=["x"], coords=dict( - x=xr.date_range("2001-01-01", periods=3, freq="YS", calendar="noleap") + x=xr.date_range("2001-01-01", periods=9, freq="YS", calendar="noleap") ), ) out = da.polyfit("x", 3, full=False) From fad554c9e7d356f779caca91bedaaa7c0cd37437 Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Tue, 29 Oct 2024 11:35:51 -0400 Subject: [PATCH 3/5] requires cftime --- xarray/tests/test_dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index f3994c15d56..6466858100e 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6735,7 +6735,7 @@ def test_polyfit_polyval(self) -> None: da_fitval = xr.polyval(da.x, out.polyfit_coefficients) np.testing.assert_allclose(da_fitval, da, rtol=1e-3) - @pytest.mark.skipif(not has_cftime, reason="Test requires cftime.") + @requires_cftime def test_polyfit_polyval_cftime(self) -> None: da = xr.DataArray( np.arange(1, 10).astype(np.float64), @@ -7258,7 +7258,7 @@ def test_differentiate_datetime(dask) -> None: assert np.allclose(actual, 1.0) -@pytest.mark.skipif(not has_cftime, reason="Test requires cftime.") +@requires_cftime @pytest.mark.parametrize("dask", [True, False]) def test_differentiate_cftime(dask) -> None: rs = np.random.RandomState(42) From 20feb4dc13677d348d7ed6cb2a0a38ba57da5c11 Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Tue, 29 Oct 2024 12:40:28 -0400 Subject: [PATCH 4/5] Ignore mypy issues with rcond --- xarray/core/dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 25ba5ffe67e..bc9360a809d 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -9072,7 +9072,7 @@ def polyfit( lhs = np.vander(x, order) if rcond is None: - rcond = x.shape[0] * np.finfo(x.dtype).eps + rcond = x.shape[0] * np.finfo(x.dtype).eps # type: ignore[assignment] # Weights: if w is not None: From 2904dafeb4f7988c27c22d2d17edafb9dde42522 Mon Sep 17 00:00:00 2001 From: Pascal Bourgault Date: Tue, 29 Oct 2024 19:18:22 -0400 Subject: [PATCH 5/5] Apply suggestions from code review Co-authored-by: Deepak Cherian --- xarray/tests/test_dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 6466858100e..95285b1c913 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6728,12 +6728,12 @@ def test_polyfit_polyval(self) -> None: out = da.polyfit("x", 3, full=False) da_fitval = xr.polyval(da.x, out.polyfit_coefficients) # polyval introduces very small errors (1e-16 here) - np.testing.assert_allclose(da_fitval, da) + xr.testing.assert_allclose(da_fitval, da) da = da.assign_coords(x=xr.date_range("2001-01-01", periods=9, freq="YS")) out = da.polyfit("x", 3, full=False) da_fitval = xr.polyval(da.x, out.polyfit_coefficients) - np.testing.assert_allclose(da_fitval, da, rtol=1e-3) + xr.testing.assert_allclose(da_fitval, da, rtol=1e-3) @requires_cftime def test_polyfit_polyval_cftime(self) -> None: