Попытка применить fftconvolve к xarray

Question

У меня есть 3D xarrayDataArray и 1D xarrayDataArray, Одно измерение трехмерного массива имеет тот же размер, что и одно измерение другого. Я хочу подать заявку scipy.signal.fftconvolve вдоль этих соответствующих размеров. Но я получаю ошибку. Вот небольшой пример кода:

import numpy as np
from scipy.signal import fftconvolve
import xarray as xr

xarr1 = xr.DataArray(np.random.random([10,20,500]), 
                     dims=('dim1', 'dim2', 'sample'))
xarr2 = xr.DataArray(np.random.random(500), 
                     dims=('sample',))
res = xr.apply_ufunc(fftconvolve, xarr1, xarr2, 
                     input_core_dims=[['sample'], ['sample']], 
                     kwargs={'mode': 'same'}, 
                     vectorize=True)

То, что я хочу сделать, является эквивалентом этого:

res = np.zeros(xarr1.shape)
for i1 in range(xarr1.shape[0]):
    for i2 in range(xarr1.shape[1]):
       res[i1, i2, :] = fftconvolve(xarr1[i1, i2, :], xarr2, mode='same')

Но я получаю следующую ошибку с apply_ufunc версия:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-174-4f45f68ab252> in <module>()
    7 xarr2 = xr.DataArray(np.random.random(500), 
    8                      dims=('sample',))
----> 9 res = xr.apply_ufunc(fftconvolve, xarr1, xarr2, input_core_dims=[['sample'], ['sample']], kwargs={'mode': 'same'}, vectorize=True)

~\Anaconda3\envs\dataproc\lib\site-packages\xarray\core\computation.py in apply_ufunc(func, *args, **kwargs)
    932                                      join=join,
    933                                      exclude_dims=exclude_dims,
--> 934                                      keep_attrs=keep_attrs)
    935     elif any(isinstance(a, Variable) for a in args):
    936         return variables_ufunc(*args)

~\Anaconda3\envs\dataproc\lib\site-packages\xarray\core\computation.py in apply_dataarray_ufunc(func, *args, **kwargs)
    209 
    210     data_vars = [getattr(a, 'variable', a) for a in args]
--> 211     result_var = func(*data_vars)
    212 
    213     if signature.num_outputs > 1:

~\Anaconda3\envs\dataproc\lib\site-packages\xarray\core\computation.py in apply_variable_ufunc(func, *args, **kwargs)
    563             raise ValueError('unknown setting for dask array handling in '
    564                              'apply_ufunc: {}'.format(dask))
--> 565     result_data = func(*input_data)
    566 
    567     if signature.num_outputs > 1:

~\Anaconda3\envs\dataproc\lib\site-packages\numpy\lib\function_base.py in __call__(self, *args, **kwargs)
2737             vargs.extend([kwargs[_n] for _n in names])
2738 
-> 2739         return self._vectorize_call(func=func, args=vargs)
2740 
2741     def _get_ufunc_and_otypes(self, func, args):

~\Anaconda3\envs\dataproc\lib\site-packages\numpy\lib\function_base.py in _vectorize_call(self, func, args)
2803         """Vectorized call to `func` over positional `args`."""
2804         if self.signature is not None:
-> 2805             res = self._vectorize_call_with_signature(func, args)
2806         elif not args:
2807             res = func()

~\Anaconda3\envs\dataproc\lib\site-packages\numpy\lib\function_base.py in _vectorize_call_with_signature(self, func, args)
2867 
2868             for output, result in zip(outputs, results):
-> 2869                 output[index] = result
2870 
2871         if outputs is None:

ValueError: setting an array element with a sequence.

Answer 1

Определение output_core_dims решает это,

res = xr.apply_ufunc(fftconvolve, xarr1, xarr2, 
                     input_core_dims=[['sample'], ['sample']], 
                     output_core_dims=[['sample']],
                     kwargs={'mode': 'same'}, vectorize=True)