Differentiable Spectral Methods for PDE-Constrained Optimization

PyTorch implementation of differentiable Chebyshev spectral solvers for PDE-constrained optimization. Gradients with respect to control or parameter variables are obtained by reverse-mode automatic differentiation (AD), eliminating the need to derive and implement adjoint equations manually.

Paper

Fuchang Wang, Huirong Cao. Differentiable Spectral Methods for PDE-Constrained Optimization. Journal of Scientific Computing, submitted 2026.

Key Features

• Machine-precision gradients: AD gradient is mathematically identical to the discrete adjoint gradient (Theorem 1)
• Exponential convergence: spectral accuracy for analytic solutions (Theorem 3)
• No adjoint derivation required: only a forward solver and loss.backward()
• Dual condition-number structure: explicit estimates for state-space and control-space systems (Theorem 2)
• Regularization error estimate: closed-form prediction matches experiments (Corollary 1)

Installation

pip install torch numpy scipy matplotlib

Requires Python 3.8+ and PyTorch 2.0+.

Quick Start

from spectral_solvers import DiffusionSolver
import torch

# 1D diffusion: -u'' = f, u(0)=u(1)=0
solver = DiffusionSolver(N=32, domain=(0, 1))

# Define objective
def objective(f):
    u = solver(f)
    u_d = torch.sin(torch.pi * solver.x_interior)
    return 0.5 * ((u - u_d)**2).sum() + 0.001 * 0.5 * (f**2).sum()

# Optimize with L-BFGS
f = torch.zeros(solver.N_interior, requires_grad=True)
optimizer = torch.optim.LBFGS([f])

def closure():
    optimizer.zero_grad()
    loss = objective(f)
    loss.backward()
    return loss

for _ in range(50):
    optimizer.step(closure)

Code Structure

code/
├── spectral_solvers.py          # Differentiable PDE solvers (diffusion, Burgers, Stokes, etc.)
├── diff_chebyshev.py            # Differentiable Chebyshev differentiation
├── run_diffusion_control.py     # Experiment 1: 1D diffusion optimal control
├── run_burgers_control.py       # Experiment 2: Burgers equation optimal control
├── run_coefficient_inversion.py # Experiment 3: diffusion coefficient inversion
├── run_stokes_control.py        # Experiment 4: 2D Stokes flow control
├── run_adjoint_comparison.py    # Experiment 5: AD vs hand-derived adjoint
├── run_diffusion_3d.py          # Experiment 6: 3D diffusion control
├── run_convection_diffusion.py  # Experiment 7: convection-diffusion control
├── run_multiple_runs.py         # Reproducibility: multiple runs with different seeds
└── plot_paper_figures.py        # Generate all paper figures

Reproducing Experiments

cd code

# Run all experiments
python run_diffusion_control.py
python run_burgers_control.py
python run_coefficient_inversion.py
python run_stokes_control.py
python run_adjoint_comparison.py
python run_diffusion_3d.py
python run_convection_diffusion.py

# Generate figures
python plot_paper_figures.py

Results are saved as JSON files in results/ and figures in figures/.

Test Problems

#	Problem	PDE	Dimension	Key Result
1	Diffusion control	$-u''=f$	1D	Gradient error 5.7e-9, 294x faster than FD
2	Burgers control	$u_t+uu_x=\nu u_{xx}+f$	1D	Gradient error 1.3e-9, 21x faster than FD
3	Coefficient inversion	$-(a(x)u')'=f$	1D	Constant coeff error <0.1%, robust to noise
4	Stokes control	$-\Delta u+\nabla p=f$	2D	Velocity error 0.046% (closed-loop test)
5	AD vs adjoint	$-u''=f$	1D	Gradients identical (9.4e-9), AD 12x faster at N=128
6	3D diffusion control	$-\Delta u=f$	3D	Error 0.99% vs theory 0.87%
7	Convection-diffusion	$\epsilon u''-u'=f$	1D	Gradient error 2.3e-10, spectral convergence 191x

Theoretical Results

• Theorem 1: AD gradient = discrete adjoint gradient (at any point where the constraint Jacobian is invertible)
• Theorem 2: Well-posedness of discrete optimality system; dual condition-number structure $\kappa(\alpha A^\top A + I)$ vs $\kappa(\alpha A^\top A + I)$
• Theorem 3: Exponential convergence for analytic data
• Corollary 1: Regularization error $|u^*-u_d|/|u_d| = \alpha\lambda/(1+\alpha\lambda)$

Citation

@article{wang2026diff,
  title={Differentiable Spectral Methods for PDE-Constrained Optimization},
  author={Wang, Fuchang and Cao, Huirong},
  journal={Journal of Scientific Computing},
  year={2026},
  note={submitted}
}

License

MIT License