JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science

Abstract

This paper introduces JAX-FEM, an open-source differentiable finite element method (FEM) library. Constructed on top of Google JAX, a rising machine learning library focusing on high-performance numerical computing, JAX-FEM is implemented with pure Python while scalable to efficiently solve problems with moderate to large sizes. For example, in a 3D tensile loading problem with 7.7 million degrees of freedom, JAX-FEM with GPU achieves around 10× acceleration compared to a commercial FEM code depending on platform. Beyond efficiently solving forward problems, JAX-FEM employs the automatic differentiation technique so that inverse problems are solved in a fully automatic manner without the need to manually derive sensitivities. Examples of 3D topology optimization of nonlinear materials are shown to achieve optimal compliance. Finally, JAX-FEM is an integrated platform for machine learning-aided computational mechanics. We show an example of data-driven multi-scale computations of a composite material where JAX-FEM provides an all-in-one solution from microscopic data generation and model training to macroscopic FE computations. The source code of the library and these examples are shared with the community to facilitate computational mechanics research. Program summaryProgram Title: JAX-FEMCPC Library link to program files:https://doi.org/10.17632/hgwshjbcw6.1Developer's repository link:https://github.com/tianjuxue/jax-am/tree/main/jax_am/femLicensing provisions: GPLv3Programming language: PythonNature of problem: Implementation of the finite element method (FEM) with several appealing features that classic FEM software typically does not have: realized with pure Python; running on CPU/GPU; differentiable for solving PDE-constrained optimization problems; seamless integration with machine learning.Solution method: Our framework JAX-FEM is based on Google JAX, a high-performance numerical computing library with automatic differentiation features and supporting both CPU/GPU. Unlike classic FEM software written in Fortran or C/C++, JAX-FEM is implemented with pure Python and can easily be installed as a Python package. We demonstrate our software by solving problems including forward PDE prediction, inverse design/optimization, and data-driven analysis.