CARDIAX-NNFE - A Cardiac Mechanics SciML Framework

Abstract

One goal of Scientific Machine Learning (SciML) is to advance traditional scientific computing frameworks with modern machine learning tools. This includes extending established methods, such as the finite element method, with cardiac function applications due to their complexity and need for very rapid execution times for real time clinical use. In this work, we present an advanced form of the Neural Network Finite Element (NNFE) method specialized for cardiac simulations, termed CARDIAX-NNFE. The NNFE method learns the parameter-to-displacement field map by training over the residual of the hyperelastic material PDE, using the domain represented by finite elements. The implementation is developed in Python using JAX to leverage its automatic differentiation, highly parallel GPU, and JIT-compilation capabilities. To demonstrate CARDIAX-NNFE effectiveness, we trained full cardiac pressure-volume responses using a simplified heart model, spanning the entire cardiac physiological functional range. Results indicated the ability to simulate a family of pressure-volume solutions with average nodal positional error of 0.023 mm and maximal error of 0.054 mm, with a single complete PV loop evaluated in 0.002 seconds. The CARDIAX-NNFE software platform thus provides for a robust platform for cardiac functional simulations. Moreover, it provides the structure for residual-based SciML methods, which can apply to a variety of physics-based biomedical problems that require high execution speed for clinical applications.