polyval: Use Horner's algorithm + support chunked inputs

doc/whats-new.rst

@@ -41,6 +41,9 @@ New Features
 - Allow passing chunks in ``**kwargs`` form to :py:meth:`Dataset.chunk`, :py:meth:`DataArray.chunk`, and
   :py:meth:`Variable.chunk`. (:pull:`6471`)
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- :py:meth:`xr.polyval` now supports :py:class:`Dataset` and :py:class:`DataArray` args of any shape,
+  is faster and requires less memory. (:pull:`6548`)
+  By `Michael Niklas <https://github.com/headtr1ck>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/core/computation.py

@@ -17,12 +17,15 @@
     Iterable,
     Mapping,
     Sequence,
+    overload,
 )
 
 import numpy as np
 
 from . import dtypes, duck_array_ops, utils
 from .alignment import align, deep_align
+from .common import zeros_like
+from .duck_array_ops import datetime_to_numeric
 from .indexes import Index, filter_indexes_from_coords
 from .merge import merge_attrs, merge_coordinates_without_align
 from .options import OPTIONS, _get_keep_attrs
@@ -1843,36 +1846,84 @@ def where(cond, x, y, keep_attrs=None):
     )
 
 
-def polyval(coord, coeffs, degree_dim="degree"):
+@overload
+def polyval(coord: DataArray, coeffs: DataArray, degree_dim: Hashable) -> DataArray:
+    ...
+
+
+@overload
+def polyval(coord: T_Xarray, coeffs: Dataset, degree_dim: Hashable) -> Dataset:
+    ...
+
+
+@overload
+def polyval(coord: Dataset, coeffs: T_Xarray, degree_dim: Hashable) -> Dataset:
+    ...
+
+
+def polyval(
+    coord: T_Xarray, coeffs: T_Xarray, degree_dim: Hashable = "degree"
+) -> T_Xarray:
     """Evaluate a polynomial at specific values
 
     Parameters
     ----------
-    coord : DataArray
-        The 1D coordinate along which to evaluate the polynomial.
-    coeffs : DataArray
-        Coefficients of the polynomials.
-    degree_dim : str, default: "degree"
+    coord : DataArray or Dataset
+        Values at which to evaluate the polynomial.
+    coeffs : DataArray or Dataset
+        Coefficients of the polynomial.
+    degree_dim : Hashable, default: "degree"
         Name of the polynomial degree dimension in `coeffs`.
 
+    Returns
+    -------
+    DataArray or Dataset
+        Evaluated polynomial.
+
     See Also
     --------
     xarray.DataArray.polyfit
-    numpy.polyval
+    numpy.polynomial.polynomial.polyval
     """
     from .dataarray import DataArray
-    from .missing import get_clean_interp_index
-
-    x = get_clean_interp_index(coord, coord.name, strict=False)
+    from .dataset import Dataset
 
     deg_coord = coeffs[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]},
+    deg_idx_sorted = np.argsort(deg_coord.values)
+    max_deg = int(deg_coord[deg_idx_sorted[-1]])
+
+    def ensure_numeric(data: DataArray) -> DataArray:
+        if data.dtype.kind in "mM":
+            return DataArray(
+                datetime_to_numeric(
+                    data, offset=np.datetime64("1970-01-01"), datetime_unit="ns"
+                ),
+                dims=data.dims,
+                coords=data.coords,
+                attrs=data.attrs,
+            )
+        return data
+
+    if isinstance(coord, Dataset):
+        coord = coord.map(ensure_numeric)
+    else:
+        coord = ensure_numeric(coord)
+
+    # using Horner's method
+    # https://en.wikipedia.org/wiki/Horner%27s_method
+    res = coeffs.isel({degree_dim: int(deg_idx_sorted[-1])}, drop=True) + zeros_like(
+        coord
     )
-    return (lhs * coeffs).sum(degree_dim)
+    deg_idx = len(deg_coord) - 2
+    for deg in range(max_deg - 1, -1, -1):
+        res *= coord
+        if deg_idx >= 0 and deg == int(deg_coord[deg_idx_sorted[deg_idx]]):
+            # this degrees coefficient is provided, if not assume 0
+            res += coeffs.isel({degree_dim: int(deg_idx_sorted[deg_idx])}, drop=True)
+            deg_idx -= 1
+
+    return res
 
 
 def _calc_idxminmax(

xarray/tests/test_computation.py

@@ -1933,37 +1933,77 @@ def test_where_attrs() -> None:
     assert actual.attrs == {}
 
 
-@pytest.mark.parametrize("use_dask", [True, False])
-@pytest.mark.parametrize("use_datetime", [True, False])
-def test_polyval(use_dask, use_datetime) -> None:
-    if use_dask and not has_dask:
-        pytest.skip("requires dask")
-
-    if use_datetime:
-        xcoord = xr.DataArray(
-            pd.date_range("2000-01-01", freq="D", periods=10), dims=("x",), name="x"
-        )
-        x = xr.core.missing.get_clean_interp_index(xcoord, "x")
-    else:
-        x = np.arange(10)
-        xcoord = xr.DataArray(x, dims=("x",), name="x")
-
-    da = xr.DataArray(
-        np.stack((1.0 + x + 2.0 * x**2, 1.0 + 2.0 * x + 3.0 * x**2)),
-        dims=("d", "x"),
-        coords={"x": xcoord, "d": [0, 1]},
-    )
-    coeffs = xr.DataArray(
-        [[2, 1, 1], [3, 2, 1]],
-        dims=("d", "degree"),
-        coords={"d": [0, 1], "degree": [2, 1, 0]},
-    )
+@pytest.mark.parametrize("use_dask", [False, True])
+@pytest.mark.parametrize(
+    ["x", "coeffs", "expected"],
+    [
+        pytest.param(
+            xr.DataArray([1, 2, 3], dims="x"),
+            xr.DataArray([2, 3, 4], dims="degree"),
+            xr.DataArray([9, 2 + 6 + 16, 2 + 9 + 36], dims="x"),
+            id="simple",
+        ),
+        pytest.param(
+            xr.DataArray([1, 2, 3], dims="x"),
+            xr.DataArray([[0, 1], [0, 1]], dims=("y", "degree")),
+            xr.DataArray([[1, 2, 3], [1, 2, 3]], dims=("y", "x")),
+            id="broadcast-x",
+        ),
+        pytest.param(
+            xr.DataArray([1, 2, 3], dims="x"),
+            xr.DataArray([[0, 1], [1, 0], [1, 1]], dims=("x", "degree")),
+            xr.DataArray([1, 1, 1 + 3], dims="x"),
+            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"),
+            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"),
+            id="sparse-index",
+        ),
+        pytest.param(
+            xr.DataArray([1, 2, 3], dims="x"),
+            xr.Dataset({"a": ("degree", [0, 1]), "b": ("degree", [1, 0])}),
+            xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [1, 1, 1])}),
+            id="array-dataset",
+        ),
+        pytest.param(
+            xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [2, 3, 4])}),
+            xr.DataArray([1, 1], dims="degree"),
+            xr.Dataset({"a": ("x", [2, 3, 4]), "b": ("x", [3, 4, 5])}),
+            id="dataset-array",
+        ),
+        pytest.param(
+            xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("y", [2, 3, 4])}),
+            xr.Dataset({"a": ("degree", [0, 1]), "b": ("degree", [1, 1])}),
+            xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("y", [3, 4, 5])}),
+            id="dataset-dataset",
+        ),
+        pytest.param(
+            xr.DataArray(pd.date_range("1970-01-01", freq="s", periods=3), dims="x"),
+            xr.DataArray([0, 1], dims="degree"),
+            xr.DataArray(
+                [0, 1e9, 2e9],
+                dims="x",
+                coords={"x": pd.date_range("1970-01-01", freq="s", periods=3)},
+            ),
+            id="datetime",
+        ),
+    ],
+)
+def test_polyval(use_dask, x, coeffs, expected) -> None:
     if use_dask:
-        coeffs = coeffs.chunk({"d": 2})
-
-    da_pv = xr.polyval(da.x, coeffs)
-
-    xr.testing.assert_allclose(da, da_pv.T)
+        if not has_dask:
+            pytest.skip("requires dask")
+        coeffs = coeffs.chunk({"degree": 2})
+    actual = xr.polyval(x, coeffs)
+    xr.testing.assert_allclose(actual, expected)
 
 
 @pytest.mark.parametrize("use_dask", [False, True])