Differentiable Spectral Element Method (DSEM)

A PyTorch framework for PDE-constrained optimization using differentiable spectral element methods.

Overview

This repository implements the first differentiable spectral element method (DSEM) that embeds Galerkin spectral element solvers into automatic differentiation (AD) computational graphs. Unlike traditional adjoint-based approaches, DSEM leverages native reverse-mode AD to compute gradients through spectral element forward solves, eliminating the need for manually derived adjoint equations.

Key Features

• 1D, 2D, and 3D spectral element solvers with Gauss-Lobatto-Legendre quadrature
• Native PyTorch autograd support — gradients through linear solves via implicit function theorem
• Matrix-free operator evaluation via sum factorization for memory efficiency
• Complex geometry support through multi-element decompositions (e.g., L-shaped domains)
• Spectral convergence — exponential error decay for smooth solutions
• Gradient accuracy — agreement with finite differences to 8+ decimal places

Method

The spectral element method decomposes the domain into elements and uses high-order polynomial bases within each element. The forward solve $u = K^{-1} M f$ is implemented as a differentiable PyTorch operation. Reverse-mode AD automatically computes gradients through the solve, equivalent to solving the discrete adjoint system.

Theoretical results

• Theorem 1: AD gradients are identical to discrete adjoint gradients
• Theorem 2: Condition number of the optimality system: $\kappa(H) = O(p^4 h^{-4}/\alpha + 1)$
• Theorem 3: $hp$-convergence with exponential rates for analytic solutions

Installation

pip install torch numpy scipy matplotlib

Quick Start

from spectral_element_2d import DifferentiableSEM2D
import torch

# 2x2 elements on [0,1]^2, polynomial order 4
elements = [(0,0.5,0,0.5), (0.5,1,0,0.5), (0,0.5,0.5,1), (0.5,1,0.5,1)]
solver = DifferentiableSEM2D(elements, p=4)

# Forward solve (differentiable)
f = torch.randn(solver.n_interior, dtype=torch.float64, requires_grad=True)
u = solver.solve(f)

# Gradient flows back through the solve
loss = u.sum()
loss.backward()
print(f.grad)  # gradient w.r.t. control f

Project Structure

├── code/
│   ├── spectral_element.py      # 1D differentiable SEM
│   ├── spectral_element_2d.py   # 2D differentiable SEM
│   ├── spectral_element_3d.py   # 3D differentiable SEM
│   ├── run_3d.py                # 3D experiments
│   ├── plot_figures.py          # Figure generation
│   └── test_*.py                # Test scripts
├── paper/
│   ├── main.tex                 # LaTeX source
│   └── figures/                 # PDF and PNG figures
└── research_plan.md             # Research plan

Citation

If you use this code, please cite:

@article{wang2026dsem,
  title={Differentiable Spectral Element Methods for PDE-Constrained Optimization},
  author={Wang, Fuchang and Cao, Huirong},
  journal={Journal of Computational Physics},
  year={2026}
}

License

MIT