A neural network finite element method for contact mechanics

Computational biomechanical models of organ systems utilizing the finite element (FE) method have been extensively applied to the study of normal and pathophysiological function. While a robust method and providing for many unique pathophysiological insights, traditional FE methods remain prohibitively slow for real-time many-query clinical applications. To meet these demanding computational requirements, we have developed a general neural network finite element (NNFE) approach for hyperelastic soft tissue organ simulations that can produce equivalent accuracy within clinically relevant time frames. However, many organ systems involve contact between structural elements (e.g. heart valve leaflets), which has not been previously addressed in the NNFE and related approaches. In the present work, we developed a NNFE-based method for contact between elastic deformable bodies. We exploited the fact that stable equilibrium solutions in nonlinear hyperelasticity minimize the potential energy to train the neural network, with contact incorporated through an potential energy penalty. This penalty was discretized and computationally implemented in JAX (Google, Inc.) in a way that could be statically allocated and jit compiled for rapid computation on modern GPU hardware. Following our previous NNFE formulation, we represented the problem domain geometry and enforced the necessary boundary conditions using a Non-Uniform Rational B-Splines (NURBS) based geometry scheme. NURBS were chosen for their ability to smoothly represent body and surface geometries. To demonstrate the method, we developed and verified two basic contact problems using a hyperelastic material model: indentation and flap. Both problems were simulated with a fully connected neural network with a parallel linear layer. Training continued until the gradient of the energy with respect to network parameters dropped below 10-2, taking less than six hours in each case studied. Once trained, the NNFE approach was capable of accurately predicting the equilibrium configurations of representative contact problems. Moreover, the NNFE contact approach required nearly six order of magnitude less time to evaluate than an equivalent FEniCSx implementation that utilized a similar mesh of cubic order Lagrange elements (0.01 s compared to 1715 seconds, respectively). Moreover, the solutions were trained over full ranges of loading and were thus able to faithfully represent families of solutions without requiring retraining. While addressing basic problems in contact, this study serves as crucial stepping stone to generate NNFE organ simulations for rapid prediction of patient specific solutions that involve contact.