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.
Related
- Benchmarking Optimizers to Solve Inverse Problems with Differentiable Physics Simulators
- Parameter-dependent inhomogeneous boundary-value problems in spaces of continuously differentiable functions
- CraftSPH: A composable framework for differentiable smoothed particle hydrodynamics implemented in PyTorch
- JAXIGA - A software framework for GPU-enabled and differentiable isogeometric analysis
- DeepTelecom: A Differentiable Multimodal Framework for Channel Twinning in Dynamic Heterogeneous Wireless Environments
BibTeX
@misc{cattaneo2026kernelpackjax,
title = {{KernelPack-JAX: a meshless differentiable GPU framework for node generation and forward and inverse PDE problems}},
author = {Alberto Cattaneo and M. Keith Ballard and Robert M. Kirby and Grady B. Wright and Varun Shankar},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.25484},
eprint = {2610.25484},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.25484},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - KernelPack-JAX: a meshless differentiable GPU framework for node generation and forward and inverse PDE problems
AU - Alberto Cattaneo
AU - M. Keith Ballard
AU - Robert M. Kirby
AU - Grady B. Wright
AU - Varun Shankar
PY - 2026
DA - 2026/10/18
PB - arXiv
DO - 10.48550/arxiv.2610.25484
UR - https://arxiv.org/abs/2610.25484
AB - 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.
ER -