Documentation · Examples · Changelog

torch-fem

torch-fem is a GPU-accelerated differentiable finite element solver for solid mechanics and heat transfer built on PyTorch. Automatic differentiation provides exact sensitivities of simulation results with respect to material parameters, geometry, loads, etc. without hand-derived adjoint formulations. It is aimed at researchers in computational mechanics who need gradients through FEM solvers for tasks such as optimization, inverse problems, and machine-learning-augmented simulation.

Features

• Differentiable: exact gradients of simulation results through efficient adjoint solves, for linear and nonlinear problems alike
• GPU-accelerated: a linear elastic model with a million degrees of freedom solves in seconds on a single GPU
• Broad scope: solid mechanics and heat transfer from trusses to shells and 3D solids, including nonlinear materials and large deformations
• Extensible: a new hyperelastic material needs only its energy function and a new element only its shape functions, as autograd derives the rest
• PyTorch-native: models are plain tensors, so FEM results plug directly into optimizers and neural networks

Installation

You may install torch-fem via pip with

pip install torch-fem[notebook]

The notebook extra adds Jupyter widgets, interactive PyVista rendering, and animation support for the example notebooks. For GPU acceleration, install PyTorch with CUDA support - see the installation guide for details.

Minimal example

This is a minimal example of how to use torch-fem to solve a very simple planar cantilever problem.

import torch
from torchfem import Planar
from torchfem.materials import IsotropicElasticityPlaneStress

torch.set_default_dtype(torch.float64)

# Material
material = IsotropicElasticityPlaneStress(E=1000.0, nu=0.3)

# Nodes and elements
nodes = torch.tensor([[0., 0.], [1., 0.], [2., 0.], [0., 1.], [1., 1.], [2., 1.]])
elements = torch.tensor([[0, 1, 4, 3], [1, 2, 5, 4]])

# Create model
cantilever = Planar(nodes, elements, material)

# Load at tip [Node_ID, DOF]
cantilever.forces[5, 1] = -1.0

# Constrained displacement at left end [Node_IDs, DOFs]
cantilever.constraints[[0, 3], :] = True

# Show model
cantilever.plot(node_markers=True, node_labels=True)

This creates a minimal planar FEM model:

# Solve for displacements u, internal forces f, stresses σ, deformation gradients F, and material states α
u, f, σ, F, α = cantilever.solve()

# Plot displacement magnitude on deformed state
cantilever.plot(u, node_property=torch.norm(u, dim=1))

This solves the model and plots the result:

To compute gradients through the FEM model, we mark the variables that require gradients. The backward pass of solve() is an adjoint solve via the implicit function theorem rather than a differentiation through solver iterations, as explained in the differentiability guide.

# Enable automatic differentiation
cantilever.thickness.requires_grad = True
u, f, _, _, _ = cantilever.solve(differentiable_parameters=cantilever.thickness)

# Compute sensitivity of compliance w.r.t. element thicknesses
compliance = torch.inner(f.ravel(), u.ravel())
torch.autograd.grad(compliance, cantilever.thickness)[0]

This returns the sensitivity of the compliance with respect to the thickness of each element:

tensor([-0.0208, -0.0053])

Both entries are negative, so adding material anywhere stiffens the structure, but the element at the clamped end is about four times as effective as the one at the tip.

The getting started guide continues from here to a minimal topology optimization, and the theory section summarizes the underlying mechanics and finite element formulation.

Examples

The subdirectories examples/basic and examples/optimization contain Jupyter notebooks for trusses, planar problems, shells, and solids, from basic analyses to topology, shape, and orientation optimization. A selection is shown below, and the full set is in the example gallery.

Plasticity in a plate with hole: Isotropic linear hardening model for plane-stress or plane-strain.
Finite strain cantilever: Hyperelastic model in Total Lagrangian Formulation.
Topology optimization of a jet engine bracket: The optimized part is cut out of the design space at an iso-value of the density.	Topology optimization of a 3D heat sink: Conductive material is distributed in a cube with a homogeneous heat source to minimize thermal compliance.
Recovery of a property field: A direct optimization and a neural field recover a graded elastic modulus from noisy observations of the displacement.	Free size optimization of a pressure vessel: Each element's shell thickness is a design variable and a fixed amount of material is redistributed to minimize compliance.

Performance

Performance is a priority concern in torch-fem: it is benchmarked on several CPUs and GPUs for the forward solve, the backward pass, and peak memory. The scripts in benchmarks/ reproduce all results, and the performance documentation reports the full set of problems and machines, including thermal and topology optimization problems from the mosaic benchmark suite and a finite-strain hyperelastic problem.

As an example, a linear elastic cube of hexahedral elements with $N$ nodes along each edge is clamped on one face and pulled on the opposite one, and the backward pass computes the gradient of the displacement sum with respect to the nodal forces. On an NVIDIA RTX 5090 in float64, using CG with a Jacobi preconditioner that suffices for this well-conditioned problem, this takes:

Alternatives

torch-fem focuses on solid mechanics and thermal problems. It provides sensitivities through PyTorch autograd, which makes it easy to drop into optimization loops and ML pipelines. It is the natural choice if you are working in the PyTorch ecosystem. Depending on your needs, one of these Python FEM tools may serve you better:

Library	Stars	Focus	Differentiable	Consider it over torch-fem when…
FEniCSx (DOLFINx)		General PDEs, UFL weak forms, MPI	via dolfin-adjoint	you need arbitrary weak forms or massively parallel distributed runs
SfePy		General multiphysics, pure Python	—	you need a broad range of PDE applications on CPU
JAX-FEM		Differentiable FEM, JAX / GPU	Yes	your stack is built on JAX rather than PyTorch
Firedrake		General PDEs, UFL weak forms	via pyadjoint	you want a UFL form language with automated adjoints for multiphysics
scikit-fem		Lightweight assembly, NumPy/SciPy	—	you want minimal dependencies and full control over custom forms
FElupe		Finite-strain solid mechanics	partially via tensortrax	you work with hyperelastic / finite-strain solids
Nutils		High-order / immersed methods	—	you research advanced or immersed discretizations including IGA
PyTorch-FEA		Biomechanics, PyTorch	Yes	you work on soft-tissue / inverse biomechanics

Not sure which to pick? The mosaic differentiable-physics benchmark suite compares several of these solvers on gradient accuracy and forward/adjoint performance under a common interface.

Citing torch-fem

If you use torch-fem in your research, please cite it as follows:

@software{torchfem,
    author = {Meyer, Nils},
    title  = {torch-fem: GPU accelerated differentiable finite elements for solid mechanics with PyTorch},
    doi    = {10.5281/zenodo.20306384},
    url    = {https://github.com/meyer-nils/torch-fem},
}

Contributing

Contributions are welcome! Please check out the contributing guide for the development workflow. Bug reports, feature requests, and usage questions are all welcome in the issue tracker - see the support guide for what to include.