核心概念

TensorMesh 是一个从底层为 PyTorch 打造的有限元库。Mesh 就是一个 torch.nn.Module,弱形式只是接收基函数张量的普通 forward 方法,而每一次线性求解都是可微的算子。同一份求解笔记本 CPU 上泊松问题的代码,也能在 GPU 上运行,并将梯度反向传播到可学习参数——无需改动任何有限元逻辑。

本页给出一套心智模型:各模块如何协同工作,以及这些选择背后的设计原则。

有限元方法工作流

在 TensorMesh 中求解偏微分方程遵循一条标准流水线:

Mesh  →  Assembler  →  SparseMatrix  →  Condenser  →  Solve
  1. Mesh 将求解域离散为点和网格单元。

  2. Assembler 将弱形式(a(u, v)l(v))转化为 SparseMatrix 或载荷向量。

  3. Condenser 通过静态凝聚施加狄利克雷边界条件,得到一个仅含内部自由度的缩减系统。

  4. 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 │
└──────────┘    └────────────┘                                          └───────────┘

每个模块包含的内容:

稀疏线性代数栈(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 中的回退路径会尽力提供支持。

下一步