polyval: Use Horner's algorithm + support chunked inputs

[headtr1ck]

• Closes xr.polyval first arg requires name attribute #6526, Better dask support in polyval #6411
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

[headtr1ck]

Several open points still:

1. Unittests for datetime values are failing, I might need some help with that since I have no idea what this means for polynomials.
2. Algorithm should work also with Datasets (any combination of DataArray and Dataset for coord and coeffs inputs). Still needs to be checked and tested (How does one define typing for such cases? I.e. DataArray + Dataset -> Dataset but DataArray + DataArray -> DataArray?)
3. It uses Horners method instead of Vandermonde matrix, should be faster and consume less memory (unless the overhead of sorting index, isel etc. is too large). Maybe some performance comparisons should be done.
4. Instead of coord the input should simply be called x or similar, however this would break backwards compatibility, so maybe we just leave it.
5. I had to add a copy(deep=True) since the broadcast returned a read-only DataArray. Any better ideas?

[headtr1ck]

I noticed that broadcasting Datasets behaves weird, see #6549, so I used a "hack" of adding an 0-valued DataArray/Dataset.
Anyone got a better idea?

[max-sixty]

This looks like excellent code @headtr1ck , especially so for a first PR. Welcome!

so I used a "hack" of adding an 0-valued DataArray/Dataset.

I think it's fine to do the hack and reference that issue; we can clean it up then (though if others have ideas then great).

I had to add a copy(deep=True) since the broadcast returned a read-only DataArray. Any better ideas?

IIRC we generally leave this up to the caller (and generally discourage mutation).

---

Someone who knows better should review the algo & opine on the datetime issue; I agree I'm not sure whether we need to support them.

[headtr1ck]

Some performance comparison:
With 5th order polynomial and 10 x-values:
old: 1.05 ms ± 15.8 µs per loop
new: 1.41 ms ± 11.6 µs per loop

With 5th order polynomial and 10000 x-values:
old: 1.46 ms ± 10.5 µs per loop
new: 1.41 ms ± 14.5 µs per loop

With 5th order polynomial and 1mio x-values:
old: 65.1 ms ± 332 µs per loop
new: 6.99 ms ± 168 µs per loop

As expected for small arrays the new method creates some overhead, but for larger arrays the speedup is quite nice.
Also, it uses in-place operations with much less memory usage.

[headtr1ck]

I added a rough support for datetime values. Someone with more knowledge of handling them should take a look, the code seems too complicated and I am sure there is a more clever solution (I could not use get_clean_interp_index since it is not an index anymore).

I agree I'm not sure whether we need to support them.

I think keeping support is nice, since they are a commonly occuring coordinates and we do not want to break anything if possible.

[dcherian]

Thanks @headtr1ck . What a great PR! This looks like a great improvement; the tests are certainly more readable.

I've left some minor comments.

One optional request: since you seem to have benchmark code, can you add it to benchmarks/polyfit.py (new file)? If this is too much, just add the code in a comment here and I'll send in a followup PR. We have some documentation on using asv here.

[max-sixty]

Benchmark did not succeed since the inputs are not compatible with the old algorith...
Do we change it such that it is compatible?

If you're confident this is faster and it's not a trivial amount of work to adjust them, I would leave it...

[headtr1ck]

Edit: nvmd, was only confusing output when the benchmark was failing. Now the benchmark looks good :)

First time working with asv...
It seems that module level variables affect all other peakmem tests (i guess memory usage of the phyton process is measured,).

We should refactor all dataarrays into the setup functions, otherwise O(n) memory algos will show wrong numbers and adding new tests will show regression on other tests.

[dcherian]

Nice this looks like an improvement for everything other than dask arrays with only 100 elements, which is not a good use-case for dask.

I was slightly concerned that the recrusive algorithm wouldn't work well with dask but it does seem to work better.

def other_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.data,
            degree_dim: np.arange(deg_coord.max() + 1)[::-1],
        },
    )
    return xr.dot(lhs, coeffs, dims=degree_dim)


coeffs = xr.DataArray(np.random.randn(2), dims="degree")
da = xr.DataArray(dask.array.random.random((10**6), chunks=(10000)), dims=("x"))
print(len(da.data.dask))
print(len(xr.polyval(da, coeffs).data.dask))
print(len(other_polyval(da, coeffs).data.dask))

100
502
1005

[headtr1ck]

One minor open point: what to do with a non-integer "degree" index?
Float type could be cast to integer (thats what is happening now).
But (nonsense) datetime etc. should raise an error?

[headtr1ck]

Personally I would allow coeffs without explicit index since I am a lazy person and would like to do coeffs = xr.DataArray([1,2], dims="degree").
But I guess with the new indexing system you want to encourage people to use them.

But I am happy with this code and look forward to use it in my projects :)

[dcherian]

Forcing the user to be explicit reduces bugs and user support requests :) so we like to do that.

Thanks again @headtr1ck this is a great PR!

[max-sixty]

Thanks @headtr1ck !