如何沿方阵第k条对角线填充子矩阵？类似scipy.linalg.block_diag功能

要实现类似scipy.linalg.block_diag()但针对非主对角线的块矩阵填充，最直接的方式是利用numpy数组切片定位块位置，或者封装自定义函数批量处理。以下是针对你需求的具体实现：

核心思路 #

你的目标矩阵是16n×16n的密集数组，由2n个8×8的子矩阵构成。主对角线已通过block_diag填充完成，接下来只需定位目标对角线对应的子矩阵切片，直接赋值即可。

方法一：直接切片赋值（快速实现） #

假设你要填充主对角线上方第一条对角线的子矩阵为h_coupling（8×8矩阵），下方第一条对角线为其转置（或自定义矩阵），代码如下：

import numpy as np
import scipy.linalg as sc

# 示例参数（替换为你的实际矩阵）
n = 3
block_size = 8
h_bulk_plane = np.random.rand(block_size, block_size)
hS_upper = np.random.rand(block_size, block_size)
h_coupling = np.random.rand(block_size, block_size)

# 1. 构造主对角线矩阵
diag_blocks = [h_bulk_plane + hS_upper] + [h_bulk_plane]*(2*(n-1)) + [h_bulk_plane - hS_upper]
H = sc.block_diag(*diag_blocks)

# 2. 填充主对角线上方第一条对角线（块位置(i, i+1)）
num_blocks = len(diag_blocks)
for i in range(num_blocks - 1):
    # 定位当前块对应的矩阵切片
    row_slice = slice(i*block_size, (i+1)*block_size)
    col_slice = slice((i+1)*block_size, (i+2)*block_size)
    H[row_slice, col_slice] = h_coupling

# 3. 填充主对角线下方第一条对角线（块位置(i+1, i)）
for i in range(num_blocks - 1):
    row_slice = slice((i+1)*block_size, (i+2)*block_size)
    col_slice = slice(i*block_size, (i+1)*block_size)
    H[row_slice, col_slice] = h_coupling.T  # 替换为你的下对角线矩阵

方法二：自定义通用函数（支持任意k条对角线） #

如果需要频繁处理不同对角线的块填充，可以封装一个类似block_diag的通用函数：

def block_offdiag(blocks, k=1):
    """
    构造块矩阵，在指定的第k条对角线填充子矩阵
    - k=1: 主对角线上方第一条
    - k=-1: 主对角线下方第一条
    - blocks: 待填充的子矩阵列表，长度需满足 len(blocks) = 总块数 - abs(k)
    """
    block_size = blocks[0].shape[0]
    dtype = blocks[0].dtype
    # 计算总块数
    if k > 0:
        total_blocks = len(blocks) + k
    else:
        total_blocks = len(blocks) - k
    
    H = np.zeros((total_blocks*block_size, total_blocks*block_size), dtype=dtype)
    for idx, block in enumerate(blocks):
        # 计算当前块在矩阵中的行、列起始索引
        if k >= 0:
            row_start = idx * block_size
            col_start = (idx + k) * block_size
        else:
            row_start = (idx - k) * block_size
            col_start = idx * block_size
        # 赋值子矩阵
        H[row_start:row_start+block_size, col_start:col_start+block_size] = block
    return H

使用时只需和主对角线矩阵相加：

# 主对角线矩阵
H_diag = sc.block_diag(*diag_blocks)

# 上方第一条对角线矩阵
offdiag_upper = block_offdiag([h_coupling]*(num_blocks-1), k=1)

# 下方第一条对角线矩阵
offdiag_lower = block_offdiag([h_coupling.T]*(num_blocks-1), k=-1)

# 合并得到最终矩阵
H = H_diag + offdiag_upper + offdiag_lower

注意事项 #

• 确保所有子矩阵的尺寸一致（此处均为8×8），否则会出现形状不匹配的错误。
• 如果处理大规模矩阵，优先考虑稀疏矩阵（如scipy.sparse.block_diag结合sparse.diags），可节省内存并提升效率。

内容的提问来源于stack exchange，提问作者user27047503