如何批量将NumPy下三角元素数组转换为对应全矩阵数组

你的代码存在三处索引逻辑问题，导致形状不匹配、运算结果不符合预期：

• 批量填充下三角时，x[:, np.tril_indices(6)]的写法会触发NumPy高级索引的广播规则，最终索引到的形状为(N,N,21,6)，和ld的(N,21)形状完全不匹配，是报错的直接原因。
• 转置操作写x[:].T会把批次维度也参与转置，无法实现“每个子矩阵单独转置补全上三角”的效果。
• 直接调用np.fill_diagonal(x[:], ...)只会对整个三维数组的主对角线赋值，不会给每个批次的6x6子矩阵单独修正对角线。

全程使用NumPy广播和高级索引，不需要Python循环即可完成批量转换，代码如下：

import numpy as np

# 输入N组6阶矩阵下三角元素，形状为(N, 21)
ld = np.arange(2*21).reshape(2, 21)
N = ld.shape[0]
mat_len = 6

# 初始化批量矩阵数组
x = np.zeros((N, mat_len, mat_len))

# 批量填充所有矩阵的下三角部分
rows, cols = np.tril_indices(mat_len)
# 构造带批次维度的索引，广播后形状与ld完全匹配
x[np.arange(N)[:, None], rows, cols] = ld

# 补全每个矩阵的上三角部分：仅转置每个子矩阵的最后两个维度，保留批次维度顺序
x = x + x.transpose(0, 2, 1)

# 批量修正对角线值（转置后对角线元素被重复相加，还原为原始下三角中的对角线值）
diag_ld_pos = [0, 2, 5, 9, 14, 20]
x[np.arange(N)[:, None], np.arange(mat_len), np.arange(mat_len)] = ld[:, diag_ld_pos]

运行上述代码得到的输出和预期结果完全一致。

关键逻辑说明 #

• 批量索引赋值时，通过np.arange(N)[:, None]把批次索引从形状(N,)转为(N,1)，和长度为21的下三角行列索引广播后，得到形状为(N,21)的索引位置，正好匹配ld的形状，不会出现广播报错。
• 多维度转置使用transpose(0,2,1)，表示保持第0维（批次维度）顺序不变，仅交换每个子矩阵的行、列维度，不会混淆不同批次的矩阵数据。
• 批量修正对角线时，同样构造三维索引，直接把每组下三角数据中对应对角线位置的原始值，写入对应子矩阵的对角线位置，实现批量赋值。

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