Bringing PDEs to JAX with forward and reverse modes automatic differentiation

Partial differential equations (PDEs) are used to describe a variety of physical phenomena. Often these equations do not have analytical solutions and numerical approximations are used instead. One of the common methods to solve PDEs is the finite element method. Computing derivative information of the solution with respect to the input parameters is important in many tasks in scientific computing. We extend JAX automatic differentiation library with an interface to Firedrake finite element library. High-level symbolic representation of PDEs allows bypassing differentiating through low-level possibly many iterations of the underlying nonlinear solvers. Differentiating through Firedrake solvers is done using tangent-linear and adjoint equations. This enables the efficient composition of finite element solvers with arbitrary differentiable programs. The code is available at github.com/IvanYashchuk/jax-firedrake.Comment: Published as a workshop paper at ICLR 2020 DeepDiffEq worksho

Neuro-symbolic partial differential equation solver

We present a highly scalable strategy for developing mesh-free neuro-symbolic partial differential equation solvers from existing numerical discretizations found in scientific computing. This strategy is unique in that it can be used to efficiently train neural network surrogate models for the solution functions and the differential operators, while retaining the accuracy and convergence properties of state-of-the-art numerical solvers. This neural bootstrapping method is based on minimizing residuals of discretized differential systems on a set of random collocation points with respect to the trainable parameters of the neural network, achieving unprecedented resolution and optimal scaling for solving physical and biological systems.Comment: Accepted for publication at NeurIPS 2022 (ML4PS workshop). arXiv admin note: substantial text overlap with arXiv:2210.1431

Neural Vortex Method: from Finite Lagrangian Particles to Infinite Dimensional Eulerian Dynamics

In the field of fluid numerical analysis, there has been a long-standing problem: lacking of a rigorous mathematical tool to map from a continuous flow field to discrete vortex particles, hurdling the Lagrangian particles from inheriting the high resolution of a large-scale Eulerian solver. To tackle this challenge, we propose a novel learning-based framework, the Neural Vortex Method (NVM), which builds a neural-network description of the Lagrangian vortex structures and their interaction dynamics to reconstruct the high-resolution Eulerian flow field in a physically-precise manner. The key components of our infrastructure consist of two networks: a vortex representation network to identify the Lagrangian vortices from a grid-based velocity field and a vortex interaction network to learn the underlying governing dynamics of these finite structures. By embedding these two networks with a vorticity-to-velocity Poisson solver and training its parameters using the high-fidelity data obtained from high-resolution direct numerical simulation, we can predict the accurate fluid dynamics on a precision level that was infeasible for all the previous conventional vortex methods (CVMs). To the best of our knowledge, our method is the first approach that can utilize motions of finite particles to learn infinite dimensional dynamic systems. We demonstrate the efficacy of our method in generating highly accurate prediction results, with low computational cost, of the leapfrogging vortex rings system, the turbulence system, and the systems governed by Euler equations with different external forces

Simulation von Fluidströmungen durch datengetriebene Evolution von Wirbelpartikeln

Fluid solvers that provide accurate and fast fluid simulations are of great importance in many scientific and engineering disciplines. Conventional numerical solvers based on the Eulerian description of the flow provide highly accurate solutions to the Navier-Stokes equations. However, there is typically a significant amount of computational effort is required to execute such Eulerian simulations. On the other hand, fluid solvers built on the Lagrangian description of the flow are more appealing in terms of its vicinity to the true physics, since it treats the actual fluid particles as the primary computational elements. A particular group of Lagrangian particle methods based on vorticity, instead of velocity, as the primary flow variable, delivers velocity field solutions, which are always divergence-free. These vortex methods have an inherent advantage that the particles need to be present only in the regions where vorticity exists, and therefore fewer fluid particles are required to execute simulations as compared to their counterparts with velocity-based formulations. Recently, deep learning solutions for fluid dynamics problems by the application of artificial neural networks has become more prominent. Neural networks encode the information about the governing laws of fluid dynamics in its parameters using the knowledge extracted from data samples during training. The aim of this work is to use deep learning to learn fluid dynamics with Lagrangian vortex particles as the primary flow representation. Solution strategies to train and evaluate the neural networks for predicting Lagrangian vortex particle dynamics for different flow scenarios are presented throughout this work. Conceptualization and implementation of an approach to model interaction between vortex particles based on the Taylor series expansion of the velocity form the core of this work. We demonstrate that our trained neural networks produce fluid simulations with reasonable accuracy for different flow scenarios while respecting appropriate constraints pertaining to fluid dynamics

Lagrangian Neural Style Transfer for Fluids

Artistically controlling the shape, motion and appearance of fluid simulations pose major challenges in visual effects production. In this paper, we present a neural style transfer approach from images to 3D fluids formulated in a Lagrangian viewpoint. Using particles for style transfer has unique benefits compared to grid-based techniques. Attributes are stored on the particles and hence are trivially transported by the particle motion. This intrinsically ensures temporal consistency of the optimized stylized structure and notably improves the resulting quality. Simultaneously, the expensive, recursive alignment of stylization velocity fields of grid approaches is unnecessary, reducing the computation time to less than an hour and rendering neural flow stylization practical in production settings. Moreover, the Lagrangian representation improves artistic control as it allows for multi-fluid stylization and consistent color transfer from images, and the generality of the method enables stylization of smoke and liquids likewise.Comment: ACM Transaction on Graphics (SIGGRAPH 2020), additional materials: http://www.byungsoo.me/project/lnst/index.htm