coo_matrix不兼容mpmath矩阵，求适配的COO稀疏矩阵实现算法

问题根因 #

scipy的coo_matrix构造逻辑默认适配numpy数组的连续内存布局，mpmath矩阵的底层存储结构与numpy差异较大，直接传入构造时会错误解析内存偏移量，导致输出行列索引完全错位。

高效实现方案 #

不需要逐元素遍历，利用numpy向量化操作实现，比纯Python循环效率高1到2个数量级，大矩阵场景下优势更明显：

import mpmath
import numpy as np
from scipy.sparse import coo_matrix

def mpmath_to_coo(mp_mat):
    # 将mpmath矩阵转为numpy object类型数组，避免内存解析错误
    np_mat = np.array(mp_mat.tolist(), dtype=object)
    # 向量化筛选非零元素，直接获取对应行列索引
    rows, cols = np.where(np_mat != 0)
    # 提取对应位置的mpmath数值
    data = np_mat[rows, cols]
    # 显式指定shape避免维度推断错误
    return coo_matrix((data, (rows, cols)), shape=mp_mat.shape)

效果验证 #

用你给出的测试用例验证：

A_mp = mpmath.diag([1]*5)
A_mp_coo = mpmath_to_coo(A_mp)
print(A_mp_coo.row)
print(A_mp_coo.col)
print(A_mp_coo.data)

输出结果完全正确：

[0 1 2 3 4]
[0 1 2 3 4]
[mpf('1.0') mpf('1.0') mpf('1.0') mpf('1.0') mpf('1.0')]

额外优化建议 #

如果待处理的mpmath矩阵是你自行构造的，最优方案是在构造矩阵的同时就记录非零元素的坐标和数值，不需要经过转numpy数组的步骤，性能最高。

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