使用xr.apply_ufunc结合scipy.signal.welch处理多维数组报错求助

问题背景 #

需要对维度为('experiment', 'trial', 'time')的3D xarray数据，为每个实验（experiment）的每个试次（trial）单独计算Welch周期图，scipy.signal.welch会返回频率数组和功率谱密度（PSD）数组两个结果，之前的尝试因维度配置错误触发报错。

正确实现方法 #

不需要手动遍历顶层维度，通过正确配置xr.apply_ufunc的核心参数即可完成向量化处理：

步骤1：生成测试数据 #

import numpy as np
import xarray as xr
from scipy import signal as sig

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),
    },
)

步骤2：正确调用xr.apply_ufunc #

# 调用welch并获取频率和PSD结果
freqs, psds = xr.apply_ufunc(
    sig.welch,
    data,
    # 指定输入的核心维度：对每个样本的time维度单独处理
    input_core_dims=[['time']],
    # 指定输出的核心维度：两个返回值都新增frequency维度
    output_core_dims=[['frequency'], ['frequency']],
    # 指定输出数据类型，避免类型推断问题
    output_dtypes=[float, float],
    # 传递welch的额外参数，比如采样率fs（根据实际数据调整）
    kwargs={'fs': 1.0}
)

结果验证 #

• freqs的维度为('frequency',)：所有实验和试次的频率轴一致（welch参数相同时）
• psds的维度为('experiment', 'trial', 'frequency')：对应每个实验-试次组合的功率谱结果

之前错误的原因分析 #

1. 第一次尝试错误：

   • 错误将input_core_dims设为[['trial', 'time']]，导致把每个实验下的所有试次+时间序列作为2D输入传给welch，不符合“每个试次单独处理”的需求
   • 启用vectorize=True会让numpy尝试将welch的数组输出视为标量，触发ValueError: setting an array element with a sequence

2. 第二次尝试错误：

   • output_core_dims配置错误，welch对1D时间序列返回的是1D频率数组，不需要包含trial维度，因此[['trial', 'frequency'], []]的配置会导致维度不匹配报错

内容的提问来源于stack exchange，提问作者subhacom