python - 根据索引对多维数组中的元素求和
问题描述
我正在处理一个非常大的多维数据,但让我以二维数组为例。给定一个每次迭代都会改变的值数组,
arr = np.array([[ 1, 2, 3, 4, 5], [5, 6, 7, 8, 9]]) # a*b
和一个一直固定的索引数组。
idx = np.array([[[0, 1, 1], [-1, -1, -1]],
[[5, 1, 3], [1, -1, -1]]]) # n*h*w, where n = a*b,
这里 -1 表示不应用索引。我希望得到一个结果
res = np.array([[1+2+2, 0],
[5+2+4, 2]]) # h*w
在实际实践中,我正在处理一个非常大的 3D 张量(n ~ 万亿),具有非常稀疏的 idx(即很多 -1)。已修复,我idx
目前的解决方案是通过填充 0 和 1 预先计算一个 *(h*w) 数组index_tensor
,然后执行
tmp = arr.reshape(1, n)
res = (tmp @ index_tensor).reshape([h,w])
它工作正常,但需要大量内存来存储index_tensor
. 有没有什么方法可以利用它的稀疏性和不可更改性idx
来降低内存成本并在 python 中保持公平的运行速度(使用 numpy 或 pytorch 会是最好的)?提前致谢!
解决方案
暂时忽略-1
复杂性,直接的索引和求和是:
In [58]: arr = np.array([[ 1, 2, 3, 4, 5], [5, 6, 7, 8, 9]])
In [59]: idx = np.array([[[0, 1, 1], [2, 4, 6]],
...: [[5, 1, 3], [1, -1, -1]]])
In [60]: arr.flat[idx]
Out[60]:
array([[[1, 2, 2],
[3, 5, 6]],
[[5, 2, 4],
[2, 9, 9]]])
In [61]: _.sum(axis=-1)
Out[61]:
array([[ 5, 14],
[11, 20]])
处理 -1 的一种方法(不一定快速或内存高效)是使用掩码数组:
In [62]: mask = idx<0
In [63]: mask
Out[63]:
array([[[False, False, False],
[False, False, False]],
[[False, False, False],
[False, True, True]]])
In [65]: ma = np.ma.masked_array(Out[60],mask)
In [67]: ma
Out[67]:
masked_array(
data=[[[1, 2, 2],
[3, 5, 6]],
[[5, 2, 4],
[2, --, --]]],
mask=[[[False, False, False],
[False, False, False]],
[[False, False, False],
[False, True, True]]],
fill_value=999999)
In [68]: ma.sum(axis=-1)
Out[68]:
masked_array(
data=[[5, 14],
[11, 2]],
mask=[[False, False],
[False, False]],
fill_value=999999)
掩码数组通过将掩码值替换为中性值来处理求和之类的操作,例如求和的情况下为 0。
(我可能会在早上重温这个)。
与矩阵积求和
In [72]: np.einsum('ijk,ijk->ij',Out[60],~mask)
Out[72]:
array([[ 5, 14],
[11, 2]])
这比掩码数组方法更直接、更快。
您还没有详细说明构建,index_tensor
所以我不会尝试比较它。
另一种可能性是用 0 填充数组,并调整索引:
In [83]: arr1 = np.hstack((0,arr.ravel()))
In [84]: arr1
Out[84]: array([0, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9])
In [85]: arr1[idx+1]
Out[85]:
array([[[1, 2, 2],
[3, 5, 6]],
[[5, 2, 4],
[2, 0, 0]]])
In [86]: arr1[idx+1].sum(axis=-1)
Out[86]:
array([[ 5, 14],
[11, 2]])
疏
第一次尝试使用稀疏矩阵:
重塑idx
为 2d:
In [141]: idx1 = np.reshape(idx,(4,3))
从中制作一个稀疏张量。首先,我将采用迭代lil
方法,尽管通常直接构建coo
(甚至csr
)输入会更快:
In [142]: M = sparse.lil_matrix((4,10),dtype=int)
...: for i in range(4):
...: for j in range(3):
...: v = idx1[i,j]
...: if v>=0:
...: M[i,v] = 1
...:
In [143]: M
Out[143]:
<4x10 sparse matrix of type '<class 'numpy.int64'>'
with 9 stored elements in List of Lists format>
In [144]: M.A
Out[144]:
array([[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 1, 0, 1, 0, 0, 0],
[0, 1, 0, 1, 0, 1, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0]])
然后可以将其用于产品总和:
In [145]: M@arr.ravel()
Out[145]: array([ 3, 14, 11, 2])
使用M.A@arr.ravel()
本质上就是你所做的。虽然M
是稀疏的,arr
但不是。对于这种情况M.A@
,比M@
.
推荐阅读
- python - 数据预处理:keras 层 vs tf.image.resize
- sql - 使用 bigrquery 的“query_exec”和“bq_table_download”之间的差异
- sql-server - SQL Server QUOTENAME 命令的替代品,没有 128 个字符的限制。如何在括号之间放置一个字符串
- openlayers - 使用 OpenLayers 放大单个点簇
- vue.js - Vue 3 的 Monaco Editor Web Worker 问题
- typescript - 如何将 TypeScript 可变元组类型用于笛卡尔积函数?
- excel - 使用 VBA 循环遍历 Excel 中的精确工作表列表
- excel - 在 Excel 报告中分析不适用于最终用户
- flutter - 创建具有两个嵌套循环的 Flutter 小部件
- grafana - 如何将变量(日期列表)默认为当前日期?