scipy 创建带状矩阵的有效方法

jv2fixgn 于 2022-11-10 发布在其他

关注(0)|答案(3)|浏览(150)

我正在尝试创建一个特定样式的艾德矩阵（参见Wikipedia）。下面的代码可以工作，但是对于大的M（~300左右），由于for循环的原因，它变得相当慢。有没有办法将它矢量化/更好地使用NumPy和/或SciPy？我在计算它对应的数学运算时遇到了麻烦，因此到目前为止我还没有成功。
我的代码如下

def banded_matrix(M):
    phis = np.linspace(0, 2*np.pi, M)
    i = 0
    ham = np.zeros((int(2*M), int(2*M)))
    for phi in phis:
        ham_phi = np.array([[1, 1],
                            [1, -1]])*(1+np.cos(phi))
        array_phi = np.zeros(M)
        array_phi[i] = 1
        mat_phi = np.diag(array_phi)
        ham += np.kron(mat_phi, ham_phi)
        i += 1
    return ham

有了%timeit banded_matrix(M=300)，在我的电脑上大约需要4秒钟。
由于代码有点不透明，我想做的是构造一个2 M × 2 M的大矩阵。在某种意义上，它在“宽度2”的对角线上有M个元素，其中的元素是依赖于phi的2x2矩阵ham_phi。矩阵随后将被对角化，所以也许人们甚至可以利用它的结构/它相当稀疏的事实来加快速度，但我不确定。
如果有人有一个想法去与此，我很乐意跟进！

scipy

来源：https://stackoverflow.com/questions/72706058/efficient-way-of-creating-a-band-or-banded-matrix

3条答案

按热度按时间

kqqjbcuj1#

您的矩阵是按块对角排列的，因此可以使用scipy.linalg.block_diag：

import numpy as np
from scipy.linalg import block_diag

def banded_matrix_scipy(M):
    ham = np.array([[1, 1], [1, -1]])
    phis = np.linspace(0, 2 * np.pi, M)
    ham_phis = ham * (1 + np.cos(phis))[:, None, None]
    return block_diag(*ham_phis)

让我们检查一下它是否工作正常，速度是否更快：
第一个

赞(0）回复(0）举报 2022-11-10

8fsztsew2#

强制性np.einsum基准

def banded_matrix_einsum(M):
    return np.einsum('ij, kl-> ikjl',
              np.eye(M)*(1 + np.cos(np.linspace(0, 2 * np.pi, M))),
              np.array([[1, 1], [1, -1]])).reshape(2*M, 2*M)

banded_matrix_einsum(4)

输出量

array([[ 2. ,  2. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],
       [ 2. , -2. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],
       [ 0. ,  0. ,  0.5,  0.5,  0. ,  0. ,  0. ,  0. ],
       [ 0. ,  0. ,  0.5, -0.5,  0. ,  0. ,  0. ,  0. ],
       [ 0. ,  0. ,  0. ,  0. ,  0.5,  0.5,  0. ,  0. ],
       [ 0. ,  0. ,  0. ,  0. ,  0.5, -0.5,  0. ,  0. ],
       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  2. ,  2. ],
       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  2. , -2. ]])

性能指标评测结果

import perfplot

perfplot.show(
    setup = lambda M: M,
    kernels = [banded_matrix_einsum, banded_matrix_scipy, banded_matrix],
    n_range = [50, 100, 150, 200, 250, 300],
    logx = False
)

scipy.linalg.block_diag与np.einsum的详细信息

perfplot.show(
    setup = lambda M: M,
    kernels = [banded_matrix_einsum, banded_matrix_scipy],
    n_range = [50, 100, 150, 200, 250, 300, 350, 400],
    logx = False
)

赞(0）回复(0）举报 2022-11-10

hsgswve43#

作为另一种方法，您可以使用numba加速器通过jitting来加速它。我提出了一个基于paime answer的等价的scipy.linalg.block_diag numba方法：

import numba as nb

@nb.njit
def block_diag_numba(result, ham_phis):
    for i in range(ham_phis.shape[0]):
        for j in range(ham_phis.shape[1]):
            result[i * 2,     i * 2:i * 2 + 2] = ham_phis[i, 0]
            result[i * 2 + 1, i * 2:i * 2 + 2] = ham_phis[i, 1]
    return result

def numba_(M):
    ham = np.array([[1, 1], [1, -1]])
    phis = np.linspace(0, 2 * np.pi, M)
    ham_phis = ham * (1 + np.cos(phis))[:, None, None]
    return block_diag_numba(np.zeros((M * ham.shape[1], M * ham.shape[1])), ham_phis)

这种方法比以前的方法快至少4-5倍，m=400（us规模）。这种方法可以调整为其他数组形状，并通过进一步优化代码（不使用 paime 答案）和将所有代码行转换为numba函数或并行化来改进。我没有进一步研究，因为 paime 答案的性能似乎可以满足OP的接受;为了说明我们可以使用numba编写更快的scipy.linalg.block_diag等效代码：

赞(0）回复(0）举报 2022-11-10

我来回答

scipy 创建带状矩阵的有效方法

3条答案

相关问题

热门标签

最新问答