Implementation of polyfit and polyval

doc/api.rst

@@ -30,6 +30,7 @@ Top-level functions
    zeros_like
    ones_like
    dot
+   polyval
    map_blocks
    show_versions
    set_options
@@ -171,6 +172,7 @@ Computation
    Dataset.quantile
    Dataset.differentiate
    Dataset.integrate
+   Dataset.polyfit
 
 **Aggregation**:
 :py:attr:`~Dataset.all`
@@ -349,6 +351,7 @@ Computation
    DataArray.quantile
    DataArray.differentiate
    DataArray.integrate
+   DataArray.polyfit
    DataArray.str
 
 **Aggregation**:

doc/computation.rst

@@ -317,6 +317,32 @@ trapezoidal rule using their coordinates,
     and integration along multidimensional coordinate are not supported.
 
 
+.. _compute.polyfit:
+
+Fitting polynomials
+===================
+
+Xarray objects provide an interface for performing linear or polynomial regressions
+using the least-squares method. :py:meth:`~xarray.DataArray.polyfit` computes the
+best fitting coefficients along a given dimension and for a given order,
+
+.. ipython:: python
+
+    x = np.arange(10)
+    a = xr.DataArray(3 + 4 * x, dims=['x'], coords={'x': x})
+    out = a.polyfit('x', 1, full=True)
+    out
+
+The method outputs a dataset containing the coefficients (and more if `full = True`).
+The inverse operation is done with :py:meth:`~xarray.polyval`,
+
+.. ipython::python
+
+    xr.polyval(x, out.polyfit_coefficients)
+
+.. note::
+    These methods replicate the behaviour of :py:func:`numpy.polyfit` and :py:func:`numpy.polyval`.
+
 .. _compute.broadcasting:
 
 Broadcasting by dimension name

doc/whats-new.rst

@@ -40,6 +40,8 @@ Breaking changes
 
 New Features
 ~~~~~~~~~~~~
+- Added :py:meth:`DataArray.polyfit` and :py:func:`xarray.polyval` for fitting polynomials. (:issue:`3349`)
+  By `Pascal Bourgault <https://github.com/aulemahal>`_.
 - :py:meth:`DataArray.sel` and :py:meth:`Dataset.sel` now support :py:class:`pandas.CategoricalIndex`. (:issue:`3669`)
   By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - Support using an existing, opened h5netcdf ``File`` with

xarray/__init__.py

@@ -15,7 +15,7 @@
 from .core.alignment import align, broadcast
 from .core.combine import auto_combine, combine_by_coords, combine_nested
 from .core.common import ALL_DIMS, full_like, ones_like, zeros_like
-from .core.computation import apply_ufunc, dot, where
+from .core.computation import apply_ufunc, dot, polyval, where
 from .core.concat import concat
 from .core.dataarray import DataArray
 from .core.dataset import Dataset
@@ -65,6 +65,7 @@
     "open_mfdataset",
     "open_rasterio",
     "open_zarr",
+    "polyval",
     "register_dataarray_accessor",
     "register_dataset_accessor",
     "save_mfdataset",

xarray/core/computation.py

@@ -1299,3 +1299,49 @@ def where(cond, x, y):
         dataset_join="exact",
         dask="allowed",
     )
+
+
+def polyval(coord, coeffs, degree_dim="degree"):
+    """Evaluate a polynomial at specific values
+
+    Parameters
+    ----------
+    coord : DataArray or array-like
+        The 1D coordinate along which to evaluate the polynomial.
+    coeffs : DataArray or array-like
+        Coefficients of the polynomials.
+    degree_dim : str or int, optional
+        Name of dimension or axis number along which the degree increases in `coeffs`.
+
+    See also
+    --------
+    `xarray.DataArray.polyfit`
+    `numpy.polyval`
+    """
+    from .dataarray import DataArray
+    from .missing import get_clean_interp_index
+
+    if not isinstance(coord, DataArray):
+        coord = DataArray(coord, dims=("x",), name="x")
+    x = get_clean_interp_index(coord, coord.name)
+
+    if isinstance(degree_dim, str):
+        deg_coord = coeffs[degree_dim]
+    else:  # Axis number
+        degree_axis = degree_dim
+        if isinstance(coeffs, DataArray):
+            degree_dim = coeffs.dims[degree_axis]
+            deg_coord = coeffs[degree_dim]
+        else:
+            degree_dim = "degree"
+            deg_coord = np.arange(coeffs.shape[degree_axis])[::-1]
+            coeffs = DataArray(coeffs)
+            coeffs.coords[coeffs.dims[degree_axis]] = deg_coord
+            coeffs = coeffs.rename({coeffs.dims[degree_axis]: degree_dim})
+
+    lhs = DataArray(
+        np.vander(x, int(deg_coord.max()) + 1),
+        dims=(coord.name, degree_dim),
+        coords={coord.name: coord, degree_dim: np.arange(deg_coord.max() + 1)[::-1]},
+    )
+    return (lhs * coeffs).sum(degree_dim)

xarray/core/dask_array_ops.py

@@ -95,3 +95,30 @@ def func(x, window, axis=-1):
     # crop boundary.
     index = (slice(None),) * axis + (slice(drop_size, drop_size + orig_shape[axis]),)
     return out[index]
+
+
+def least_squares(lhs, rhs, rcond=None, skipna=False):
+    import dask.array as da
+
+    lhs_da = da.from_array(lhs, chunks=(rhs.chunks[0], lhs.shape[1]))
+    if skipna:
+        added_dim = rhs.ndim == 1
+        if added_dim:
+            rhs = rhs.reshape(rhs.shape[0], 1)
+        results = da.apply_along_axis(
+            nputils._nanpolyfit_1d,
+            0,
+            rhs,
+            lhs_da,
+            dtype=float,
+            shape=(lhs.shape[1] + 1,),
+            rcond=rcond,
+        )
+        coeffs = results[:-1, ...]
+        residuals = results[-1, ...]
+        if added_dim:
+            coeffs = coeffs.reshape(coeffs.shape[0])
+            residuals = residuals.reshape(residuals.shape[0])
+    else:
+        coeffs, residuals, _, _ = da.linalg.lstsq(lhs_da, rhs)
+    return coeffs, residuals

xarray/core/dataarray.py

@@ -1,3 +1,4 @@
+import datetime
 import functools
 import warnings
 from numbers import Number
@@ -18,7 +19,6 @@
     cast,
 )
 
-import datetime
 import numpy as np
 import pandas as pd
 
@@ -3234,6 +3234,43 @@ def map_blocks(
 
         return map_blocks(func, self, args, kwargs)
 
+    def polyfit(
+        self,
+        dim: Hashable,
+        deg: int,
+        skipna: bool = None,
+        rcond: float = None,
+        w: Union[Hashable, Any] = None,
+        full: bool = False,
+        cov: bool = False,
+    ):
+        """
+        Least squares polynomial fit.
+
+        This replicates the behaviour of `numpy.polyfit` but differs by skipping
+        invalid values when `skipna = True`.
+
+        Parameters
+        ----------
+        dim : hashable
+            Coordinate along which to fit the polynomials.
+        skipna : bool, optional
+            If True, removes all invalid values before fitting each 1D slices of the array
+            Defaults is True if data is stored in as dask.array or if there is any
+            invalid values, False otherwise.
+
+        See documentation of `numpy.polyfit`.
+
+        Returns
+        -------
+        A single dataset with the coefficients of the best fit for this DataArray.
+        The residuals, singular values and rank are included if `full` is True.
+        The covariance matrix is included if `full` is False and `cov` is True.
+        """
+        return self._to_temp_dataset().polyfit(
+            dim, deg, skipna=skipna, rcond=rcond, w=w, full=full, cov=cov
+        )
+
     # this needs to be at the end, or mypy will confuse with `str`
     # https://mypy.readthedocs.io/en/latest/common_issues.html#dealing-with-conflicting-names
     str = property(StringAccessor)