A Differentiable, GPU-Accelerated Finite-Element Formulation in JAX for Large-Deformation Shells with Contact: Application to Patient-Specific Aortic Valve Simulation

Patient-specific aortic-valve models can help assess how tissue properties and geometry influence leaflet deformation and coaptation, with potential to inform intervention planning. Large deformations, near-incompressibility, and evolving contact make these simulations computationally demanding and can challenge nonlinear convergence. Differentiable, GPU-accelerated methods offer efficient derivative evaluation and parallel computation, but their matrix-free application to soft-tissue shells with contact remains less established. We developed a forward solver using tatva's energy-based finite-element framework in JAX. The solver combines finite-strain extensible-director shells, hyperelastic constitutive models, and a finite-range repulsive contact energy. Nonlinear equilibrium is solved using a regularized, matrix-free Newton-Krylov method with two-level preconditioning. Automatic differentiation evaluates tangent-vector products without forming or storing the primary tangent matrix, while GPU execution parallelizes element and quadrature-point computations. Constitutive, shell, contact, and self-contact responses were assessed against analytical, semi-analytical, and published benchmarks. At a contact-free plate equilibrium, implicit sensitivities of the converged apex displacement to constitutive parameters, thickness, and pressure agreed with finite differences from independent nonlinear simulations. Simulations of three patient-specific aortic-valve models (two tricuspid and one bicuspid) achieved final relative equilibrium residuals below 10^-4. Median GPU runtimes were approximately 15-25 minutes, with 5.3-6.5-fold speedups relative to corresponding CPU runs on the tested hardware. These results provide component-level verification and demonstrate GPU-accelerated forward valve simulation in the tested cases, supporting future gradient-based inverse studies.