Как применить xarray u_function над NetCDF и вернуть 2D-массив (несколько новых переменных) в DataSet

Question

Я пытаюсь использовать xarray «apply_ufunc», чтобы применить данную функцию «f» ко всем парам координат (то есть: пикселей) в наборе данных.

Функция «f» возвращает двумерный массив(Матрица NxN) как результат. Поэтому результирующий DataSet будет иметь несколько новых переменных после анализа: всего «M» новых переменных.

Функция «f» работает просто отлично. Таким образом, ошибка, похоже, не исходит от нее.

Возможной проблемой может быть структура, которую двумерный массив возвращает из "f". Насколько я понимаю, xarray apply_ufunc требует, чтобы результирующий массив был структурирован в кортежи. Итак, я даже пытался преобразовать 2D-массив в набор массивов, но пока ничего не получалось.

Ситуацию можно проверить в других работах Работы , а также. В этой настоящей ссылке автор должен дважды запустить одну и ту же функцию подгонки линейной регрессии над исходным набором данных, чтобы извлечь все параметры из регрессии (бета_0 и альфа).

Поэтому я хотел бы знать,, способен ли xarray apply_ufunc работать с функциями сокращения, как в приведенной выше ссылке (или в фрагменте кода ниже), которая возвращает несколько новых переменных.

Ниже я представляю воспроизводимый код, включающий обсуждаемую проблему. Обратите внимание, что функция «f» возвращает 2D-массив. Глубина второго измерения равна 4. Поэтому я ожидаю, что после всей обработки у меня будет результирующий DataSet с 4 новыми переменными.

import numpy as np
import xarray as xr


x_size = 10
y_size = 10
time_size = 30

lon = np.arange(50, 50+x_size)

lat = np.arange(10, 10+y_size)

time = np.arange(10, 10+time_size)

array = np.random.randn(y_size, x_size, time_size)

ds = xr.DataArray(data=array, coords = {'lon':lon, 'lat':lat, 'time':time}, dims=('lon', 'lat', 'time'))

ds.isel({'time':0}).plot(x='lon', y='lat')

def f (x):


    return (x, x**2, x**3, x**4)


def f_xarray(ds, dim=['time'], 
              dask='allowed', 
              new_dim_name=['predicted']):   


    filtered = xr.apply_ufunc(f,
                                  ds,
                                      dask=dask,
                                      vectorize=True,
                                      input_core_dims=[dim],
                                      #exclude_dims = dim, # This must not be setted.
                                      output_core_dims= [['x', 'x2', 'x3', 'x4']], #[new_dim_name],
                                      #kwargs=kwargs,
                                      #output_dtypes=[float],
                                      #dataset_join='outer',
                                      #dataset_fill_value=np.nan,
                                      ).compute()

    return filtered


ds2 = f_xarray(ds)

# Error message returned: 

   # ValueError: wrong number of outputs from pyfunc: expected 1, got 4