如何用Numpy高效实现对称二阶张量与全张量的互转？

问题背景 #

处理应力相关的对称二阶张量时，需要将n×6的对称张量数组转换为n×3×3的全量张量进行变换，之后再还原为n×6的对称张量。现有代码可运行，但针对数百万级别的数据量，需要更高效简洁的数组操作方案。

原始代码问题分析 #

原始代码通过手动构造全张量、多次转置和reshape实现转换，存在冗余操作，且手动索引容易出错，效率较低。

优化方案 #

1. 从n×6对称张量转n×3×3全量张量 #

直接利用Numpy的切片赋值，一步完成对称张量的全量展开，避免不必要的转置和reshape操作：

import numpy as np

# 修正输入格式为n×6（示例n=2）
sym_tens_in = np.array([[0,1,2,3,4,5], [9,10,11,12,13,14]])  # 对应(S11,S22,S33,S12,S23,S13)
n = sym_tens_in.shape[0]

# 初始化全量张量
tens_full = np.zeros((n, 3, 3), dtype=sym_tens_in.dtype)

# 赋值对角线元素（S11, S22, S33）
tens_full[:, 0, 0] = sym_tens_in[:, 0]
tens_full[:, 1, 1] = sym_tens_in[:, 1]
tens_full[:, 2, 2] = sym_tens_in[:, 2]

# 赋值对称非对角线元素（利用对称特性一次赋值两个位置）
tens_full[:, 0, 1] = tens_full[:, 1, 0] = sym_tens_in[:, 3]  # S12 = S21
tens_full[:, 1, 2] = tens_full[:, 2, 1] = sym_tens_in[:, 4]  # S23 = S32
tens_full[:, 0, 2] = tens_full[:, 2, 0] = sym_tens_in[:, 5]  # S13 = S31

# 这里执行张量变换操作...

2. 从n×3×3全量张量转回n×6对称张量 #

直接提取所需的对称分量，无需reshape为一维数组，操作更直接高效：

# 初始化输出对称张量
sym_tens_out = np.empty((n, 6), dtype=tens_full.dtype)

# 提取对角线分量
sym_tens_out[:, 0] = tens_full[:, 0, 0]
sym_tens_out[:, 1] = tens_full[:, 1, 1]
sym_tens_out[:, 2] = tens_full[:, 2, 2]

# 提取非对角线对称分量（任意取对称位置即可）
sym_tens_out[:, 3] = tens_full[:, 0, 1]
sym_tens_out[:, 4] = tens_full[:, 1, 2]
sym_tens_out[:, 5] = tens_full[:, 0, 2]

效率优势 #

• 避免了多次转置、reshape操作，减少内存拷贝和中间数组的创建，降低内存开销。
• 利用Numpy的向量化切片赋值/提取操作，比手动构造数组的循环或索引方式更快，尤其在处理数百万个张量时，性能提升明显。
• 代码逻辑更清晰，减少手动索引出错的概率。

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