TensorMesh 文档¶

████████╗███████╗███╗   ██╗███████╗ ██████╗ ██████╗ ███╗   ███╗███████╗███████╗██╗  ██╗
╚══██╔══╝██╔════╝████╗  ██║██╔════╝██╔═══██╗██╔══██╗████╗ ████║██╔════╝██╔════╝██║  ██║
   ██║   █████╗  ██╔██╗ ██║███████╗██║   ██║██████╔╝██╔████╔██║█████╗  ███████╗███████║
   ██║   ██╔══╝  ██║╚██╗██║╚════██║██║   ██║██╔══██╗██║╚██╔╝██║██╔══╝  ╚════██║██╔══██║
   ██║   ███████╗██║ ╚████║███████║╚██████╔╝██║  ██║██║ ╚═╝ ██║███████╗███████║██║  ██║
   ╚═╝   ╚══════╝╚═╝  ╚═══╝╚══════╝ ╚═════╝ ╚═╝  ╚═╝╚═╝     ╚═╝╚══════╝╚══════╝╚═╝  ╚═╝

一个快速、可微、无 JIT 且易于调试的有限元库——将现代 PyTorch 的使用体验带入科学计算。

TensorMesh 是一个基于 PyTorch 构建的有限元方法（FEM）库，旨在优雅而高效地求解偏微分方程（PDE）。通过与 PyTorch 生态原生集成，它在不牺牲经典有限元严谨性的前提下，提供自动微分、GPU 加速以及直观的 Pythonic 接口。

核心优势¶

GPU 原生与可微：完全基于 PyTorch 从底层构建。只需一行代码即可将整个有限元工作流迁移到 GPU——其后续的每一次装配、求解和梯度计算都会自动继承设备，无需单独的后端或数据搬运步骤。原生 autograd 可无缝贯穿装配与求解，从而实现端到端可微的 PDE 流水线。
高性能张量化装配：由 TensorGalerkin 驱动的全张量化 Map-Reduce 算法，将逐单元运算融合为单体的 GPU 内核，消除了 Python 层级的循环，相较于基于 CPU 的有限元技术栈带来数量级的加速。
无 JIT 且易于调试：即时执行（eager execution），无编译开销。动态网格、自适应加密以及交互式工作流均可开箱即用——没有重新编译的延迟，也没有难以理解的追踪信息。
全面的单元与网格支持：支持三角形、四面体、四棱锥和三棱柱单元，可为常见几何自动生成网格，并与 Gmsh / VTK-HDF5 实现无缝 I/O。
灵活的求解器：由 torch-sla 驱动，这是我们配套的可微稀疏线性代数库。线性、非线性和特征值求解器可在 CPU 和 GPU 上跨多种后端运行，并完整支持 autograd、批量求解以及分布式多 GPU 扩展。
Pythonic API：用纯 Python 编写自定义弱形式——无需独立的 DSL，也无需 form 编译器。只要会写 PyTorch，就能写有限元。

功能对比¶

功能	FEniCS	scikit-fem	JAX-FEM	torch-fem	TensorMesh
自定义弱形式（Pythonic）	⚠️	✅	❌	❌	✅
易于安装	❌	✅	⚠️	✅	✅
易于调试	❌	✅	❌	✅	✅
易于输入输出	❌	❌	❌	❌	✅
大规模网格	✅	✅	❌	❌	✅
GPU 支持	✅	❌	✅	✅	✅
效率	✅	❌	✅	⚠️	✅
端到端 autograd	⚠️	❌	✅	✅	✅
深度学习集成（PyTorch）	❌	❌	⚠️	✅	✅
成熟度	✅	✅	⚠️	⚠️	⚠️

备注

自定义弱形式（Pythonic）——支持直接在 Python 中定义双线性／线性形式，无需 UFL 等独立的 DSL。端到端 autograd——梯度可原生贯穿整个流水线；FEniCS 通过外部的 dolfin-adjoint 包来支持这一点。成熟度——反映项目的存续时间、生态规模以及生产环境部署情况。

引用¶

TensorMesh 是 TensorGalerkin 框架中的有限元求解器组件。如果您在研究中使用了 TensorMesh，请引用 TensorGalerkin 论文：

@article{wen2026tensorgalerkin,
  title   = {Learning, Solving and Optimizing PDEs with {TensorGalerkin}:
             an Efficient High-Performance Galerkin Assembly Algorithm},
  author  = {Wen, Shizheng and Chi, Mingyuan and Yu, Tianwei and
             Moseley, Ben and Michelis, Mike Yan and Ren, Pu and
             Sun, Hao and Mishra, Siddhartha},
  journal = {arXiv preprint arXiv:2602.05052},
  year    = {2026}
}

如果您的工作还依赖于 torch-sla（TensorMesh 的求解器后端），请额外引用：

@article{chi2026torchsla,
  title   = {torch-sla: Differentiable Sparse Linear Algebra with Adjoint
             Solvers and Sparse Tensor Parallelism for PyTorch},
  author  = {Chi, Mingyuan and Wen, Shizheng},
  journal = {arXiv preprint arXiv:2601.13994},
  year    = {2026}
}

联系我们¶

TensorMesh 基于 Apache License 2.0 发布。如需合作与洽谈，请通过 shizheng.wen@sam.math.ethz.ch 联系 Shizheng Wen。