TorchFSM, A Differentiable PDE Solver Built on PyTorch

We’re excited to announce the release of TorchFSM, a lightweight, flexible library for solving partial differential equations (PDEs) using the Fourier spectral method, fully integrated with PyTorch.

TorchFSM is designed with modern research in mind, especially for those working in physics-informed machine learning, differentiable simulations, and scientific computing.

Why TorchFSM?

🔧 Modular by design: Build your own PDEs using intuitive operator blocks like divergence, gradient, curl, and Laplacian. Mix and match them to create complex residuals or simulations with ease.

⚡ GPU-accelerated: Powered by PyTorch and optimized with FFTs, TorchFSM can simulate 3D PDEs in just a few minutes on modern GPUs.

🧠 Differentiable out of the box: TorchFSM integrates seamlessly with PyTorch’s autograd. It’s ideal for residual-based learning, PINNs, or generating synthetic datasets for deep learning models.

📦 Built for batched simulations: Run simulations across multiple initial conditions in parallel using PyTorch’s batching capabilities—perfect for ensemble experiments or parameter sweeps.

Get Started Ready to explore? Head over to the GitHub repo to get started: 👉 https://github.com/qiauil/torchfsm

Prefer JAX? Check out our JAX-based solver, Exponax: 👉 https://github.com/Ceyron/exponax