Better dask support in polyval

[dcherian]

Is your feature request related to a problem?

polyval does not handle dask inputs well.

nt = 8772 // 4
ny = 489
nx = 655
# chunks like the data is stored on disk
# small in time, big in space
# because the chunk sizes are -1 along lat, lon;
# reshaping this array to (time, latlon) prior to fitting is pretty cheap
chunks = (8, -1, -1)

da = xr.DataArray(
    dask.array.random.random((nt, ny, nx), chunks=chunks),
    dims=("ocean_time", "eta_rho", "xi_rho"),
)

dim = "ocean_time"
deg = 1

p = da.polyfit(dim="ocean_time", deg=1, skipna=False)

# create a chunked version of the "ocean_time" dimension
chunked_dim = xr.DataArray(
    dask.array.from_array(da[dim].data, chunks=da.chunksizes[dim]), dims=dim, name=dim
)
xr.polyval(chunked_dim, p.polyfit_coefficients)

Describe the solution you'd like

Here's a partial solution. It does not handle datetime inputs (polyval handles this using get_clean_interp_index which computes dask inputs). But I've replaced the call to np.vander and used xr.dot.

def polyval(coord, coeffs, degree_dim="degree"):
    x = coord.data

    deg_coord = coeffs[degree_dim]
    N = int(deg_coord.max()) + 1

    lhs = xr.DataArray(
        np.stack([x ** (N - 1 - i) for i in range(N)], axis=1),
        dims=(coord.name, degree_dim),
        coords={coord.name: coord, degree_dim: np.arange(deg_coord.max() + 1)[::-1]},
    )
    return xr.dot(lhs, coeffs, dims=degree_dim)

polyval(chunked_dim, p.polyfit_coefficients)

This looks like what I expected

cc @aulemahal

Describe alternatives you've considered

No response

Additional context

No response