Question

numpy或scipy中是否存在特定功能，以便轻松创建双对角或三对角矩阵？

到目前为止，我想出的最简单的解决方案是：

main_diag = [1,2,3,4,5]
off1 = [1,2,3,4]
np.diag(main_diag) + np.diag(off1, 1)

array([[1, 1, 0, 0, 0],
       [0, 2, 2, 0, 0],
       [0, 0, 3, 3, 0],
       [0, 0, 0, 4, 4],
       [0, 0, 0, 0, 5]])

是否有一些功能可以一步完成？

Answer 1

您可以通过scipy使用sparse.diags。这是OP的简单解决方案。

from scipy import sparse

diags = [range(1, 6), range(1, 5)] 
sparse.diags(diags, [0, 1]).toarray()

array([[1., 1., 0., 0., 0.],
       [0., 2., 2., 0., 0.],
       [0., 0., 3., 3., 0.],
       [0., 0., 0., 4., 4.],
       [0., 0., 0., 0., 5.]])

概括地说，

size = 5
num_diags = 2 # should be < size
diags = [range(1, size+1-i) for i in range(num_diags)] 

sparse.diags(diags, range(num_diags)).toarray()

array([[1., 1., 0., 0., 0.],
       [0., 2., 2., 0., 0.],
       [0., 0., 3., 3., 0.],
       [0., 0., 0., 4., 4.],
       [0., 0., 0., 0., 5.]])

Answer 2

这是您可以使用NumPy实现的另一种方法：

import numpy as np

def make_diags(diags):
    # Make a linear array for the whole matrix
    n = len(diags[0])
    a = np.zeros(n * n, dtype=diags[0].dtype)
    # Assign each diagonal to the right stride
    step = n + 1
    for i, diag in enumerate(diags):
        a[i:(n - i) * n:step] = diag
    # Reshape
    return a.reshape(n, n)

print(make_diags([np.arange(1, 6), np.arange(1, 5), np.arange(1, 4)]))
# [[1 1 1 0 0]
#  [0 2 2 2 0]
#  [0 0 3 3 3]
#  [0 0 0 4 4]
#  [0 0 0 0 5]]

它似乎比使用scipy.sparse.diags更快，尽管幅度不大，所以我不确定是否足以弥补额外的复杂性：

import scipy.sparse

def make_diags_sparse(diags):
    return scipy.sparse.diags(diags, range(len(diags)), dtype=diags[0].dtype).toarray()

diags = [np.arange(10000 - i) for i in range(1000)]
%timeit make_diags(diags)
# 129 ms ± 439 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit make_diags_sparse(diags)
# 298 ms ± 1.97 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

编辑：可以对该函数进行一点编辑，以添加如下所示堆叠较低对角线的功能：

import numpy as np

def make_diags(diags, lower=False):
    n = len(diags[0])
    nn = n * n
    a = np.zeros(nn, dtype=diags[0].dtype)
    step = n + 1
    start, end = (n, 1) if lower else (1, n)
    for i, diag in enumerate(diags):
        a[i * start:nn - i * end:step] = diag
    return a.reshape(n, n)

print(make_diags([np.arange(1, 6), np.arange(1, 5), np.arange(1, 4)], lower=True))
# [[1 0 0 0 0]
#  [1 2 0 0 0]
#  [1 2 3 0 0]
#  [0 2 3 4 0]
#  [0 0 3 4 5]]

编辑2：另一个进一步的编辑，以允许任意的对角线索引和形状，更类似于scipy.sparse.diags：

import numpy as np

def make_diags(diags, offsets, shape=None):
    if shape is None:
        shape = (len(diags[0]) + abs(offsets[0]),) * 2
    rows, cols = shape
    size = rows * cols
    a = np.zeros(size, dtype=diags[0].dtype)
    step = cols + 1
    d = cols - rows
    for offset, diag in zip(offsets, diags):
        start = max(offset, 0) - min(offset, 0) * cols
        offset2 = d - offset
        end = size - (max(offset2, 0) - min(offset2, 0) * cols)
        a[start:end:step] = diag
    return a.reshape(rows, cols)

print(make_diags([np.arange(1, 4), np.arange(1, 5), np.arange(1, 3)], (-1, 2, 4), (4, 6)))
# [[0 0 1 0 1 0]
#  [1 0 0 2 0 2]
#  [0 2 0 0 3 0]
#  [0 0 3 0 0 4]]

创建一个对角矩阵

2 个答案: