TensorMesh: Differentiable Finite Element Method Library for PyTorch

TensorMesh is organized into specialized modules for different stages of the FEM pipeline:

• tensormesh.mesh: Mesh data structures, built-in generators (gen_rectangle, gen_circle, gen_cube, gen_L, etc.), and I/O (Gmsh, VTK-HDF5).
• tensormesh.element: Shape functions, quadrature rules, and element transformations (geometric orders 1–4).
• tensormesh.assemble: Assemblers for matrices and vectors (ElementAssembler, NodeAssembler, FacetAssembler, FacetBilinearAssembler) and multi-field systems (MixedElementAssembler).
• tensormesh.sparse: SparseMatrix (a torch_sla.SparseTensor subclass) supporting various backends (SciPy, PyTorch Krylov, cuDSS, STRUMPACK, PyAMG, AmgX).
• tensormesh.operator: Boundary condition handling via Condenser (Dirichlet), BlochReducer (periodic), or wave operators (robin_operator, port_source).
• tensormesh.ode: Time integrators (Euler, midpoint, Runge–Kutta).
• tensormesh.dataset: Generation of parametric PDE datasets (Poisson, Heat, Wave, linear elasticity).
• tensormesh.visualization: Plotting via Matplotlib and PyVista.
• tensormesh.functional: FEM-specific tensor utilities (elasticity, Voigt notation).
• tensormesh.material: Material property definitions.
• tensormesh.optimizer: Optimization algorithms like OC for topology optimization.

TensorMesh is a fast, differentiable, and JIT-free finite element method (FEM) library built on PyTorch. It is designed to solve partial differential equations (PDEs) using modern PyTorch ergonomics, providing native automatic differentiation (autograd), GPU acceleration, and a Pythonic interface.

Key features include:

• GPU-Native: Entire workflows can be moved to the GPU with a single line of code, with all assembly and solves inheriting the device automatically.
• Tensorized Assembly: Uses a Map-Reduce algorithm (TensorGalerkin) to fuse element-wise operations into monolithic GPU kernels, avoiding Python loops.
• JIT-Free: Uses eager execution, making it suitable for dynamic meshes and interactive workflows without compilation overhead.
• Flexible Elements: Supports triangular, tetrahedral, pyramid, and prismatic elements.
• Differentiable Solvers: Powered by torch-sla, supporting linear, nonlinear, and eigenvalue solvers with full autograd and multi-GPU scaling.

The tensormesh.operator module provides the following first-class boundary condition implementations:

• Dirichlet Conditions: Implemented via tensormesh.Condenser using static condensation.
• Bloch-Floquet Periodic Conditions: Implemented via tensormesh.BlochReducer. This is used for tying opposite faces with a complex phase factor (common in band-structure computations).
• Wave Boundary Operators: Includes tensormesh.robin_operator and tensormesh.port_source for impedance and plane-wave-port boundary conditions, built on tensormesh.FacetBilinearAssembler.
• Energy-based Terms: Penalty contact, surface tension, and pressure loads are supported via tensormesh.assemble.ContactAssembler (a FacetAssembler subclass) using the element_energy hook.

Note on unsupported/manual methods:

• Lagrange multipliers for Dirichlet are not provided as a helper class; they must be implemented manually by augmenting the matrix.
• Plain periodic boundary conditions (zero-phase) can be handled by BlochReducer, but general non-matched periodic meshes require manual DOF identification and row/column merging.

The TensorMesh User Guide provides in-depth, topic-oriented documentation for the FEM (Finite Element Method) pipeline. Key areas of focus include:

• Core Concepts: Understanding the FEM pipeline and module interactions.
• Meshes: Building, inspecting, loading/saving meshes, and managing per-node/per-cell data.
• Elements & Quadrature: Working with basis functions, quadrature rules, and element types.
• Forms: Writing weak forms in pure Python using assembler base classes.
• Mixed Assembly: Implementing multi-field block systems (e.g., Taylor-Hood Stokes or coupled physics).
• Boundary Conditions: Applying Dirichlet BCs via static condensation and handling Neumann BCs.
• Sparse Solvers: Performing linear and nonlinear sparse solves via torch-sla (supporting five backends, batched RHS, and Newton/Picard/Anderson methods).
• Time Integration: Using explicit and implicit-linear Runge-Kutta schemes for transient problems.
• Batched Workflows: Leveraging memory chunking, batched RHS, and ML datasets.
• Differentiability: Computing end-to-end gradients through the assemble $\rightarrow$ solve pipeline for inverse problems and topology optimization.
• Distributed FEM: Partitioning meshes and assembling per GPU to solve systems across multiple devices via torch-sla.

TensorMesh is organized into several functional modules covering the full finite element analysis (FEA) and physics-informed machine learning pipeline. Use the following modules to build your simulation or learning workflow:

Core FEA Workflow

• Mesh: Manage mesh data structures and geometric operations.
• Element: Define finite elements and handle quadrature rules.
• Assemble: Perform matrix and vector assembly routines.
• Operator: Define differential operators and apply boundary conditions.
• Sparse: Execute sparse matrix operations and linear solvers.

Physics & Integration

• ODE: Handle time integration and Ordinary Differential Equation solvers.
• Functional: Define functional forms and variational methods.
• Material: Access isotropic material models and library presets.
• Distributed: Implement multi-GPU mesh partitioning and parallel assembly.

Machine Learning & Data

• Neural Network (nn): Use neural network modules specifically designed for physics-informed learning.
• Dataset: Access built-in datasets and specialized data loaders.

Utilities

• Optimizer: Use optimization algorithms for solving inverse problems.
• Visualization: Utilize plotting and visualization tools for results analysis.

A map of the core modules in the tensormesh package:

The TensorMesh Example Gallery provides runnable scripts that map mathematical formulations to concrete code. Examples range from FEM building blocks to research-grade simulations in fluid, solid, and electromagnetic mechanics.

Available example categories include:

• Basics & Visualization: Basis points, shape functions, element ordering, and mesh generation.
• PDE Solves:
  • Poisson Equation: Elliptic PDEs in 2D/3D, batched RHS, and h-adaptive refinement.
  • Diffusion: Transient heat equations and nonlinear Allen-Cahn phase-field evolution.
  • Wave Equation: Hyperbolic PDEs with explicit central-difference time stepping.
  • Complex-Valued FEM: Time-harmonic Helmholtz equations with complex coefficients.
• Specialized Physics:
  • Phononic Crystals: Bloch-Floquet band structures and slab transmission.
  • Open-Domain Wave: PML (Perfectly Matched Layers) and plane-wave ports for acoustic and photonic resonators.
  • Modal Analysis: Generalized eigenproblems for dielectric waveguide modes and cantilever vibration.
  • Solid Mechanics: Linear elasticity, hyperelasticity, contact, and plasticity.
  • Fluid Mechanics: Incompressible Navier-Stokes using Taylor-Hood mixed elements (Stokes, cavity, cylinder flow, etc.).
  • Magnetostatics: 3D Maxwell equations using stabilized nodal curl-curl formulations.
• Advanced Applications:
  • Inverse Design & Identification: Differentiable FEM for coefficient-field identification and topology optimization.
  • Physics-Informed Learning: Training neural networks to minimize the assembled Galerkin residual ||K u - F||².
  • ML Datasets: Batch generation of heat and wave snapshots for training neural operators.

The notebooks are categorized by problem type to help you find specific FEM implementations and workflows:

Basics

• basics.ipynb: Elements, Lagrange basis functions, and mesh generation.
• poisson.ipynb: The standard workflow: Mesh $\rightarrow$ Assembler $\rightarrow$ Condenser $\rightarrow$ Solve.

Scalar Problems

• poisson_3d.ipynb: 3D Poisson problem on tetrahedra.
• poisson_h_adaptivity.ipynb: h-adaptive refinement (solve $\rightarrow$ estimate $\rightarrow$ mark $\rightarrow$ remesh).
• helmholtz.ipynb: Complex-valued FEM.

Time-Dependent Problems

• heat.ipynb: Implicit Euler with factorised stepping operators.
• wave.ipynb: Central differences with energy-conservation monitoring.
• allen_cahn.ipynb: Nonlinear phase field using Newton solves.

Fluids (Mixed Assembly)

• stokes_taylor_hood.ipynb: P2-P1 Taylor-Hood mixed elements.
• cavity.ipynb: Steady Navier-Stokes via Picard iteration.
• flow_obstacles.ipynb: Flow past obstacles using CSG meshing.

Solid Mechanics

• cantilever_beam.ipynb: 3D linear elasticity.
• hyperelastic_beam.ipynb: Large deformation using strain energy and autograd/LBFGS.

Differentiable FEM

• coefficient_identification.ipynb: Inverse problems solved via autograd through the FEM solve.
• poisson_galerkin.ipynb: Physics-informed learning using the assembled Galerkin residual.