4.7 Article

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

期刊

COMPUTER PHYSICS COMMUNICATIONS
卷 291, 期 -, 页码 -

出版社

ELSEVIER
DOI: 10.1016/j.cpc.2023.108802

关键词

Differentiable simulation; Neural networks; Design and optimization

向作者/读者索取更多资源

This paper introduces JAX-FEM, an open-source differentiable finite element method (FEM) library that is implemented with pure Python and scalable for solving moderate to large problems. It achieves significant acceleration compared to commercial FEM codes and enables automatic solving of inverse problems. JAX-FEM also serves as an integrated platform for machine learning-aided computational mechanics.
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 10x 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.& COPY; 2023 Elsevier B.V. All rights reserved.

作者

我是这篇论文的作者
点击您的名字以认领此论文并将其添加到您的个人资料中。

评论

主要评分

4.7
评分不足

次要评分

新颖性
-
重要性
-
科学严谨性
-
评价这篇论文

推荐

暂无数据
暂无数据