在 Numba 支持下组装块矩阵

Question

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

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

接下来是各种计算，涉及大量的 numpy 操作和 for 循环。因此，我想使用 Numba 来优化代码。但是，Numba 似乎不支持 numpy 块功能。矩阵A和B不是很大，所以我可以使用一个可能不如np.block优化的函数，但我仍然想以块矩阵的方式组装H为了可读性。是否有任何功能可以实现此目的？

Answer 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