TorchFSM: A Differentiable Fourier Spectral Solver for Machine Learning–Enabled PDE Simulation
Please login to view abstract download link
As the convergence of machine learning and numerical simulation accelerates, there is an increasing demand for simulation tools that are both computationally efficient and natively compatible with modern AI frameworks. We present TorchFSM, an open-source software package designed to bridge the gap between traditional partial differential equation (PDE) solvers and the machine learning ecosystem. TorchFSM is based on the Fourier Spectral Method $^{[1]}$, enabling highly accurate spatial discretization for a broad class of PDEs. Built on a PyTorch backend, the framework adopts a modular architecture that provides core mathematical operators—such as gradients, divergence, and convection—as intuitive, composable building blocks. This design enables the rapid development of custom solvers while significantly reducing the implementation overhead typically associated with complex physical models. Beyond usability, TorchFSM is engineered for high-performance scientific computing. By leveraging native GPU acceleration, the framework enables high-fidelity three-dimensional simulations to be executed in minutes rather than hours. In addition, TorchFSM supports batched simulation execution, a critical capability for parameter studies and uncertainty quantification that is often absent from legacy solvers. Meanwhile, TorchFSM's fully differentiable feature allows physical solvers to be seamlessly integrated into machine learning workflows, enabling gradient-based optimization for inverse problems and the training of physics-informed neural operators $^{[2]}$. By providing a flexible, ML-ready, and open-source environment, TorchFSM aims to empower the scientific community to develop the next generation of high-performance, differentiable multiphysics simulations.
