如何在Python Scipy稀疏COO矩阵中保留显式零值并获取其索引

问题描述 #

创建了一个COO矩阵，其data数组中包含显式零值。查询data数组能看到这些零值，但调用nonzero()方法时，这些零值的索引会被过滤掉。如何获取这些显式零值的索引？这是否是COO实现的bug？

复现代码如下，最终断言会失败，因为显式元素数量为7，但nonzero()返回的索引仅6个：

import scipy.sparse as sp

sparse_simple = sp.coo_matrix(
    [
        [1.1, 0, 1.1],
        [0, 1.1, 4.1],
        [1.1, 4.1, 1.1]
    ]
)

sparse_simple_data = sparse_simple.data
sparse_simple_nz = sparse_simple.nonzero()
sparse_simple_data[1] = 0
(n_rows, n_cols) = sparse_simple.shape
sparse_simple_with_explicit_close_to_zero = sp.coo_matrix(
    (sparse_simple_data, (sparse_simple_nz[0], sparse_simple_nz[1])),
    shape=(n_rows, n_cols)
)
num_explicit_vals = len(sparse_simple_with_explicit_close_to_zero.data)
nz_idcs = sparse_simple_with_explicit_close_to_zero.nonzero()
num_nzs = len(nz_idcs[0])

assert num_explicit_vals == num_nzs

当前有个临时hack方案：给data数组所有值加极小数值，但不认可这种方法，希望找到正规解决方案。

解答 #

这不是bug #

Scipy稀疏矩阵的设计初衷是通过仅存储非零元素来节省空间，因此nonzero()、count_nonzero()等方法会自动过滤掉显式存储的零值，这是预期行为，并非实现bug。

正规解决方案 #

COO矩阵的核心结构由三个数组组成：row（所有显式元素的行索引）、col（所有显式元素的列索引）、data（对应元素的值）。这三个数组的长度完全一致，且row和col不会过滤零值，直接利用这两个属性即可获取所有显式元素的索引：

1. 获取所有显式元素的索引（包括零值）

# 获取所有显式元素的行、列索引
all_row_indices = sparse_simple_with_explicit_close_to_zero.row
all_col_indices = sparse_simple_with_explicit_close_to_zero.col

# 验证长度与data一致
assert len(all_row_indices) == len(sparse_simple_with_explicit_close_to_zero.data)

2. 单独筛选显式零值的索引
   通过布尔索引过滤data数组，即可提取出零值对应的行、列索引：

# 创建零值掩码
zero_mask = sparse_simple_with_explicit_close_to_zero.data == 0

# 获取显式零值的行、列索引
zero_row_indices = sparse_simple_with_explicit_close_to_zero.row[zero_mask]
zero_col_indices = sparse_simple_with_explicit_close_to_zero.col[zero_mask]

这种方法无需修改矩阵的实际值，完全符合COO矩阵的设计规范，是替代hack方案的正规做法。

内容的提问来源于stack exchange，提问作者JStrahl