核心概念¶
TensorMesh 是一个从底层为 PyTorch 打造的有限元库。Mesh 就是一个 torch.nn.Module,弱形式只是接收基函数张量的普通 forward 方法,而每一次线性求解都是可微的算子。同一份求解笔记本 CPU 上泊松问题的代码,也能在 GPU 上运行,并将梯度反向传播到可学习参数——无需改动任何有限元逻辑。
本页给出一套心智模型:各模块如何协同工作,以及这些选择背后的设计原则。
有限元方法工作流¶
在 TensorMesh 中求解偏微分方程遵循一条标准流水线:
Mesh → Assembler → SparseMatrix → Condenser → Solve
Mesh 将求解域离散为点和网格单元。
Assembler 将弱形式(
a(u, v)或l(v))转化为SparseMatrix或载荷向量。Condenser 通过静态凝聚施加狄利克雷边界条件,得到一个仅含内部自由度的缩减系统。
Solve 将缩减系统分派给某个稀疏线性代数后端(通过 torch-sla 包)。
快速开始 用大约 30 行 Python 端到端地走完这条流水线。本指南后续的每一章都会聚焦于其中一个阶段。
模块全景¶
整个库沿流水线清晰划分。箭头表示数据流向,而非导入方向:
┌──────────┐ ┌────────────┐ ┌──────────────┐ ┌───────────┐ ┌─────────┐
│ Mesh │ → │ Assembler │ → │ SparseMatrix │ → │ Condenser │ → │ Solve │
│ (nn. │ │ (Element / │ │ (torch_sla. │ │ (Dirichlet│ │ (torch- │
│ Module) │ │ Node / │ │ SparseTensor│ │ static │ │ sla │
│ │ │ Facet) │ │ + spmm / @) │ │ cond.) │ │ .solve) │
└──────────┘ └────────────┘ └──────────────┘ └───────────┘ └─────────┘
↑ ↑ │
│ │ ▼
┌──────────┐ ┌────────────┐ ┌───────────┐
│ element │ │ functional │ │ Postproc │
│ (Triangle│ │ (voigt, │ │ visualize │
│ Hex,…) │ │ strain,…) │ │ ode step │
└──────────┘ └────────────┘ └───────────┘
每个模块包含的内容:
tensormesh.mesh——Mesh及其内置生成器(gen_rectangle、gen_circle、gen_cube等);meshio 输入/输出;邻接、分区与图着色辅助工具。tensormesh.element—— 参考单元形状(Triangle、Quadrilateral、Tetrahedron、Hexahedron、Prism、Pyramid、Line)、基函数求值、求积规则,以及 Gmsh/VTK ↔ TensorMesh 排序约定。tensormesh.assemble—— 三个弱形式基类ElementAssembler、NodeAssembler、FacetAssembler,以及针对最常见弱形式(Laplace、质量、线弹性、新胡克等)的内置实现。tensormesh.sparse——SparseMatrix(torch_sla.SparseTensor的子类),因此线性系统通过K.solve(b)求解、非线性系统通过K.nonlinear_solve(residual, u0, *params)求解,二者均由torch-sla分派。库内自带的spsolve/nonlinear_solve()自由函数是遗留入口点,已计划移除。tensormesh.operator——Condenser,用于通过静态凝聚施加狄利克雷边界条件。tensormesh.ode—— 用于瞬态问题的显式与隐式线性时间积分器(Euler、中点法、龙格-库塔)。tensormesh.functional—— Voigt 弹性辅助函数、应变 / 应力,以及其他在forward方法内部使用的张量工具。tensormesh.dataset——MeshGen以及用于生成训练数据集的预置多频方程类。tensormesh.material——IsotropicMaterial以及材料库预设(Steel、Aluminum、Rubber、Glass)。tensormesh.optimizer——OCOptimizer(最优性准则法),用于基于柔度的拓扑优化。tensormesh.visualization—— matplotlib(2D)与 PyVista(3D)后端;由plot()延迟导入。tensormesh.distributed—— 跨多个进程(rank)的图分区分布式装配(高级特性;参见示例库)。
稀疏线性代数栈(SparseMatrix、.solve / .nonlinear_solve、感知梯度的伴随反向传播)被委托给一个独立的包 torch-sla,并与同一生态系统中的其他项目共享。
设计原则¶
PyTorch 原生。 Mesh 继承自 torch.nn.Module。它的 points 与逐单元连接性都是 buffer;逐节点字段同样是 buffer。装配器本身也是 nn.Module。这里没有单独的“有限元内核”抽象层——一切都是张量,流经我们熟悉的 PyTorch 机制(.to(device)、.double()、state_dict、autograd、JIT 追踪)。
纯 Python 弱形式。 用户唯一需要编写的、与具体偏微分方程相关的代码,就是一个返回求积点处被积函数的 forward 方法:
class LaplaceAssembler(ElementAssembler):
def forward(self, gradu, gradv):
return gradu @ gradv
库会负责参考单元求值、几何、求积权重,以及全局装配到稀疏矩阵的步骤。同样的模式也适用于载荷向量(NodeAssembler)和边界积分(FacetAssembler)。
张量化装配。 不存在 Python 层面对单元的循环。在 __call__ 内部,基函数和求积点针对整个网格只求值一次;用户的 forward 在一个已经具备单元与求积维度、可直接广播的张量上运行;全局装配则是一次稀疏散射。其结果是:装配就是单个 GPU 内核。
生而可微。 SparseMatrix.solve 是一个带有自定义反向传播(伴随稀疏求解)的 torch.autograd.Function。因此梯度可以端到端地从损失出发,反向穿过线性求解、装配,以及任何参与其中的参数——无论是材料系数、狄利克雷取值,还是神经网络的预测。参见 可微性。
模块化线性代数。 求解器层是一个独立的包 torch-sla。同一份有限元代码只需改动一个关键字参数,即可在 SciPy(CPU)、Eigen(CPU)、原生 PyTorch(CPU/GPU)、cuDSS(GPU)和 CuPy(GPU)之间切换。PETSc 与 Hypre 已列入 torch-sla 的路线图;在它们正式发布之前,若本地已安装这些库,tensormesh.sparse 中的回退路径会尽力提供支持。
下一步¶
网格 —— 构建、检查与加载网格;逐节点与逐单元数据;meshio 往返读写。
单元与求积 —— 单元家族以及基函数 / 求积接口。
形式 —— 依照
ElementAssembler/NodeAssembler/FacetAssembler的约定编写你自己的弱形式。快速开始 —— 同一条流水线,作为完整的实战示例。