Python实现3D数组分箱统计对应粒子数的效率问题咨询

核心是用NumPy内置的向量化直方图函数替代手动遍历，利用底层C实现的运算逻辑大幅提升效率，8e6规模的样本运算耗时仅需几毫秒。

核心代码实现 #

import numpy as np

# 替换为你的实际3D速度数组、粒子数组
vel_3d = np.random.rand(200,200,200) * 80  # 示例速度数组，取值范围0~80
particle_3d = np.random.randint(1, 12, size=(200,200,200))  # 示例粒子数数组

# 直接计算加权直方图，weights参数传入对应位置的粒子数，自动完成分箱求和
bin_particle_sum, bin_edges = np.histogram(vel_3d, bins=50, weights=particle_3d)

返回值说明 #

• bin_edges：长度为51的一维数组，记录50个分箱的左右边界，第i个分箱的区间为[bin_edges[i], bin_edges[i+1])，最后一个分箱为闭区间
• bin_particle_sum：长度为50的一维数组，第i个元素对应第i个速度分箱内的粒子总数量

可选扩展 #

如果需要自定义分箱边界而非按速度最大最小值自动划分，可自行指定bins参数：

# 示例：自定义速度分箱范围为0~100，生成51个边界点对应50个等宽分箱
custom_bins = np.linspace(0, 100, 51)
bin_particle_sum, bin_edges = np.histogram(vel_3d, bins=custom_bins, weights=particle_3d)

如果需要做更多分箱维度的统计，可通过np.digitize获取每个元素的分箱索引后自行计算：

# 获取每个速度值对应的分箱索引
bin_idx = np.digitize(vel_3d, bins=bin_edges)
# 示例：统计每个分箱的平均速度
bin_avg_vel = [vel_3d[bin_idx == i].mean() for i in range(1, len(bin_edges))]

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