使用 Numba 支持组装块矩阵答案

【问题标题】：Assembling block matrix with Numba support使用 Numba 支持组装块矩阵
【发布时间】：2019-08-19 16:20:33
【问题描述】：

在我的代码（用 Python 2.7 编写）中，我创建了两个 numpy 数组，A 和 B。然后我用它们组装一个更大的矩阵H，代码如下

H = np.block([A, B], [-B, -A])

随后进行了各种计算，涉及大量的 numpy 操作和 for 循环。因此，我想使用 Numba 来优化代码。但是，Numba 似乎不支持 numpy 块功能。矩阵A 和B 不是很大，所以我可以使用可能不如np.block 优化的函数，但我仍然想以块矩阵方式组装H可读性的目的。有没有什么功能可以做到这一点？

【问题讨论】：

block的工作核心在np.core.shape_base._block。它在每个轴上执行递归 _block 和 concatenate。在您的情况下，这是 3 个连接。 concatenate 已经是编译代码了。

标签： python numpy numba

【解决方案1】：

只是为了使@hpaulj 的评论具体化，对您的输入数据做出一些基本假设，而不进行任何类型的错误检查：

@nb.njit
def nb_block(X):
    xtmp1 = np.concatenate(X[0], axis=1)
    xtmp2 = np.concatenate(X[1], axis=1)
    return np.concatenate((xtmp1, xtmp2), axis=0)

以下也有效：

@nb.njit
def nb_block2(X):
    xtmp1 = np.hstack(X[0])
    xtmp2 = np.hstack(X[1])
    return np.vstack((xtmp1, xtmp2))

并且两者的性能随着不同大小的数组而不同。您应该对自己的应用程序进行基准测试。

然后调用：

A = np.zeros((50,30))
B = np.ones((50,30))

X = np.block([[A, B], [-B, -A]])
Y = nb_block(((A, B), (-B, -A)))  # Note the tuple-of-tuples vs list-of-lists
np.all_close(X, Y)  # True

【讨论】：