KernelPack-JAX: a meshless differentiable GPU framework for node generation and forward and inverse PDE problems

Abstract

We present KernelPack-JAX, a meshless, differentiable, GPU framework in JAX for solving forward and inverse PDE problems in domains whose boundaries, including internal obstacles, are given only by samples of their smooth pieces. Our node generator is based on a new stacked representation of piecewise-smooth boundaries, in which a parametric partition of unity of each smooth piece induces the supports of a curl-free implicit partition of unity of the whole boundary. This representation keeps the edges and corners of the boundary and classifies every point of space without a global fit. Boundary and interior nodes are then selected by a parallel greedy independent-set algorithm that reproduces sequential Poisson-disk sampling exactly. We discretize elliptic and parabolic problems with radial basis function-generated finite differences (RBF-FD) on stencils selected by a new visibility heuristic to avoid crossing the boundary, with ghost nodes for the boundary conditions. In addition, our framework exposes a GPU sparse direct solver to JAX as a differentiable linear solve, so that Jacobians of time-dependent problems cost a small multiple of a forward solve. We derive error estimates that separate consistency from stability and bound the parameter error by the forward error. We demonstrate high-order convergence for Darcy flow, reaction-diffusion, and advection-diffusion in 2D and 3D domains, including one with an internal obstacle, and show that the recovered parameters largely follow the forward errors. We also present timings for every stage.