有限元程序加速小技巧 | Python花式索引

【问】：在组装整体刚度时，常用的方法是先索引单元节点的局部自由度号在整体刚度矩阵的自由度号，然后通过双重循环进行装配，比如这样：

for i in range(len(ElementNodeDOF)):
    for j in range(len(ElementNodeDOF)):
        KKG[ElementNodeDOF[i], ElementNodeDOF[j]] += ElementStiffnessMatrix[i, j]

这时候有没有其他方法能够避免双重循环结构呢？

【答】：有的，可以尝试 Python 的花式索引。

【木木分析】：有关花式索引的概念可以在网上稍微搜索一下，木木可以提供一个思路，这里直接以一个案例来简要说明花式索引在有限元整体刚度组装中的作用。

# -*- coding:UTF-8 -*-
import numpy as np

# 创建一个 6x6 的全局刚度矩阵，初始化为全零
KKG = np.zeros((6, 6))

# 假设这是一个元素刚度矩阵，2x2
ElementStiffnessMatrix = np.array([[10, 2],
                                   [ 2, 5]])
# 对应的自由度索引 (假设这个元素影响全局矩阵的第2和第4个自由度)
ElementNodeDOF = np.array([1, 3])

# 使用 np.ix_ 创建适合花式索引的网格
KKG[np.ix_(ElementNodeDOF, ElementNodeDOF)] += ElementStiffnessMatrix

print("组装后的全局刚度矩阵 KKG:")
print(KKG)

运行结果：

组装后的全局刚度矩阵 KKG:
[[ 0.  0.  0.  0.  0.  0.]
 [ 0. 10.  0.  2.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.]
 [ 0.  2.  0.  5.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.]]

原来的双重循环结构只需要一条语句：

KKG[np.ix_(ElementNodeDOF, ElementNodeDOF)] += ElementStiffnessMatrix

即可完成，避免了双重循环的结构，直接通过索引切片的方式，大大提升了装配效率。

注：在编程的过程中，并不是循环越多、条件结构越复杂，程序就越优！