将 numpy.bmat 与 numba 一起使用答案

【问题标题】：Using numpy.bmat with numba将 numpy.bmat 与 numba 一起使用
【发布时间】：2020-06-04 23:41:40
【问题描述】：

我正在尝试在我的 numba 优化 python 程序中使用 np.bmat。为此，我必须手动定义一个 jitted 函数 bmat，因为不支持来自 numpy 的本机函数：

@njit
def _bmat_2d(matrices):
    arr_rows = []
    for row in matrices:
        arr_rows.append(np.concatenate(row, axis=-1))
    return np.array(np.concatenate(arr_rows, axis=0))

（此代码或多或少是 numpy 代码的简化副本）

但是：

numba 只接受 np.concatenate [1] 输入中的元组
numba 非常不擅长将任意列表转换为元组 [2]

你对此有什么想法吗？

参考：

【问题讨论】：

标签： python numpy numba

【解决方案1】：

以下内容是否适合您的目的？

import numpy as np
import numba as nb

@nb.njit
def _bmat_2d(m):
    out = np.hstack(m[0])
    for row in m[1:]:
        x = np.hstack(row)
        out = np.vstack((out, x))

    return out

A = np.random.randint(10, size=(3,2))
B = np.random.randint(10, size=(3,1))
C = np.random.randint(10, size=(3,3))
D = np.random.randint(10, size=(4,6))

a = np.bmat(((A, B, C), (D,)))
b = _bmat_2d(((A, B, C), (D,)))

print(np.allclose((a, b))  # True

请注意，您必须传入元组而不是列表列表，否则您将收到“反射列表”错误，因为当前版本中的 Numba 无法处理列表列表。

【讨论】：

不错！并且性能明显优于 numpy（例如，您的示例中的 timeit 为 numpy 提供 10.4 µs ± 68.6 ns 而对于 jitted 函数为 1.69 µs ± 3.66 ns）
请注意_bmat_2d(((A, B, C),)) 会失败，因为它假定不止一行。我不知道如何满足这种情况下的 numba 编译器。

【解决方案2】：

由于np.hstack 对我来说与numba 不兼容，我不得不编写自己的解决方案。也许你们中的一些人觉得这很有用。它不漂亮，但它可以完成工作。

这基本上与J = np.bmat([[J_1, J_2], [J_3, J_4]]) 做同样的事情。

请务必更改 J = np.zeros((8, len(J_1[0])*2)) 以适合您想要的输出数组：

import numpy as np
import numba

@numba.njit
def main():
    J_1 = np.array([[-64., 25.6, 25.6, 12.8], [25.6, -25.6, 0., 0.], [25.6, 0., -25.6, 0.], [12.8, 0., 0., -652.8]])
    J_2 = np.array([[-85.33333333, 34.13333333, 34.13333333, 17.06666667], [34.13333333, -34.13333333, 0., 0.], [34.13333333, 0., -34.13333333, 0.], [17.06666667, 0., 0., -870.4]])
    J_3 = np.array([[85.33333333, -34.13333333, -34.13333333, -17.06666667], [-34.13333333, 34.13333333, -0., -0.], [-34.13333333, -0., 34.13333333, -0.], [-17.06666667, -0., -0., 870.4]])
    J_4 = np.array([[-64., 25.6, 25.6, 12.8], [25.6, -25.6, 0., 0.], [25.6, 0., -25.6, 0.], [12.8, 0., 0., -652.8]])

    J = np.zeros((8, len(J_1[0])*2))
    for idx, _ in enumerate(J_1[0]):
        J[0][idx], J[1][idx], J[2][idx], J[3][idx], J[4][idx], J[5][idx], J[6][idx], J[7][idx] = J_1[0][idx], J_1[1][idx], J_1[2][idx], J_1[3][idx], J_3[0][idx], J_3[1][idx], J_3[2][idx], J_3[3][idx]
        J[0][idx+len(J_1[0])], J[1][idx+len(J_1[0])], J[2][idx+len(J_1[0])], J[3][idx+len(J_1[0])], J[4][idx+len(J_1[0])], J[5][idx+len(J_1[0])], J[6][idx+len(J_1[0])], J[7][idx+len(J_1[0])] = J_2[0][idx], J_2[1][idx], J_2[2][idx], J_2[3][idx], J_4[0][idx], J_4[1][idx], J_4[2][idx], J_4[3][idx]

    print(J)

if __name__ == '__main__':
    main()

编辑：

一个人在另一个线程上帮助我用这个简单的替换 np.bmat，它在 numba.njit 内部工作：

J = np.vstack((np.hstack((J_1, J_2)), np.hstack((J_3, J_4))))

【讨论】：