Rapid Cardiac Simulations Using JAX

Introduction: : Myocardial infarction (MI) is a leading cause of death affecting nearly 1 million people annually, resulting in the death of nearly 30% of those affected. High fidelity simulations of MI can play a pivotal role in understanding the cascade of events that lead to progressive heart failure and develop novel therapies. While the accuracy of current computational cardiac models has improved markedly, they remain infeasible for direct clinical use due to prohibitively long execution times. In recent years GPU technology has made major advances thanks to the rise of machine learning and artificial intelligence. While parallelization across CPUs has been accomplished in both open source and commercial codes (e.g. FEniCSx), GPUs provide substantially faster speed gains in solving the large linear systems of equations commonly encountered in cardiac mechanics. Google’s JAX (https://github.com/jax-ml/jax/) is a library that allows for easy deployment of Python code onto GPUs to vectorize and parallelize them. In the present work, we extended JAX-FEM [1] a JAX-based finite element library, to focus on cardiac mechanics by creating CARDIAX. CARDIAX was developed using an established ovine biventricular model by solving a pressure-volume (PV) loop over the cardiac cycle. Simulation results were then compared with a similar model run in ABAQUS [2]. The performance of CARDIAX was also demonstrated on a benchmark problem to show the speeds of CARDIAX, JAX-FEM, and FEniCSx.

Materials and

Methods: : The ovine cardiac model was developed using experimental data which consisted of a comprehensive dataset from a single ovine heart, including pressure volume loops and diffusion tensor magnetic image resonance (DTMRI) data. An anisotropic material model was used where the fiber directions were obtained through the DTMRI data. Simulations were run to solve the PV loop, showing the importance of compressibility in the myocardium where the volumetric penalty decreases in the systolic phase. The PV loop was solved by determining the amount myofibers must contract to match a given volume at a certain pressure, which means solving many inverse problems in sequence. The CARDIAX code takes full advantage of GPU processing power through vectorization and just-in-time compilation via the JAX library. We also took advantage of computing the adjoint to obtain gradients, and other derivatives are obtained through automatic differentiation. The workflow to solve the chain of inverse problems is now much more streamlined (Figure 1) and used a more realistic spring-like boundary condition compared with the original approach. The PV loop was solved for the healthy state and with an apical infarction. Along with comparing total execution times, we also created a basic benchmark problem to more directly compare JAX-FEM, FEniCSx, and CARDIAX in a standardized manner. The benchmark is a unit cube with a neo-Hookean material model loaded in compression. The cube is then discretized finer and finer to increase the number of degrees of freedom for the problem.

Results, Conclusions, and Discussions:: The simulations performed in CARDIAX revealed similar trends to those done in ABAQUS. The level of active contraction increases in the infarcted heart to match the same PV loop as the healthy heart. The levels of compressibility are consistent as well, where the tissue compresses more near the endocardial surface compared to the epicardium. Not only did we get consistent results with the ABAQUS model, but we were able to obtain them in a much faster time. The time it took to solve the simulations in ABAQUS took well over a day, whereas the CARDIAX simulation took approximately 48 minutes. Moreover, we were able to simulate 100 data points on the PV loop as opposed to the ABAQUS simulation which only solved for 20. For the benchmark problem of the neo-Hookean cube, CARDIAX was the fastest solver for problems ranging from 1000s to 100000s degrees of freedom (Figure 2). It was able to solve smaller problems more than 30x faster than FEniCSx, and it could solve the larger problems more than 25x faster than the GPU code JAX-FEM. This shows the advantage of solving finite element problems with current GPU technology. We now have a package that can solve problems faster, is entirely in Python, and gives the user as much control as they desire. In doing so, simulations have the potential to start being used in clinical setting by removing the bottleneck of solve time.

Acknowledgements and/or References (Optional):: Acknowledgements:
The author gratefully acknowledges the financial support from NIH grant HL073021.

References:
1. Xue, T, et al., Comp Phys Comm 291 (2023): 108802
2. Liu, H, et al. Sci Rep 11 (2021): 13466