从 scipy CSR 矩阵索引到 numpy 数组的最有效方法？

Question

我有一个形状为 (4000, 3) 的 numpy ndarray X，其中 X 中的每个样本都是一个 3D 坐标 (x,y,z)。

我有一个形状为 (4000, 4000) 的 scipy csr 矩阵 nn_rad_csr，它是从 sklearn.neighbors.radius_neighbors_graph(X, 0.01, include_self=True).

生成的最近邻图

nn_rad_csr.toarray()[i] 是一个形状为 (4000,) 的稀疏向量，其二进制权重（0 或 1）与节点 X[i] 的最近邻图中的边相关联。

例如，如果 nn_rad_csr.toarray()[i][j] == 1 则 X[j] 在 X[i] 的最近邻居半径内，而值 0 表示它不是邻居。

我想要做的是有一个函数 radius_graph_conv(X, rad)，它 returns 一个数组 Y，它是 X，由它的邻居的值平均。我不确定如何利用 CSR 矩阵的稀疏性来有效地执行 radius_graph_conv。我在下面有两个 graph conv 的简单实现。

import numpy as np
from sklearn.neighbors import radius_neighbors_graph, KDTree

def radius_graph_conv(X, rad):
    nn_rad_csr = radius_neighbors_graph(X, rad, include_self=True)
    csr_indices = nn_rad_csr.indices
    csr_indptr  = nn_rad_csr.indptr
    Y = np.copy(X)
    for i in range(X.shape[0]):
        j, k = csr_indptr[i], csr_indptr[i+1]
        neighbor_idx = csr_indices[j:k]
        rad_neighborhood = X[neighbor_idx] # ndim always 2
        Y[i] = np.mean(rad_neighborhood, axis=0)
    return Y

def radius_graph_conv_matmul(X, rad):
    nn_rad_arr = radius_neighbors_graph(X, rad, include_self=True).toarray()
    # np.sum(nn_rad_arr, axis=-1) is basically a count of neighbors

    return np.matmul(nn_rad_arr / np.sum(nn_rad_arr, axis=-1), X)

有更好的方法吗？对于 knn 图，它是一个非常简单的函数，因为邻居的数量是固定的，你可以只索引到 X，但是对于基于半径或密度的最近邻图，你必须使用 CSR，（或数组数组，如果你使用的是 kd 树）。

Answer 1

这是利用csr格式的直接方法。您的 matmul 解决方案可能在幕后做类似的事情。但是我们通过利用它是一个邻接矩阵来保存一次查找（来自 .data 属性）；此外，diffing .indptr 应该比将等量的 1 相加更有效率。

>>> import numpy as np
>>> from scipy import sparse
>>> 
# create mock data
>>> A = np.random.random((100, 100)) < 0.1
>>> A = (A | A.T).view(np.uint8)
>>> AS = sparse.csr_matrix(A)
>>> X = np.random.random((100, 3))
>>> 
# dense solution for reference
>>> Xa = A @ X / A.sum(axis=-1, keepdims=True)
# sparse solution
>>> XaS = np.add.reduceat(X[AS.indices], AS.indptr[:-1], axis=0) / np.diff(AS.indptr)[:, None]
>>> 
# check they are the same
>>> np.allclose(Xa, XaS)
True

从 scipy CSR 矩阵索引到 numpy 数组的最有效方法？

Most efficient way to index into a numpy array from a scipy CSR matrix?

python

numpy

nearest-neighbor

scipy

scikit-learn