在 xarray apply_ufunc 中坐标的新维度

Question

我有一个 3D 信号数组，其尺寸为

('experiment', 'trial', 'time')

。 我正在尝试使用

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 进行一维输入：

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

投掷

TypeError：只有长度为 1 的数组可以转换为 Python 标量

上述异常是导致以下异常的直接原因：

回溯（最近一次调用最后一次）： .... numpy/lib/function_base.py”，第 2506 行，在 _vectorize_call_with_signature 中 输出[索引] = 结果

ValueError: setting an array element with a sequence.

numpy 向量化可能期望返回值是标量？ 由于 sig.welch 可以处理 2D 数组，因此另一次尝试：

    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 中 引发值错误（ ValueError：一维参数没有足够的维度用于所有核心维度（'dim2'，'dim0'）

有没有一种方法可以进行矢量化，或者必须循环遍历顶级维度？

Answer 1

问题在于

scipy.signal.welch

返回两个不同维度的数组。 因此，解决方法可以是编写包装函数：

def wrapper(v, **kwargs):
    _, psd = signal.welch(v, **kwargs)
    return psd

仅返回功率谱密度并忽略频率。 那么向量化就可以是

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

那么

ret

的尺寸将为

['experiment', 'trial', 'frequency']

，因为

wrapper

只会将

'time'

更改为

'frequency'

尺寸。 您可以在此之后将坐标值分配给

frequency

维度：

f, psd = sig.welch(data[0,0])
ret['frequency'] = f