Reimplement .polyfit() with apply_ufunc

[slevang]

• Closes Unexpected chunking of 3d DataArray in polyfit() #4554
• Closes Polyfit performance on large datasets - Suboptimal dask task graph #5629
• Closes polyfit with weights alters the DataArray in place #5644
• Tests added
• Passes pre-commit run --all-files
• User visible changes (including notable bug fixes) are documented in whats-new.rst

Reimplement polyfit using apply_ufunc rather than dask.array.linalg.lstsq. This should solve a number of issues with memory usage and chunking that were reported on the current version of polyfit. The main downside is that variables chunked along the fitting dimension cannot be handled with this approach.

There is a bunch of fiddly code here for handling the differing outputs from np.polyfit depending on the values of the full and cov args. Depending on the performance implications, we could simplify some by keeping these in apply_ufunc and dropping later. Much of this parsing would still be required though, because the only way to get the covariances is to set cov=True, full=False.

A few minor departures from the previous implementation:

1. The rank and singular_values diagnostic variables returned by np.polyfit are now returned on a pointwise basis, since these can change depending on skipped nans. np.polyfit also returns the rcond used for each fit which I've included here.
2. As mentioned above, this breaks fitting done along a chunked dimension. To avoid regression, we could set allow_rechunk=True and warn about memory implications.
3. Changed default skipna=True, since the previous behavior seemed to be a limitation of the computational method.
4. For consistency with the previous version, I included a transpose operation to put degree as the first dimension. This is arbitrary though, and actually the opposite of how curvefit returns ordering. So we could match up with curvefit but it would be breaking for polyfit.

No new tests have been added since the previous suite was fairly comprehensive. Would be great to get some performance reports on real-world data such as the climate model detrending application in #5629.

[github-actions]

Unit Test Results

         6 files           6 suites   58m 4s ⏱️
16 290 tests 14 551 ✔️ 1 739 💤 0 ❌
90 936 runs  82 738 ✔️ 8 198 💤 0 ❌

Results for commit 62b4637.

[mathause]

Looks good to me!

[mathause]

You could do

-        skipna : bool, optional
+        skipna : bool, default: True

instead.

[mathause]

numpy recommends to use Polynomial.fit <numpy.polynomial.polynomial.Polynomial.fit> class - did you consider switching to this (no requirement just a question).

[mathause]

It might not be worth it, but you could avoid a if conditional in the inner loop as follows:

def _wrapper_skipna(x, y):
    ...

def _wrapper_noskipna(x, y):
    ...

if skipna:
    _wrapper = _wrapper_skipna
else:
    _wrapper = _wrapper_noskipna

[mathause]

@slevang are you still interested to continue this PR? I think that would be a worthwhile addition and should not be too much left to do. (What would be nice, however, are tests for the issues this fixes.)

[slevang]

@slevang are you still interested to continue this PR? I think that would be a worthwhile addition and should not be too much left to do. (What would be nice, however, are tests for the issues this fixes.)

Definitely! I got distracted is all, and @dcherian posted a nice solution in #5629 that could allow us to preserve the ability to fit along a chunked dimension using blockwise operations and the dask lstsq implementation used by the existing polyfit code.

I'm happy to pick this back up and finish it off if there is consensus on the right way forward, but the blockwise approach seemed promising so I put this on hold.

[Illviljan]

@slevang Yeah, the blockwise approach seems indeed nice. You're very welcome to continue with the blockwise approach in a different PR if you want to.

[slevang]

Not sure I understand the blockwise approach well enough to make a PR, but maybe I'll give it a try at some point.

[Illviljan]

I think you can use a lot of @dcherian's code as a base and then for starters see if it simply passes all the tests (including the ones you added here). If you make a draft PR here it's easier to help out as well if you're getting stuck.

[dcherian]

Underlying problem should be solved by #9766