Gpu-native software speeds up simulation and learning for complex physics models
diff-soup: GPU-native Differentiable Simulation and Inference for Large-scale Stiff Dynamical Systems
Mathematical Software
Summary
Simulating physical processes often involves solving complex math problems that computers find hard and slow, especially when these involve stiff dynamical systems. To make this faster and easier to learn from, the authors created diff-soup, a software tool that uses powerful graphics processors and advanced math methods to speed up both running these simulations and adjusting models based on data. Diff-soup works well for very large and complex problems and helps connect traditional physics equations with modern machine learning. The authors tested it on tough examples and showed it to be much faster than existing tools.
What this means in practice
- •For computational engineers: Run large-scale physics simulations faster and compute gradients efficiently to improve model calibration and design optimization.
- •For chemical process developers: Train neural models of catalyst activity in complex reactors using differentiable simulation to accelerate reaction design cycles.
Authors
Samuel Degnan-Morgenstern, Martin Z. Bazant
Abstract
Learning physical models from spatiotemporal data requires simulating systems of partial differential equations (PDEs) and differentiating through their solution. Upon spatial discretization, these models give rise to stiff dynamical systems that require repeatedly solving large, sparse linear systems. Traditional scientific computing frameworks are largely CPU-based and require problem-specific tuning to accelerate, scale, and differentiate. In this work, we introduce diff-soup, a differentiable, sparse-oriented, unified PDE suite for GPU-native simulation and inference of large, stiff dynamical systems. Built in JAX, diff-soup provides 1) a fully differentiable design for problems with sparse structure, 2) scaling to large problems through the NVIDIA cuDSS sparse direct solver with model and data multi-GPU parallelism, and 3) end-to-end PDE-constrained optimization, from classical statistical inference to embedded deep learning. On stiff PDE benchmarks, diff-soup outperforms established CPU- and GPU-based frameworks by one to two orders of magnitude in both forward solves and gradients. We demonstrate both a classical Bayesian inference workflow, identifying constitutive relationships, and a modern deep learning workflow, training a neural field for the catalyst activity of a three-dimensional packed bed reactor. We provide diff-soup as an open-source repository to facilitate research in computational science and engineering.