Python构建含2X2子矩阵的稀疏矩阵报错，求解决方案

问题背景 #

要构建一个2J×2J的分块矩阵C，其中每个子块都是2×2矩阵：

[ 0 A 0 0 0 0 ...       0 0 
  B 0 A 0 0 0 ...       0 0
  0 B 0 A 0 0 ...       0 0
  0 0 B 0 A 0 ....      0 0
            ..........
  0 0 0 0 0   ....      B 0]

这里的0是全零2×2矩阵，A、B是给定的2×2矩阵，分别位于分块的上对角线和下对角线。

尝试用scipy稀疏矩阵实现时，写出了如下代码（J=5）：

def TheMatrix(J,alpha):
    A = [[0.5, alpha],[ alpha,0.5]]
    B =[[0.5,-alpha],[-alpha,0.5]]

    upperDiag = A*np.ones(J-1)
    lowerDiag = B*np.ones(J-1)

    diagonals = [upperDiag, lowerDiag]  # non-zero diagonals
    ASparse = sps.diags(diagonals,offsets=[1, -1])
    ASparse = sps.csc_matrix(ASparse)

运行后触发报错：

File "myprg.py", line 16, in TheMatrix(J,alpha)
upperDiag = A*np.ones(J-1)
            ~^~~~~~~~~~~~~
ValueError: operands could not be broadcast together with shapes (2,2) (5,)

问题原因 #

直接用列表形式的A、B和一维numpy数组相乘，无法实现“重复J-1次分块”的效果——numpy广播规则不支持(2,2)和(J-1,)形状的数组相乘。

解决方案 #

以下两种方法可以解决这个问题，分别适合不同场景：

方法1：用克罗内克积快速生成分块结构 #

利用numpy.kron（克罗内克积）可以快速生成重复的分块矩阵，再转为稀疏矩阵：

import numpy as np
import scipy.sparse as sps

def TheMatrix(J, alpha):
    # 将A、B转为numpy数组（必须，列表无法参与克罗内克积计算）
    A = np.array([[0.5, alpha], [alpha, 0.5]])
    B = np.array([[0.5, -alpha], [-alpha, 0.5]])
    
    # 生成分块上对角线：J-1个A沿上对角线排列
    upper_blocks = np.kron(np.eye(J, k=1), A)
    # 生成分块下对角线：J-1个B沿下对角线排列
    lower_blocks = np.kron(np.eye(J, k=-1), B)
    
    # 合并后转为稀疏矩阵
    return sps.csc_matrix(upper_blocks + lower_blocks)

说明：np.eye(J, k=1)是J×J的单位矩阵，上对角线为1，和A做克罗内克积后，就会把每个1的位置替换成A矩阵，正好对应分块上对角线的结构。

方法2：直接构造稀疏矩阵的索引与数据（内存高效） #

如果J很大，先生成全矩阵会占用过多内存，可以直接构造稀疏矩阵的data、row、col数组：

import numpy as np
import scipy.sparse as sps

def TheMatrix(J, alpha):
    A = np.array([[0.5, alpha], [alpha, 0.5]])
    B = np.array([[0.5, -alpha], [-alpha, 0.5]])
    
    # 分块对角线的偏移量是±2（每个子块是2×2，所以元素级偏移为2）
    upper_offset = 2
    lower_offset = -2
    
    # 生成上对角线的所有元素：把A展平后重复J-1次
    upper_data = np.tile(A.flatten(), J-1)
    # 生成下对角线的所有元素：把B展平后重复J-1次
    lower_data = np.tile(B.flatten(), J-1)
    
    # 计算上对角线元素的行、列索引
    upper_rows = []
    upper_cols = []
    for i in range(J-1):
        # 第i个A块的行起始位置是i*2，列起始位置是i*2 + upper_offset
        block_rows = np.arange(i*2, i*2 + 2)
        block_cols = block_rows + upper_offset
        upper_rows.extend(block_rows)
        upper_cols.extend(block_cols)
    
    # 计算下对角线元素的行、列索引
    lower_rows = []
    lower_cols = []
    for i in range(J-1):
        # 第i个B块的行起始位置是i*2 + 2，列起始位置是i*2
        block_rows = np.arange(i*2 + 2, i*2 + 4)
        block_cols = block_rows + lower_offset
        lower_rows.extend(block_rows)
        lower_cols.extend(block_cols)
    
    # 合并所有数据和索引
    all_data = np.concatenate([upper_data, lower_data])
    all_rows = np.concatenate([upper_rows, lower_rows])
    all_cols = np.concatenate([upper_cols, lower_cols])
    
    # 构造稀疏矩阵
    return sps.csc_matrix((all_data, (all_rows, all_cols)), shape=(2*J, 2*J))

说明：这种方法只存储非零元素的位置和值，内存占用远低于全矩阵，适合J较大的场景。

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