TorchFSM

Fourier Spectral Method with PyTorch

[ Documentation & Examples]

TL;DR

TorchFSM is a PyTorch-based library for solving PDEs using Fourier spectral method. It is designed for physics-based deep learning and differentiable simulations.

Simple Example

Build a Burger equation with provided operators (the Burgers equation is also available at torchfsm.pde):

from torchfsm.operator import Operator, Convection, Laplacian

def Burgers(nu:float) -> Operator:
    return nu*Laplacian()-Convection()
burgers=Burgers(0.01)

Simulate the Burgers equation at 1D mesh:

from torchfsm.mesh import MeshGrid
from torchfsm.plot import plot_traj
from torchfsm.traj_recorder import AutoRecorder
import torch

device='cuda' if torch.cuda.is_available() else 'cpu'
L=1.0; N=128; 
mesh=MeshGrid([(0,L,N)],device=device)
x=mesh.bc_mesh_grid()
u_0=torch.sin(2*torch.pi*x/L)+0.5

traj=burgers.integrate(u_0=u_0,mesh=mesh,
    dt=0.01,step=200,
    trajectory_recorder=AutoRecorder(),
)
plot_traj(traj,animation=False)

Directly simulate the Burgers on 2D mesh with the same operator:

mesh=MeshGrid([(0,L,N),(0,L,N)],device=device)
x,y=mesh.bc_mesh_grid()
u_0=torch.cat([torch.sin(2*torch.pi*x/L)*torch.cos(4*torch.pi*y/L),
               torch.cos(2*torch.pi*x/L)*torch.sin(4*torch.pi*y/L)],
              dim=1)

traj=burgers.integrate(u_0=u_0,mesh=mesh,
    dt=0.01,step=50,
    trajectory_recorder=AutoRecorder(),
)

plot_traj(traj,animation=False)

Feature

• Modular by design: TorchFSM offers a modular architecture with essential mathematical operators—like divergence, gradient, and convection—so you can build custom solvers like stacking building blocks, quickly and intuitively.

• GPU-accelerated: TorchFSM leverages GPU computing to speed up simulations dramatically. Run complex 3D PDEs in minutes, not hours, with seamless hardware acceleration.

• Batched simulation support: Built on PyTorch, TorchFSM enables batched simulations with varied initial conditions—ideal for parameter sweeps, uncertainty quantification, or ensemble analysis.

• Differentiable and ML-ready: Fully differentiable by design, TorchFSM integrates naturally with machine learning workflows—for residual operators, differentiable physics, or dataset generation.

Installation

• Install through pip: pip install torchfsm or via conda: conda install conda-forge::torchfsm
• Install the development version through pip: pip install git+https://github.com/qiauil/torchfsm@dev
• Install locally: Download the repository using git clone https://github.com/qiauil/torchfsm.git, cd torchfsm, and run pip install .

Documentations

Check 👉 here.

Citation

We will be happy if you use TorchFSM in your research. If you find our work useful, please consider citing our repository:

@software{liu_2025_15350210,
  author       = {Liu, Qiang and
                  Koehler, Felix and
                  Thuerey, Nils},
  title        = {TorchFSM: Fourier Spectral Method with PyTorch},
  month        = may,
  year         = 2025,
  publisher    = {Zenodo},
  doi          = {10.5281/zenodo.15350210},
  url          = {https://doi.org/10.5281/zenodo.15350210},
}

Further Information

• TorchFSM is mainly built by Qiang Liu, Felix Koehler, and Nils Thuerey at Physics Based Simulation Group, Technical University of Munich . We are grateful for all contributors!
• If you are more familiar with Jax, please check our Jax alternative Exponax