是否有一种节省内存的方法可以将零列和行插入到 numpy 数组中?

Is there a memory-efficient way of inserting zero-columns and rows into a numpy array?

我有一个大的(对称的)numpy 矩阵 arr。它的形状是arr.shape = (50_000, 50_000)。我想插入一些 zero-rows/columns (以对称方式)。假设我要插入的行/列的数量可能是 123.

小例子

import numpy as np

# Create symmetric square matrix
size = 3
arr = np.array(list(range(size**2))).reshape(size, size)
arr = arr + arr.T

# insert zero columns
# Each "1" represents a column from the original matrix, e.g.
# the first 1 is the first column of arr, the second 1 the second column of arr
# and so on
insert_cols = [1, 0, 0, 1, 0, 1, 0, 0]

# insert the zero rows / columns
current_index = 0
for col in insert_cols:
    if col == 0:
        arr = np.insert(arr, current_index, 0, axis=0)
        arr = np.insert(arr, current_index, 0, axis=1)
    current_index += 1

print(arr)

如果我np.insert理解正确,那么这段代码会创建一个数组副本,并一直复制内容。

问题

我认为使用 sparse matrix classes 之一可能会更简单/更有效?还有其他方法可以提高效率吗?

给定 insert_cols,我们可以这样做 -

n = len(insert_cols)
out = np.zeros((n,n),arr.dtype)
idx = np.flatnonzero(insert_cols)
out[np.ix_(idx,idx)] = arr # or out[idx[:,None],idx] = arr

或者,使用布尔版本进行索引。因此-

insert_cols_bool = np.asarray(insert_cols, dtype=bool)

然后,使用 insert_cols_bool 代替 idx


稀疏矩阵

为了提高内存效率,我们可以将输出存储为稀疏矩阵 -

from scipy.sparse import coo_matrix

l = len(idx)
r,c = np.broadcast_to(idx[:,None],(l,l)).ravel(),np.broadcast_to(idx,(l,l)).ravel()
out = coo_matrix((arr.ravel(), (r,c)), shape=(n,n))