如何加速大稀疏矩阵中的digamma计算？

Question

我有两个大的numpy矩阵X和r（约5500x3500）。我需要计算digamma(X+r)和digamma(r)，我使用scipy.special.psi进行计算。计算发生在我通过梯度下降拟合的渐变函数中。在优化期间，渐变函数将被调用很多次，因此它必须快速运行。现在，到实际问题：

• digamma(r)需要大约11秒才能运行
• digamma(X+r)需要约50秒才能运行

出于好奇，我跑了一些其他的例子：

• digamma(X.astype(np.float64))，需要〜1分钟才能运行
• digamma(X)，需要〜5-6分钟才能运行

为什么会发生这种情况，我该怎么做才能改善运行时间？

• X包含从0到164的整数（numpy.int64）。
• r包含从0.到~4.6的浮点数（numpy.float64）。
• X非常稀疏（约75％为0）。
• r不稀疏（约10％为0）

编辑：更正X.as_float