Support multidimensional weights in polyfit

[mgunyho]

Is your feature request related to a problem?

I have a set of measurement data with errors that vary across the measurements, and I would like to fit a polynomial to them. However, the w argument of DataArray.polyfit only accepts a one-dimensional array:

xarray/xarray/core/dataset.py

Lines 9165 to 9166 in 7fd572d

	if w.ndim != 1:
	raise TypeError("Expected a 1-d array for weights.")

Here's a minimum example code snippet:

import xarray as xr

data = xr.DataArray([[1, 2, 3], [4, 5, 6]], coords=[("x", [1, 2]), ("y", [0.1, 0.2, 0.3])])

# dummy 2D weights
w = 0.1 * data

data.polyfit("y", deg=2, w=w)

Describe the solution you'd like

The w argument of DataArray.polyfit should accept a DataArray, which would be broadcast against the array that is being fit. Perhaps for Dataset, w could also be the name of a variable in the dataset that is then used as the weights.

Describe alternatives you've considered

I tried to work around this using curvefit, but that doesn't support multidimensional weights / sigma either, as they are passed through kwargs to scipy, and not broadcast on the xarray side. And even if it worked, I would need to manually define a polynomial function and to use polyval, I would have to juggle with the degree dimension.

My current workaround is to manually use apply_ufunc with vectorize=True, but there I ran into #9845.

Additional context

No response