多分组特征下子集数据均值的高效向量化计算方案咨询

需求背景 #

基于两个分组特征（分别包含4个和5个分组），计算每个(i,j)组合分组对应的observables数据均值。

非向量化实现代码 #

import numpy as np

N = 110
observables = np.random.rand(N)

I = np.random.randint(np.array([0,0])[:,None], np.array([4,5])[:,None], 
                     size=(2,N))

# 可将I理解为对数据的两种分组标记，一组有4个分组，另一组有5个分组
# 需要计算所有i（第一组索引）和j（第二组索引）组合对应的均值
averages = np.zeros((4,5))

# 非向量化实现
for i in range(4):
 for j in range(5):
   J = np.argwhere((I[0,:]==i) & (I[1,:]==j)).flatten()
   averages[i,j] = np.mean(observables[J])

尝试的低效向量化实现 #

当N较大（如N=1000）时，因最小公倍数（LCM）导致数据量激增，效率极低：

import math

# 低效的向量化实现
J = [np.argwhere((I[0,:]==i) & (I[1,:]==j)).flatten() for i in range(4)
     for j in range(5)]
lengths = [len(j) for j in J]

L = math.lcm(*lengths)

J = np.array([np.tile(j,int(L/len(j))) for j in J])

averages_vectorized = np.reshape(np.mean(observables[J], axis=1), (4,5))

其他尝试（本质仍含循环） #

列表推导式实现，本质是隐藏的循环：

J = [np.argwhere((I[0,:]==i) & (I[1,:]==j)).flatten() for i in range(4)
     for j in range(5)]

averages_list_comrehension = np.reshape(np.array([np.mean(observables[j]) for j in J]), (4,5))

待验证思路 #

曾考虑在observables末尾添加nan值，扩展索引至统一长度后用np.nanmean计算，但需适配PyTensor张量场景，不确定PyTensor是否支持nan张量及nanmean操作。

高效向量化解决方案 #

方案1：利用np.bincount实现分组统计 #

将二维的(i,j)分组转换为一维唯一索引，通过bincount分别计算每组总和与数量，最终相除得到均值：

# 将(i,j)转换为唯一的一维索引：i*5 + j（第二组共5个分组）
group_ids = I[0] * 5 + I[1]

# 计算每组的观测值总和
sum_vals = np.bincount(group_ids, weights=observables, minlength=4*5)
# 计算每组的元素数量
counts = np.bincount(group_ids, minlength=4*5)

# 处理空分组（避免除以0），若需保留空值为nan可跳过此步
counts[counts == 0] = 1
averages = sum_vals / counts

# 重塑为(4,5)的结果矩阵
averages = averages.reshape(4,5)

该方法完全向量化，无显式/隐式循环，时间复杂度O(N)，大N场景下效率极高。

方案2：利用np.add.at进行累加 #

通过广播索引直接对结果矩阵累加观测值和计数，再计算均值：

averages = np.zeros((4,5))
counts = np.zeros((4,5))

# 按分组索引累加观测值和计数
np.add.at(averages, (I[0], I[1]), observables)
np.add.at(counts, (I[0], I[1]), 1)

# 计算均值，自动处理空分组（空分组结果为nan）
averages = np.divide(averages, counts, where=counts!=0)

此方法同样完全向量化，逻辑直观，性能与bincount方案接近。

针对PyTensor场景的适配 #

PyTensor支持bincount操作，无需依赖nan值处理，适配代码如下：

import pytensor.tensor as pt

# 假设observables和I为PyTensor张量
group_ids = I[0] * 5 + I[1]
sum_vals = pt.bincount(group_ids, weights=observables, minlength=20)
counts = pt.bincount(group_ids, minlength=20)

# 处理除以0，空分组设为0（可根据需求调整）
averages = pt.where(counts != 0, sum_vals / counts, pt.zeros_like(sum_vals))
averages = averages.reshape((4,5))

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