Use the same function to floatize coords in polyfit and polyval (pyda…

…ta#9691) * use same function to floatize coords in polyfit and polyval * Add whats new - fit typing - avoid warnings in tests * requires cftime * Ignore mypy issues with rcond * Apply suggestions from code review Co-authored-by: Deepak Cherian <[email protected]> --------- Co-authored-by: Deepak Cherian <[email protected]>
dcherian · Oct 29, 2024 · cdec18f · cdec18f
1 parent 0c6cded
commit cdec18f
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 13 deletions.
diff --git 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 <https://github.com/shoyer>`_.
+- Fix regression in the interoperability of :py:meth:`DataArray.polyfit` and :py:meth:`xr.polyval` for date-time coordinates. (:pull:`9691`).
+  By `Pascal Bourgault <https://github.com/aulemahal>`_.
 
 Documentation
 ~~~~~~~~~~~~~

diff --git 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,22 +9065,14 @@ 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 = np.asarray(_ensure_numeric(self.coords[dim]).astype(np.float64))
 
         xname = f"{self[dim].name}_"
         order = int(deg) + 1
         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:

diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
@@ -6720,6 +6720,34 @@ 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(
+            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)
+        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)
+        xr.testing.assert_allclose(da_fitval, da, rtol=1e-3)
+
+    @requires_cftime
+    def test_polyfit_polyval_cftime(self) -> None:
+        da = xr.DataArray(
+            np.arange(1, 10).astype(np.float64),
+            dims=["x"],
+            coords=dict(
+                x=xr.date_range("2001-01-01", periods=9, 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
@@ -7230,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)