Fix polyfit fail on deficient rank (#4193)

* Fix polyfit fail on deficient rank

* Add docs and RankWarning

* Fix deficient ranks outputs | workaround dask bug | add tests

* Add a note to the doc | whats new entry

* Update xarray/core/nputils.py

Apply suggestion from review.

Co-authored-by: Mathias Hauser <mathause@users.noreply.github.com>

* Fix test and catch warnings

* forgot to run black

* adapt polyfit test to properly test issue #4190

* Fix syntax in doc/whats-new.rst

Co-authored-by: keewis <keewis@users.noreply.github.com>

Co-authored-by: Mathias Hauser <mathause@users.noreply.github.com>
Co-authored-by: keewis <keewis@users.noreply.github.com>

Author: 3 people

doc/whats-new.rst

@@ -65,6 +65,7 @@ Bug fixes
   unchunked dimensions (:pull:`4312`) By `Tobias Kölling <https://github.com/d70-t>`_.
 - Fixed a bug in backend caused by basic installation of Dask (:issue:`4164`, :pull:`4318`)
   `Sam Morley <https://github.com/inakleinbottle>`_.
+- Fixed a few bugs with :py:meth:`Dataset.polyfit` when encountering deficient matrix ranks (:issue:`4190`, :pull:`4193`). By `Pascal Bourgault <https://github.com/aulemahal>`_.
 - Fixed inconsistencies between docstring and functionality for :py:meth:`DataArray.str.get`
   and :py:meth:`DataArray.str.wrap` (:issue:`4334`). By `Mathias Hauser <https://github.com/mathause>`_.
 - Fixed overflow issue causing incorrect results in computing means of :py:class:`cftime.datetime`

xarray/core/dask_array_ops.py

@@ -131,5 +131,7 @@ def least_squares(lhs, rhs, rcond=None, skipna=False):
             coeffs = coeffs.reshape(coeffs.shape[0])
             residuals = residuals.reshape(residuals.shape[0])
     else:
+        # Residuals here are (1, 1) but should be (K,) as rhs is (N, K)
+        # See issue dask/dask#6516
         coeffs, residuals, _, _ = da.linalg.lstsq(lhs_da, rhs)
     return coeffs, residuals

xarray/core/dataarray.py

@@ -3478,7 +3478,8 @@ def polyfit(
             polyfit_coefficients
                 The coefficients of the best fit.
             polyfit_residuals
-                The residuals of the least-square computation (only included if `full=True`)
+                The residuals of the least-square computation (only included if `full=True`).
+                When the matrix rank is deficient, np.nan is returned.
             [dim]_matrix_rank
                 The effective rank of the scaled Vandermonde coefficient matrix (only included if `full=True`)
             [dim]_singular_value

xarray/core/dataset.py

@@ -5957,13 +5957,21 @@ def polyfit(
                 The coefficients of the best fit for each variable in this dataset.
             [var]_polyfit_residuals
                 The residuals of the least-square computation for each variable (only included if `full=True`)
+                When the matrix rank is deficient, np.nan is returned.
             [dim]_matrix_rank
                 The effective rank of the scaled Vandermonde coefficient matrix (only included if `full=True`)
+                The rank is computed ignoring the NaN values that might be skipped.
             [dim]_singular_values
                 The singular values of the scaled Vandermonde coefficient matrix (only included if `full=True`)
             [var]_polyfit_covariance
                 The covariance matrix of the polynomial coefficient estimates (only included if `full=False` and `cov=True`)
 
+        Warns
+        -----
+        RankWarning
+            The rank of the coefficient matrix in the least-squares fit is deficient.
+            The warning is not raised with in-memory (not dask) data and `full=True`.
+
         See also
         --------
         numpy.polyfit
@@ -5997,10 +6005,6 @@ def polyfit(
         degree_dim = utils.get_temp_dimname(self.dims, "degree")
 
         rank = np.linalg.matrix_rank(lhs)
-        if rank != order and not full:
-            warnings.warn(
-                "Polyfit may be poorly conditioned", np.RankWarning, stacklevel=4
-            )
 
         if full:
             rank = xr.DataArray(rank, name=xname + "matrix_rank")
@@ -6009,7 +6013,7 @@ def polyfit(
             sing = xr.DataArray(
                 sing,
                 dims=(degree_dim,),
-                coords={degree_dim: np.arange(order)[::-1]},
+                coords={degree_dim: np.arange(rank - 1, -1, -1)},
                 name=xname + "singular_values",
             )
             variables[sing.name] = sing
@@ -6018,11 +6022,14 @@ def polyfit(
             if dim not in da.dims:
                 continue
 
-            if skipna is None:
-                if isinstance(da.data, dask_array_type):
-                    skipna_da = True
-                else:
-                    skipna_da = np.any(da.isnull())
+            if isinstance(da.data, dask_array_type) and (
+                rank != order or full or skipna is None
+            ):
+                # Current algorithm with dask and skipna=False neither supports
+                # deficient ranks nor does it output the "full" info (issue dask/dask#6516)
+                skipna_da = True
+            elif skipna is None:
+                skipna_da = np.any(da.isnull())
 
             dims_to_stack = [dimname for dimname in da.dims if dimname != dim]
             stacked_coords: Dict[Hashable, DataArray] = {}
@@ -6040,9 +6047,15 @@ def polyfit(
             if w is not None:
                 rhs *= w[:, np.newaxis]
 
-            coeffs, residuals = duck_array_ops.least_squares(
-                lhs, rhs.data, rcond=rcond, skipna=skipna_da
-            )
+            with warnings.catch_warnings():
+                if full:  # Copy np.polyfit behavior
+                    warnings.simplefilter("ignore", np.RankWarning)
+                else:  # Raise only once per variable
+                    warnings.simplefilter("once", np.RankWarning)
+
+                coeffs, residuals = duck_array_ops.least_squares(
+                    lhs, rhs.data, rcond=rcond, skipna=skipna_da
+                )
 
             if isinstance(name, str):
                 name = "{}_".format(name)

xarray/core/nputils.py

@@ -232,10 +232,17 @@ def _nanpolyfit_1d(arr, x, rcond=None):
     out = np.full((x.shape[1] + 1,), np.nan)
     mask = np.isnan(arr)
     if not np.all(mask):
-        out[:-1], out[-1], _, _ = np.linalg.lstsq(x[~mask, :], arr[~mask], rcond=rcond)
+        out[:-1], resid, rank, _ = np.linalg.lstsq(x[~mask, :], arr[~mask], rcond=rcond)
+        out[-1] = resid if resid.size > 0 else np.nan
+        warn_on_deficient_rank(rank, x.shape[1])
     return out
 
 
+def warn_on_deficient_rank(rank, order):
+    if rank != order:
+        warnings.warn("Polyfit may be poorly conditioned", np.RankWarning, stacklevel=2)
+
+
 def least_squares(lhs, rhs, rcond=None, skipna=False):
     if skipna:
         added_dim = rhs.ndim == 1
@@ -248,16 +255,21 @@ def least_squares(lhs, rhs, rcond=None, skipna=False):
                 _nanpolyfit_1d, 0, rhs[:, nan_cols], lhs
             )
         if np.any(~nan_cols):
-            out[:-1, ~nan_cols], out[-1, ~nan_cols], _, _ = np.linalg.lstsq(
+            out[:-1, ~nan_cols], resids, rank, _ = np.linalg.lstsq(
                 lhs, rhs[:, ~nan_cols], rcond=rcond
             )
+            out[-1, ~nan_cols] = resids if resids.size > 0 else np.nan
+            warn_on_deficient_rank(rank, lhs.shape[1])
         coeffs = out[:-1, :]
         residuals = out[-1, :]
         if added_dim:
             coeffs = coeffs.reshape(coeffs.shape[0])
             residuals = residuals.reshape(residuals.shape[0])
     else:
-        coeffs, residuals, _, _ = np.linalg.lstsq(lhs, rhs, rcond=rcond)
+        coeffs, residuals, rank, _ = np.linalg.lstsq(lhs, rhs, rcond=rcond)
+        if residuals.size == 0:
+            residuals = coeffs[0] * np.nan
+        warn_on_deficient_rank(rank, lhs.shape[1])
     return coeffs, residuals
 
 

xarray/tests/test_dataarray.py

@@ -4301,8 +4301,14 @@ def test_polyfit(self, use_dask, use_datetime):
         ).T
         assert_allclose(out.polyfit_coefficients, expected, rtol=1e-3)
 
+        # Full output and deficient rank
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", np.RankWarning)
+            out = da.polyfit("x", 12, full=True)
+            assert out.polyfit_residuals.isnull().all()
+
         # With NaN
-        da_raw[0, 1] = np.nan
+        da_raw[0, 1:3] = np.nan
         if use_dask:
             da = da_raw.chunk({"d": 1})
         else:
@@ -4317,6 +4323,11 @@ def test_polyfit(self, use_dask, use_datetime):
         assert out.x_matrix_rank == 3
         np.testing.assert_almost_equal(out.polyfit_residuals, [0, 0])
 
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", np.RankWarning)
+            out = da.polyfit("x", 8, full=True)
+            np.testing.assert_array_equal(out.polyfit_residuals.isnull(), [True, False])
+
     def test_pad_constant(self):
         ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5))
         actual = ar.pad(dim_0=(1, 3))

xarray/tests/test_dataset.py

@@ -5621,6 +5621,16 @@ def test_polyfit_output(self):
         out = ds.polyfit("time", 2)
         assert len(out.data_vars) == 0
 
+    def test_polyfit_warnings(self):
+        ds = create_test_data(seed=1)
+
+        with warnings.catch_warnings(record=True) as ws:
+            ds.var1.polyfit("dim2", 10, full=False)
+            assert len(ws) == 1
+            assert ws[0].category == np.RankWarning
+            ds.var1.polyfit("dim2", 10, full=True)
+            assert len(ws) == 1
+
     def test_pad(self):
         ds = create_test_data(seed=1)
         padded = ds.pad(dim2=(1, 1), constant_values=42)