JAXIGA - A software framework for GPU-enabled and differentiable isogeometric analysis
Abstract
In this paper, we present JAXIGA, an open-source library for isogeometric analysis based on JAX. The library separates the construction of the mesh from the computations which are traced for compilation and differentiation. This allows geometry evaluation, assembly and the supported solvers to be combined in a just-in-time compiled program, for which first-order reverse-mode derivatives can be computed. We consider three numerical methods obtained from a single pointwise energy density. The Galerkin and energy-minimization methods use the same discrete potential and have the same stationary points, while collocation enforces the corresponding strong form. The sensitivities of a converged solution are obtained by solving one adjoint problem, and their accuracy depends on the tolerances of both the primal and adjoint problems. We also describe hierarchical extraction for local refinement of the variational methods and the use of smooth spline spaces for fourth-order problems. The numerical results are compared with analytical solutions and independent FEniCSx and JAX-FEM computations. For the reported GPU benchmark with 36,448 unknowns, assembly is twenty-four times faster than on the host CPU. When the conjugate-gradient solver is also run on the GPU, the total solution time is reduced from 31.8 to 1.38 seconds, and the result agrees with the direct solution to 5×10−11. These results show that the choice of solver and the reuse of compilation are important for obtaining an overall reduction in computational cost.
Related
- CraftSPH: A composable framework for differentiable smoothed particle hydrodynamics implemented in PyTorch
- KernelPack-JAX: a meshless differentiable GPU framework for node generation and forward and inverse PDE problems
- DeepTelecom: A Differentiable Multimodal Framework for Channel Twinning in Dynamic Heterogeneous Wireless Environments
- Software Defect Prediction: A Systematic Review and a Conceptual Framework for Stable, Actionable Explanations (AISD-F)
- Movable Antenna Enabled Anti-Jamming: From Local Optimality Guaranteed Design to a Unified Antenna Position Optimization Framework
BibTeX
@misc{anitescu2026jaxiga,
title = {{JAXIGA - A software framework for GPU-enabled and differentiable isogeometric analysis}},
author = {Cosmin Anitescu and Yongjie Jessica Zhang},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.09083},
eprint = {2610.09083},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.09083},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - JAXIGA - A software framework for GPU-enabled and differentiable isogeometric analysis
AU - Cosmin Anitescu
AU - Yongjie Jessica Zhang
PY - 2026
DA - 2026/10/06
PB - arXiv
DO - 10.48550/arxiv.2610.09083
UR - https://arxiv.org/abs/2610.09083
AB - In this paper, we present JAXIGA, an open-source library for isogeometric analysis based on JAX. The library separates the construction of the mesh from the computations which are traced for compilation and differentiation. This allows geometry evaluation, assembly and the supported solvers to be combined in a just-in-time compiled program, for which first-order reverse-mode derivatives can be computed. We consider three numerical methods obtained from a single pointwise energy density. The Galerkin and energy-minimization methods use the same discrete potential and have the same stationary points, while collocation enforces the corresponding strong form. The sensitivities of a converged solution are obtained by solving one adjoint problem, and their accuracy depends on the tolerances of both the primal and adjoint problems. We also describe hierarchical extraction for local refinement of the variational methods and the use of smooth spline spaces for fourth-order problems. The numerical results are compared with analytical solutions and independent FEniCSx and JAX-FEM computations. For the reported GPU benchmark with 36,448 unknowns, assembly is twenty-four times faster than on the host CPU. When the conjugate-gradient solver is also run on the GPU, the total solution time is reduced from 31.8 to 1.38 seconds, and the result agrees with the direct solution to 5×10^−11. These results show that the choice of solver and the reuse of compilation are important for obtaining an overall reduction in computational cost.
ER -