CARDIAX: A JAX-based platform for Rapid Cardiac Functional Simulations

Computational pipelines of increasing sophistication are actively being developed for both basic studies and patient-specific cardiac simulations for direct clinical utilization. A major barrier continues to be simulation execution times, limiting levels of scale integration and performing such simulations in clinically relevant time frames. While machine-learning based surrogate models are an active research area, they remain early in development. As an alternative approach we have developed CARDIAX, a JAX-based traditional finite element platform that leverages GPU technologies for accelerated cardiac simulations. A high-fidelity cardiac model was developed in the CARDIAX framework to simulate a complete pressure-volume loop in a normal and infarcted heart. The cardiac model was based on an extant comprehensive dataset acquired from a single ovine heart, which accounted for realistic biventricular geometry and fiber orientations, myocardial mechanical behavior including physiologically realistic compressibility behavior, and realistic kinematic constraints. We were able to run simulations considerably faster in the JAX framework using GPUs compared to similar CPU based software platforms making at least a 10 fold improvement in runtime. In a cardiac disease model, CARDIAX demonstrated a mean increase in active stress of ∼30 kPa in the presence of an apical infarction compared to the normal case. To better conduct objective performance benchmarking, we developed and utilized a standard benchmark using an cuboidal test problem with a hyperelastic isotropic material model applied to demonstrate the performance characteristics of each software platform evaluated. CARDIAX exhibited well over a ten fold speed gain across a range of mesh sizes. Regarding scalability, for larger problem sizes with 100,000 degrees of freedom CARDIAX demonstrated more than 35 times faster execution time. We conclude that use of GPU-based software technologies can be effective platforms for high speed cardiac mechanics simulations.