只有效地计算对称外积矩阵的唯一元素

Question

我有一大组N个d维向量(驻留在矩阵中)，我通过取自身外积(即每个向量I乘以自身)来提升这些向量。对于每个向量，这产生一个对称矩阵，其中(d+1)选择2个唯一条目。对于整个数据，这是一个N x d x d张量。我想只计算唯一的(d+1)，从每个张量切片的下对角线选择2个条目，并将它们存储在向量中。我希望在Python中使用尽可能少的内存占用和尽可能快的速度来实现这一点，包括使用C绑定。

如果使用标准的numpy方法，它会分配每个矩阵的整体。这大约是实际所需内存复杂性的两倍。

对于这里的尺度感，考虑N=20k和d=20k的情况。则每个元素N*d^2*~8字节=(2*10^4)^3*8字节=64 TB

如果我们只计算对唯一条目进行编码的向量，则我们有(20001选择2)*20k*8=200001000*2000*8字节=32 TB

有没有一种快速的方法可以做到这一点，而不必求助于缓慢的方法(比如用python编码我自己的外部产品)？

编辑：我会注意到，在numpy 中创建外部产品数组时也提出了类似的问题

我已经知道如何使用einsum来计算这个(就像上面的问题一样)。然而，如果没有额外的(d选择2)计算和分配，就无法得出这样做的答案

编辑2：这个线程如何在Numpy(或其他Python解决方案)中利用外部产品的对称性？问了一个相关的问题，但没有解决内存的复杂性。最重要的答案仍然是为每个外部乘积分配一个d x d数组。

这个线程Numpy Performance-向量的外积及其转置也解决了自外积的计算问题，但没有达到内存高效的解决方案。

编辑3：如果要分配整个数组，然后提取元素，np.tril_indices或scipy.spatial.distance.squareform就可以了。

Answer 1

不确定您想要的输出方式，但始终可以选择使用Numba:

import numpy as np
import numba as nb
# Computes unique pairwise products
@nb.njit(parallel=True)
def self_outer_unique(a):
n, d = a.shape
out = np.empty((n, (d * d + d) // 2), dtype=a.dtype)
for i in nb.prange(n):
for j1 in range(d):
for j2 in range(j1, d):
idx = j1 * (2 * d - j1 + 1) // 2 + j2 - j1
out[i, idx] = a[i, j1] * a[i, j2]
return out

这将为您提供一个数组，每行上都有所有唯一的乘积(即完整输出的扁平上三角形)。

import numpy as np
a = np.arange(12).reshape(4, 3)
print(a)
# [[ 0  1  2]
#  [ 3  4  5]
#  [ 6  7  8]
#  [ 9 10 11]]
print(self_outer_unique(a))
# [[  0   0   0   1   2   4]
#  [  9  12  15  16  20  25]
#  [ 36  42  48  49  56  64]
#  [ 81  90  99 100 110 121]]

就性能而言，它比使用NumPy计算完整的外积更快，尽管从中重新创建完整的数组需要更长的时间。

import numpy as np
def np_self_outer(a):
return a[:, :, np.newaxis] * a[:, np.newaxis, :]
def my_self_outer(a):
b = self_outer_unique(a)
n, d = a.shape
b_full = np.zeros((n, d, d), dtype=a.dtype)
idx0 = np.arange(n)[:, np.newaxis]
idx1, idx2 = np.triu_indices(d)
b_full[idx0, idx1, idx2] = b
b_full += np.triu(b_full, 1).transpose(0, 2, 1)
return b_full
n, d = 1000, 100
a = np.arange(n * d).reshape(n, d)
print(np.all(np_self_outer(a) == my_self_outer(a)))
# True
%timeit np_self_outer(a)
# 24.6 ms ± 248 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit self_outer_unique(a)
# 6.32 ms ± 69.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit my_self_outer(a)
# 124 ms ± 770 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)