Python调整3D数组大小展平时遇广播错误的问题求助

问题根源 #

你用pivot处理后得到的是形状为(2,1)的object数组，每个单元格单独存储一个(200,)的psth数组，导致最终嵌套结构为(3,2,1)，而非预期的(2,200)。后续sum()报错是因为你试图把(200,)的数组（或其求和结果）塞进原数组的(1,)位置，触发了广播不匹配。

修正方案 #

不需要使用pivot，直接对每个trial分组后的psth列提取数组，用np.stack堆叠成(2,200)的结构即可：

import numpy as np
import pandas as pd

data = {
    'trial_id': [1, 1, 2, 2, 3, 3],
    'unit': [1, 2, 1, 2, 1, 2],
    'psth': [np.random.rand(200) for _ in range(6)]
}

AE_0_df = pd.DataFrame(data)

trial_data = []

for trial_id, group in AE_0_df.groupby('trial_id'):
    # 按unit排序确保unit1、unit2的顺序固定
    sorted_group = group.sort_values(by='unit')
    # 将两个(200,)的psth数组堆叠成(2,200)的数组
    unit_psth = np.stack(sorted_group['psth'].values)
    trial_data.append(unit_psth)

# 转换为numpy数组后形状就是(3,2,200)
final_array = np.array(trial_data)
print(final_array.shape)  # 输出: (3, 2, 200)

额外说明 #

如果你的需求是对每个unit的psth数组求和（得到标量），可以在堆叠后直接指定维度求和：

# 对每个unit的200个数据求和，得到(3,2)的结果
summed_data = np.array(trial_data).sum(axis=2)
print(summed_data.shape)  # 输出: (3, 2)

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