DeepAI AI Chat
Log In Sign Up

JAX-DIPS: Neural bootstrapping of finite discretization methods and application to elliptic problems with discontinuities

by   Pouria Mistani, et al.

We present a scalable strategy for development of mesh-free hybrid neuro-symbolic partial differential equation solvers based on existing mesh-based numerical discretization methods. Particularly, this strategy can be used to efficiently train neural network surrogate models for the solution functions and operators of partial differential equations while retaining the accuracy and convergence properties of the state-of-the-art numerical solvers. The presented neural bootstrapping method (hereby dubbed NBM) is based on evaluation of the finite discretization residuals of the PDE system obtained on implicit Cartesian cells centered on a set of random collocation points with respect to trainable parameters of the neural network. We apply NBM to the important class of elliptic problems with jump conditions across irregular interfaces in three spatial dimensions. We show the method is convergent such that model accuracy improves by increasing number of collocation points in the domain. The algorithms presented here are implemented and released in a software package named JAX-DIPS (, standing for differentiable interfacial PDE solver. JAX-DIPS is purely developed in JAX, offering end-to-end differentiability from mesh generation to the higher level discretization abstractions, geometric integrations, and interpolations, thus facilitating research into use of differentiable algorithms for developing hybrid PDE solvers.


page 4

page 9

page 19

page 21

page 22

page 24


Neuro-symbolic partial differential equation solver

We present a highly scalable strategy for developing mesh-free neuro-sym...

Large-scale Neural Solvers for Partial Differential Equations

Solving partial differential equations (PDE) is an indispensable part of...

M2N: Mesh Movement Networks for PDE Solvers

Mainstream numerical Partial Differential Equation (PDE) solvers require...

Message Passing Neural PDE Solvers

The numerical solution of partial differential equations (PDEs) is diffi...

Performance prediction of finite-difference solvers for different computer architectures

The life-cycle of a partial differential equation (PDE) solver is often ...

Learning a Mesh Motion Technique with Application to Fluid-Structure Interaction and Shape Optimization

Mesh degeneration is a bottleneck for fluid-structure interaction (FSI) ...