JAXIGA - A software framework for GPU-enabled and differentiable isogeometric analysis

1. Introduction

• Spline bases represent geometry and PDE solutions, with NURBS exactly representing conics and suitable continuity supporting fourth-order weak formulations and design optimization.Higher continuity can also reduce degrees of freedom as polynomial degree increases; sensitivity computation through geometry, assembly, and the discrete solution is needed for inverse problems and optimization.Source: §1 · pp. 1–2

1.1. Related work and software

• JAXIGA shares interface conventions with JAX-FEM but uses NURBS and hierarchical splines rather than Lagrange elements; its energy method uses splines rather than the neural-network trial space of the deep energy method.The paper also situates the library among IGA and general discretization frameworks, including PetIGA, GeoPDEs, G+Smo, and Python-based packages.Source: §1.1 · p. 2
• Earlier GPU assembly studies use different problems, implementations, and hardware, so they do not directly compare total solution cost with JAXIGA.The paper also identifies weighted quadrature, sum factorization, and low-rank assembly as possible alternatives to its current procedures.Source: §1.1 · p. 2
• The capability comparison notes that some frameworks offer broader mesh handling, element libraries, adaptivity, solvers, and parallel scaling, while JAXIGA targets a smaller set of problem classes with a common extension interface.The comparison defines framework-native differentiation through the solver and native GPU support as distinct capabilities.Source: §1.1 · p. 3, Table 1 · p. 3

1.2. Contribution and scope

• JAXIGA derives Galerkin, energy-minimization, and collocation methods from one energy definition and supports first-order reverse-mode sensitivities through compiled geometry evaluation, assembly, and supported solvers.The mesh construction and refinement are kept separate from computations traced by JAX; differentiation is supported under stated conditions.Source: §1 · p. 2, §1.2 · p. 3
• The variational methods support locally refined spaces, while collocation is currently restricted to tensor-product spaces.The current implementation also has limitations discussed in Section 5.Source: §1.2 · p. 4
• The experiments compare computed solutions with analytical results and independent FEniCSx and JAX-FEM computations, and check sensitivities against finite differences.The paper also examines discretization accuracy and implementation cost.Source: §1.2 · p. 4

2.1. Spline discretization

• Spline discretization uses tensor-product B-splines and rational weights for geometry and solution fields; Bézier extraction also accommodates truncated hierarchical refinement.Interior-knot multiplicity controls continuity, while knot insertion and degree elevation preserve the geometry map.Source: §2.1 · p. 4

2.2. Variational formulation

• The variational formulation combines internal energy and prescribed loads into a discrete potential; its gradient and Hessian give the residual and tangent, and quadratic potentials yield linear systems.For nonconvex potentials, stationarity does not imply minimality, and Newton iteration and energy minimization may reach different stationary points.Source: §2.2 · p. 5

2.3. Methods and coupling

• Galerkin solves for a stationary point of the shared discrete potential, energy minimization searches for its minimum, and collocation enforces the associated strong form at Greville points.Neumann and interface flux equations can make the collocation system overdetermined, in which case it is solved in the least-squares sense.Source: §2.3 · pp. 5–6
• Conforming multipatch coupling identifies matching trace degrees of freedom and gives C0 continuity across interfaces; nonconforming and C1 coupling are not implemented.The matching requires compatible trace bases and orientations, along with coincident control points.Source: §2.3 · p. 6

2.4. Solution strategies

• Linear Galerkin systems use sparse LU, a dense direct solver, or Krylov methods; nonlinear Galerkin, energy minimization, and collocation use Newton, L-BFGS, and Gauss–Newton strategies, respectively.Automatic CG selection is restricted to qualifying positive-definite linear problems; the GPU size threshold was calibrated on the plate benchmark and may not generalize to other degrees or hardware.Source: §2.4 · p. 6

3.1. The static/traced separation

• JAXIGA separates host-side mesh setup from traced computation, allowing compiled programs to be reused when parameters change on a fixed mesh.Changing the mesh requires rebuilding setup data and may require recompilation; connectivity and other discretization data are stored as fixed metadata.Source: §3 · p. 6, §3.1 · p. 7, Figure 1 · p. 7

3.2. Geometry and multipatch coupling

• Gmsh cells become tensor-product patches coupled with C0 continuity, while native NURBS retain their knots, weights, control net, and within-patch continuity.CAD-aware enrichment refits boundary samples, so general CAD input remains approximate; differentiation begins from the resulting fixed patch layout.Source: §3.2 · p. 8, §3.3 · pp. 8–9

3.4. Problems: one physics definition

• A problem combines a function space and boundary conditions with a pointwise energy density, from which flux and mass maps can be derived.For nonvariational problems, flux and mass maps may instead be specified directly for collocation.Source: §3.4 · p. 9

3.5. Deriving three methods

• Galerkin solves for a stationary point of the discrete potential, energy minimization uses that potential, and collocation solves its corresponding strong-form residual.Galerkin tangents come from the energy Hessian for user-defined energies, while collocation appends interface-flux and Neumann equations on multipatch spaces.Source: §3.5 · p. 10

3.6. Solvers and implicit differentiation

• First-order reverse-mode sensitivities require a converged, smooth solution on a fixed discretization, and their accuracy depends on both primal and adjoint convergence.One adjoint system suffices regardless of parameter count; for nonlinear least-squares collocation, the implicit correction uses the full Hessian rather than only the Gauss–Newton term.Source: §3.6 · p. 11

3.7. Postprocessing

• Physical-point evaluation differentiates through inverse-map Newton iteration while holding the host-selected element index fixed.The sensitivity is local: the point must remain inside that element and inversion must converge with a nonsingular geometry Jacobian; a changed containing element requires a new host search.Source: §3.7 · p. 12

3.8. Hierarchical refinement

• THB-splines enable local refinement while preserving partition of unity and the original geometry, with mirrored interface refinement and class-2 admissibility controlling the hierarchy.Adaptivity runs in an outer Python loop because marking and refinement change array shapes; gradients do not propagate across refinement cycles. The volume-only residual indicator is unreliable for general problems with repeated knots, interfaces, coefficient discontinuities, or boundary residuals.Source: §3.8 · pp. 12–13

3.9. Fourth-order problems

• Maximally smooth degree-p splines provide C1 continuity for primal fourth-order formulations, while the framework also supports matrix-free tangents, batching, and differentiated time-stepping.For shells, parametric first and second derivatives feed the energy; matrix-free Krylov storage is O(ndofs), and Newmark-β and generalized-α trajectories are differentiated through lax.scan.Source: §3.9 · pp. 13–14, §3.10 · p. 14

4.1. Convergence of the Galerkin discretization; 4.2. Imported CAD and NURBS geometries

• Galerkin Poisson convergence approached p + 1 in one and two dimensions; the examples also demonstrated imported CAD plate and native NURBS connecting-rod geometries.For p = 2, 3, 4, the finest one-dimensional rates were 3.00, 4.00, 5.00 and the two-dimensional rates 3.02, 4.02, 5.02; the three-dimensional rates remained slightly above asymptotic values. The plate boundary-distance measure fell from 3.16×10^-3 to 1.35×10^-5 after CAD refitting, and the rod's 25 patches were coupled across 39 interfaces.Source: Table 4 · p. 15, §4.1 · pp. 15–16, §4.2 · p. 16

4.3. Agreement between the three methods

• Across nine Poisson cases, Galerkin and energy minimization differed by at most 5.1 × 10^-15, while collocation was less accurate and converged at degree-dependent rates.For a hyperelastic problem, Newton and energy minimization differed by 1.5 × 10^-13; on a curved NURBS annulus, two strong-form residual evaluations agreed to 10^-15.Source: §4.3 · pp. 17–18, Table 5 · p. 17

4.4. Verification against an independent implementation

• On the 36,448-unknown plate benchmark, A100 assembly was 24× faster, while device-based CG reduced total time to 1.38 s and agreed with the direct solution to 5 × 10^-11.The full solution timing depends on running the linear solver on the device; the reported JAXIGA-versus-FEniCSx and JAX-FEM comparisons also show expected convergence, and JAXIGA degree four reached 7.8 × 10^-5 L2 error with 2,112 unknowns versus FEniCSx P3 at 1.1 × 10^-4 with 4,800 unknowns.Source: §4.4 · pp. 18–19, Table 6 · p. 19
• For the Kirsch plate, adaptive refinement reached 3.4 × 10^-5 energy-norm error with 3,856 degrees of freedom, versus about 5,400 for uniform refinement.The expected stress concentration factor of three was recovered; separate checks found basis data agreeing with a NumPy reference within 2 × 10^-16 to 4 × 10^-15 and verified geometric and differential identities.Source: §4.4 · pp. 20–21, Figure 5 · p. 20

4.5. Differentiability; 4.6. Computational cost

• Adjoint gradients differed from central finite differences by at most 2.3 × 10^-9, and each backward pass solved exactly one adjoint system.Compilation dominates single-call costs at these problem sizes, whereas compiling the full parameter-to-solution function removes additional Python dispatch and setup work; gradient time was close to twice solution time.Source: §4.5 · pp. 21–22, Table 8 · p. 21, §4.6 · p. 22

4.7. Inverse problem: permeability identification

• In the Darcy permeability inverse problem, L-BFGS reduced the objective from 7.3 × 10^-5 to 4.1 × 10^-7 in 30 iterations, recovering the field with 0.59% relative L2 error.The compiled, differentiated objective included the Darcy solve and pressure evaluation at scattered points, with one adjoint problem per gradient evaluation.Source: §4.7 · p. 22, Figure 6 · p. 23

4.8. Material identification from full-field measurements

• Under nested five-fold cross-validation over 90 specimens, inversion-derived features increased mean Dice from 0.801 to 0.839 (paired gain 0.038; 95% bootstrap interval [0.012, 0.070]) and material-pair accuracy from 94.4% to 97.8%.JAXIGA required 31 s per specimen on one Apple M1 Pro CPU core and 16 s on an NVIDIA A10 GPU using Jacobi-preconditioned conjugate gradients with device-side solves. Direct thresholding overestimated inclusion width; the C0 model fit no better (mean misfit 0.343 in both cases) and lowered thresholded-field Dice from 0.721 to 0.710.Source: §4.8 · pp. 23–25

4.9. Adaptive local refinement on a corner singularity

• For the L-shaped corner singularity, uniform refinement has energy-norm convergence O(N^-1/3) regardless of degree, while grading toward the corner can recover O(N^-p/2).The refinement procedure solves the discrete problem, evaluates a residual indicator, marks elements by 50% Dörfler bulk marking, and refines the marked region while mirroring refinement across interfaces.Source: §4.9 · p. 25
• Adaptive refinement gives fitted slopes −1.13 and −1.56 for p=2,3 versus uniform refinement near −1/3; at p=3 it is four times more accurate with approximately twelve times fewer unknowns.Rates are least-squares fits over the final eight cycles; Dörfler marking refines in bursts, making shorter fitting windows sensitive to their position within a step.Source: §4.9 · pp. 25–26, Figure 8 · p. 26, Table 11 · p. 26
• Boundary-data interpolation limits p=3 accuracy: error stalls at 3.7 × 10^-3 although the estimator continues to fall, so that run starts from two additional uniform refinements.The residual estimator omits Dirichlet-data error; the p=2 run is unaffected, but resolving the boundary data is necessary for p=3.Source: §4.9 · pp. 26–27

4.10. A fourth-order problem: the Kirchhoff–Love shell

• The Kirchhoff–Love shell uses an H2-conforming displacement space; on a flat patch, its solution agrees with the plate solution to relative error 4.7 × 10^-15.The linear shell strains are obtained by forward-mode differentiation, while the geometrically nonlinear shell uses the strains themselves.Source: §4.10 · p. 27
• 0.30059 is the converged downward displacement magnitude −uz at A for p=4 on 8×8 elements (432 degrees of freedom), versus the reference value 0.3024.The p=2 displacement is 0.3002 on 32×32 elements; its slower convergence is consistent with membrane locking.Source: §4.10 · p. 27, Table 12 · p. 28

4.11. Assembly variants, batching and time integration

• The subsection examines assembly and batching procedures with a Poisson problem, then tests time integration and spatially varying stiffness identification using longitudinal waves in an elastic bar.Source: §4.11 · p. 28

Assembly and batching checks

• Quadratic splines on a uniform 16 × 16 mesh with a0 = 1 are used to compare matrix-free tangent products and solutions with assembled counterparts.The Poisson problem has exact solution u = sin(2πx) sin(2πy)/a0.Source: Assembly and batching checks · p. 28
• On an 8 × 8 mesh, vmap solutions for eight a0 values from 0.5 to 4.0 are compared with separate solves, with round-off agreement expected.The source term is fixed; chunked assembly uses eight elements per chunk.Source: Assembly and batching checks · p. 29
• Table 13 compares alternative assembly and solution paths in its first four rows, then checks analytical frequencies and discrete-energy conservation.The supplied table caption identifies these checks but gives no cell values.Source: Table 13 · p. 29

Free vibration of an elastic bar

• The model is an undamped elastic bar with longitudinal displacement and no body force; its strain-energy density is ψ = 1/2 E(x)u_,x^2.After imposing Dirichlet constraints, the Galerkin discretization uses free displacement coefficients d, consistent mass matrix M and boundary-load vector b.Source: Free vibration of an elastic bar · pp. 29–30
• The maximum relative frequency error is 2.6×10−7 for cubic splines (p = 3) on 32 uniform elements.The check uses a bar with L = 1, E = 2 and ρ = 3, with both ends fixed, and compares the first five natural frequencies with the generalized eigenvalue problem.Source: Free vibration of an elastic bar · p. 30
• The discrete energy is conserved to a relative difference of 5.1×10−15, approximately round-off error.The free-vibration check uses the Newmark method with β = 1/4, γ = 1/2, ∆t = 2×10−3 and 1,000 steps.Source: Free vibration of an elastic bar · p. 30

Identification of a stiffness inclusion

• The inversion estimates contrast α and center xc from an end-velocity trace, differentiating the discrete time trajectory through 360 Newmark steps with reverse-mode gradients.The synthetic observations are noise-free and start from α = 0.20 and xc = 0.50.Source: Identification of a stiffness inclusion · pp. 30–31
• 7.9% of the observed trace norm is the homogeneous-bar difference available to identify the inclusion; cycle skipping can mislead distant initial guesses.The misfit oscillates as the reflected pulse arrival time changes.Source: Identification of a stiffness inclusion · p. 31

4.12. Phase-field fracture

• The fracture example models a plane-strain plate with a damage-seeded notch; ϕ is zero for intact material and one for fully broken material.The notch is represented by damage rather than removed from the geometry.Source: §4.12 · p. 31, Formulation · p. 31
• The degradation function weights elastic strain energy, while a running-maximum history field of tensile strain energy introduces damage irreversibility.The degradation function has residual stiffness k ≪ 1.Source: Formulation · p. 32
• Fully broken crack faces carry no compressive stiffness in this treatment, so crack closure is not represented.Source: Formulation · p. 32
• The staggered load loop alternates between the displacement and damage equations and updates the history field, but does not generally correspond to coordinate descent of the stated energy.When ψ− is nonzero, setting H = ψ+ does not recover that energy's damage equation.Source: Formulation · pp. 32–33

• Peak loads are 771, 641 and 616 N for ℓ/h = 1, 2 and 4, respectively; the coarsest discretization does not fracture in the loading range.The authors do not regard the finest curve alone as evidence of spatial convergence.Source: Discretization and staggered solution · p. 33
• The 1,200-increment solve takes 2,422 staggered sweeps, reusing basis data and sparsity patterns; SuperLU factors are reused only when iterative refinement meets a 10^-5 residual tolerance.The load schedule uses 44 increments of 10^-4 mm followed by 1,156 of 10^-6 mm.Source: Discretization and staggered solution · p. 33
• 616 N at 4.86 × 10^-3 mm is within 0.7% of the reference's 620 N at 5.07 × 10^-3 mm; reaction falls to about 12 N.The final damage field and load–displacement response are shown in Figure 12.Source: Results · p. 33
• The reference curve's model is unidentified, so the softening difference cannot be attributed to mesh resolution alone; quantitative validation requires matched models and load-step and spatial convergence studies.The reference implementation contains multiple second- and fourth-order phase-field models and degradation functions.Source: Results · p. 34

• For unknown crack paths, refinement triggers where quadrature-point damage exceeds 1/2 and expands within ℓ in 2D or 1.5ℓ in 3D; rebuilding the discretization adds cost.The compared tests use different loading and mesh settings for tension and shear.Source: Crack-following adaptivity · p. 34
• Tension peaks differ by 0.23 N (642.74 versus 642.51 N), but adaptive runtime is 247 s versus 126 s and final mesh size is 2,740 versus 2,920 elements.Source: Crack-following adaptivity · p. 34
• Shear peak loads agree within 0.1% (551.06 versus 551.59 N), while 543 versus 1,810 s gives a 3.3× reduction using 1,306–4,213 adaptive elements versus 25,600 uniform.Source: Crack-following adaptivity · p. 34
• At 12.44 µm, reactions are 11.25 and 253.91 N, a 43.99% difference relative to the uniform peak; outside two increments, the largest discrepancy is 1.18%.The adaptive load drop precedes the uniform-mesh drop by two increments.Source: Crack-following adaptivity · p. 34

• At the largest case, factorization takes over 20× longer than assembly; Jacobi-preconditioned CG avoids fill-in, but iteration counts need not be mesh-independent.Source: Three dimensions · p. 34, Table 15 · p. 35
• At ℓ = 1/16 mm, the 1×0.2×1 mm block peaks at 11.63 N, between thickness-scaled plane-stress and plane-strain results of 10.19 and 11.91 N.The block has a seeded through-thickness crack and traction-free lateral faces.Source: Three dimensions · p. 34, Figure 14 · p. 36
• Halving ℓ while keeping ℓ/h = 2 raises the peak to 14.87 N; the finer run uses 17,740 elements and 56,112 displacement unknowns.Remeshing, setup and recompilation take 742 s of the 2,083 s total, or 36%.Source: Three dimensions · p. 36

5. Discussion

• For repeated solves on a fixed mesh, compilation can be reused, while device-side implicit gradients and linear algebra reduce sensitivity-analysis and solution costs.Adaptive refinement can reduce unknowns, but rebuilding the discretization and recompiling must also be accounted for.Source: §5 · p. 36
• Accuracy work remains for boundary-data projection, error indicators incorporating boundary and interface residuals, and fracture convergence studies using the same model in both implementations.Cost comparisons among stronger preconditioners or assembly methods should use the same solution and gradient accuracy, and sensitivities remain subject to the conditions in Section 3.6.Source: §5 · pp. 36–37
• Multipatch coupling requires conforming interfaces and provides only C0 continuity, so primal fourth-order problems must use suitable single patches.Collocation is limited to tensor-product spaces with dense normal equations; local hierarchical refinement is unsupported with collapsed-boundary constraints.Source: §5 · p. 36

6. Conclusions

• GPU assembly is accelerated by 24×, while total solution time falls by 23× relative to the reported CPU direct solver when a suitable device solver is used.The authors conclude that reducing assembly cost alone is insufficient for a comparable reduction in total computational cost.Source: §6 · p. 37
• JAXIGA obtains three numerical methods from a common energy density, with spline geometry and first-order implicit sensitivities.Computed solutions were compared with analytical results and independent implementations, and sensitivities were checked against finite-difference approximations.Source: §6 · p. 37
• Hierarchical extraction enables local refinement of the variational methods, and smooth single-patch spaces are used for fourth-order problems.Source: §6 · p. 37

Appendix A. Installation and Usage

• Installation requires Python 3.10 or later, JAX, NumPy, SciPy, and optax; VTK output, plotting, and Gmsh dependencies are optional.The stated editable-install command is pip install -e ".[post,plot]".Source: Appendix A · p. 37
• Sixty-four-bit floating point is enabled on import for the accuracy needed by the reported discretization and gradient checks.Source: Appendix A · p. 37
• The complete Poisson example uses Galerkin, and the same problem can use energy minimization or collocation by changing the method argument.The example solves the problem, evaluates its L2 error against an exact solution, and writes VTK output.Source: Appendix A · pp. 37–38

Appendix B. Geometry Import and Refinement

• For 2D CAD input, the meshing procedure attempts quadrilateral algorithms and can subdivide remaining triangles; 3D tetrahedra are subdivided into hexahedra when needed.The sizing heuristic halves its target size until quality and sampled CAD-distance checks pass, but does not minimize cell count or certify geometric error.Source: Appendix B · p. 38
• High-order conversion fits Bernstein coefficients to Gmsh nodal coordinates, while CAD-aware enrichment projects boundary Greville samples to determine control points.Polynomial cell checks combine sampled distortion with Bernstein bounds on the Jacobian determinant, rejecting cells when bounds are inconclusive; shared traces must remain compatible.Source: Appendix B · p. 38
• The rational-geometry option reconstructs supported conic and cylindrical maps, which ordinary homogeneous refinement then preserves exactly.Bundled native NURBS data let the standalone plate, connecting-rod, and blade examples run without MATLAB or IGAPack.Source: Appendix B · p. 38