将 scipy 稀疏矩阵与 3d numpy 数组相乘

Question

我有以下矩阵

a = sp.random(150, 150)
x = np.random.normal(0, 1, size=(150, 20))

我基本上想实现以下公式

我可以这样计算内部差异

diff = (x[:, None, :] - x[None, :, :]) ** 2
diff.shape  # -> (150, 150, 20)

a.shape  # -> (150, 150)

我基本上想广播 scipy 稀疏矩阵和每个内部 numpy 数组之间的逐元素乘法。

如果 A 可以很稠密，那么我可以简单地做

np.einsum("ij,ijk->k", a.toarray(), (x[:, None, :] - x[None, :, :]) ** 2)

但 A 很稀疏，而且可能很大，所以这不是一个选择。当然，我可以重新排序轴并diff使用 for 循环遍历数组，但是使用 numpy 是否有更快的方法？

正如 @hpaulj 指出的，当前的解决方案还形成了一个 shape 数组(150, 150, 20)，这也会立即导致内存问题，所以这个解决方案也不好。

Answer 1

import numpy as np
import scipy.sparse
from numpy.random import default_rng

rand = default_rng(seed=0)

# \sigma_k = \sum_i^N \sum_j^N A_{i,j} (x_{i,k} - x_{j,k})^2

# Dense method
N = 100
x = rand.integers(0, 10, (N, 2))
A = np.clip(rand.integers(0, 100, (N, N)) - 80, a_min=0, a_max=None)
diff = (x[:, None, :] - x[None, :, :])**2
product = np.einsum("ij,ijk->k", A, diff)

# Loop method
s_loop = [0, 0]
for i in range(N):
    for j in range(N):
        for k in range(2):
            s_loop[k] += A[i, j]*(x[i, k] - x[j, k])**2
assert np.allclose(product, s_loop)

# For any i,j, we trivially know whether A_{i,j} is zero, and highly sparse matrices have more zeros
# than nonzeros. Crucially, do not calculate (x_{i,k} - x_{j,k})^2 at all if A_{i,j} is zero.
A_i_nz, A_j_nz = A.nonzero()
diff = (x[A_i_nz, :] - x[A_j_nz, :])**2
s_semidense = A[A_i_nz, A_j_nz].dot(diff)
assert np.allclose(product, s_semidense)

# You can see where this is going:
A_sparse = scipy.sparse.coo_array(A)
diff = (x[A_sparse.row, :] - x[A_sparse.col, :])**2
s_sparse = A_sparse.data.dot(diff)
assert np.allclose(product, s_sparse)

看起来相当快；这大约在一秒钟内完成：

N = 100_000_000
print(f'Big test: initialising a {N}x{N} array')
n_values = 10_000_000
A = scipy.sparse.coo_array(
    (
        rand.integers(0, 100, n_values),
        rand.integers(0, N, (2, n_values)),
    ),
    shape=(N, N),
)
x = rand.integers(0, 100, (N, 2))

print('Big test: calculating')
s = A.data.dot((x[A.row, :] - x[A.col, :])**2)

print(s)