Новое измерение с координатами в xarray apply_ufunc ⇐ Python

[Anonymous]

У меня есть трехмерный массив сигналов с размерами («эксперимент», «испытание», «время»).
Я пытаюсь векторизовать вычисление периодограммы Уэлча для каждого испытания каждого эксперимента. используя xr.apply_ufunc с scipy.signal.welch, но не могу заставить его работать с размерами. scipy.signal.welch возвращает два массива: частоты и спектр PSD/мощности.
Создание случайных данных:

Код: Выделить всё

    data = xr.DataArray(
np.random.random((3, 2, 1024)),
dims=['experiment', 'trial', 'time'],
coords={
'experiment': np.arange(3),
'trial': np.arange(2),
'time': np.arange(1024),
},
)

Теперь применяем scipy.sig.welch для ввода 1D:

Код: Выделить всё

    ret = xr.apply_ufunc(
sig.welch,
data,
input_core_dims=[['trial', 'time']],
output_core_dims=[['frequency'], []],
vectorize=True,
)

выдает

TypeError: только массивы длины 1 могут быть преобразованы в скаляры Python
Вышеупомянутое исключение было непосредственной причиной следующего исключения:
Traceback (последний вызов):
....
numpy/lib /function_base.py", строка 2506, в _vectorize_call_with_signature
вывод[index] = результат

Код: Выделить всё

ValueError: setting an array element with a sequence.

Возможно, векторизация numpy предполагает, что возвращаемые значения будут скалярами?
Поскольку sig.welch может работать с 2D-массивами, еще одна попытка:

Возможно, векторизация numpy предполагает, что возвращаемые значения будут скалярами? р>

Код: Выделить всё

    ret = xr.apply_ufunc(
sig.welch,
data,
input_core_dims=[['trial', 'time']],
output_core_dims=[['trial', 'frequency'], []],
vectorize=True,
)

и это выдает:

... numpy/lib/function_base.py", строка 2050, в _update_dim_sizes
raise ValueError(
ValueError: одномерный аргумент не имеет достаточного количества измерений для всех основных измерений («dim2», «dim0»)

Есть ли способ выполнить векторизацию или необходимо пройти через измерение верхнего уровня?

Подробнее здесь: https://stackoverflow.com/questions/768 ... pply-ufunc