Source-linked AI summary
JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science
Tianju Xue, Shuheng Liao, Zhengtao Gan, Chanwook Park, Xiaoyu Xie, Wing Kam Liu, Jian Cao
TL;DR
JAX-FEM addresses the need for efficient and less labor-intensive PDE-based computational mechanics by combining GPU-accelerated FEM, automatic differentiation, and machine learning in one open-source library. It solves forward and inverse problems, including topology optimization, and reports substantial GPU speedups on a 7,703,841-DOF problem while supporting integrated data-driven workflows.
Problem
PDE-based computational analysis requires efficient numerical algorithms and human-resource use, while inverse sensitivity derivations and multi-tool machine-learning workflows can be difficult or cumbersome.
Method
JAX-FEM uses pure-Python FEM on Google JAX, combining GPU-accelerated forward solving, automatic differentiation for inverse design, and integrated machine-learning functionality.
Results
JAX-FEM solved a 7,703,841-DOF problem in 523 s on GPU, achieving 16.1× and 9.1× acceleration over Abaqus on CPU and with MPI, respectively.
Takeaways & Limitations
JAX-FEM provides an open-source platform for forward and inverse FEM research and data-driven computational mechanics.
Takeaways & Limitations
The reported inverse problems are deterministic, and a single 48 GB NVIDIA GPU supports around 10 million DOF before memory becomes insufficient.
Abstract
from arXiv · showhide
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$\times$ 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.
1 Introduction
JAX-FEM is an open-source FEM library built on JAX to improve computational efficiency and human productivity in PDE-based analysis. It combines GPU-accelerated forward solving, differentiable inverse design, and integrated machine-learning workflows in a pure-Python framework.
- JAX-FEM targets both computational efficiency and human productivity in PDE-based computational analysis.The paper frames overall efficiency as depending on hardware use and efficient use of human resources.
- Feature 1: GPU acceleration enables efficient forward PDE solutions while retaining a pure-Python frontend for moderate-sized problems.JAX-FEM uses JAX’s XLA backend to provide competitive performance without restricting users to small problems.
- Feature 2: Automatic differentiation supports inverse design without requiring manual sensitivity derivations for nonlinear forward problems.The paper identifies sensitivity computation as a key but non-trivial part of gradient-based inverse optimization.
- Feature 3: JAX-FEM integrates machine learning with computational mechanics to avoid transferring data among separate simulation, training, and FEM tools.Its JAX foundation provides access to machine-learning functionality within the same platform.
- The library is presented as an open-source platform whose code and examples support research in computational mechanics.The paper also describes comparisons with other FEM software and example applications across forward, inverse, and machine-learning workflows.
2 Solving forward problems
JAX-FEM solves representative three-dimensional elasticity, hyperelasticity, and plasticity problems using automatic differentiation and GPU-oriented array programming. Its results agree with FEniCSx, while GPU execution substantially accelerates large finite-element problems.
- Problem classes and verification: JAX-FEM formulates representative three-dimensional solid-mechanics problems and verifies them against FEniCSx for linear elasticity, hyperelasticity, and plasticity.The verification uses uniaxial tensile loading with fixed bottom and prescribed top displacement boundaries.
- Implementation features: Array programming with JAX vectorizes element operations instead of explicitly writing conventional finite-element loops, supporting GPU acceleration.The implementation uses jax.vmap to perform vectorized operations.
- Implementation features: Automatic differentiation computes element Jacobians directly from the weak form, avoiding manual derivation of difficult tangent tensors.This is especially useful for constitutive models such as plasticity, where the fourth-order tangent tensor is nontrivial to derive.
- Verification results: The linear-elastic force–displacement results agree well between JAX-FEM and FEniCSx.The comparison uses quasi-static incremental loading from 0 to 0.1 mm over 10 steps.
- Verification results: The hyperelastic force–displacement results agree well between JAX-FEM and FEniCSx under quasi-static loading from 0 to 2 mm.The model is a neo-Hookean solid representing an almost incompressible isotropic elastomer.
- Verification results: The path-dependent elasto-plastic stress–displacement results match exactly between JAX-FEM and FEniCSx.The loading increases from 0 to 0.1 mm and then unloads to 0; automatic differentiation removes the need to derive the elastoplastic tangent tensor manually.
- Performance and scalability: 523 s for 7,703,841 DOF on JAX-FEM GPU versus 8409 s on Abaqus CPU and 4769 s with Abaqus MPI, corresponding to 16.1× and 9.1× acceleration.The benchmark uses a standard ASTM D638 Type 1 tensile specimen with linear elastic material.
- Performance and scalability: Abaqus MPI speedup decreases from 9.1× at 2,344,230 DOF to 1.8× at 7,703,841 DOF as message-passing and storage delays increase.The paper attributes the larger-DOF slowdown to transient variables exceeding CPU memory and being stored on local storage.
3 Solving inverse problems
JAX-FEM formulates inverse problems as PDE-constrained optimization and uses automatic differentiation to compute the derivatives needed for gradient-based optimization. Demonstrations cover sparse full-field reconstruction, topology optimization, and gradient verification.
- Formulation: PDE-constrained optimization treats the finite element solution U as constrained by the discretized governing PDE C(U, θ) = 0.The objective J(U, θ) is optimized over the solution vector U and parameter vector θ.
- Automatic differentiation: JAX-FEM computes derivative information automatically, avoiding manually derived sensitivity expressions for inverse problems.Automatic differentiation is applied to derivative-related computations, including vector-Jacobian products.
- Full-field reconstruction: In sparse full-field reconstruction, the objective drops to nearly zero within 20 optimization steps, while 250 observed points yield about 12.0% relative L2 error.The 250 observations comprise less than 1% of the total points.
- Full-field reconstruction: Increasing observations from 250 to 2500 reduces the full-field inference error from about 12.0% to about 1.4%.The paper reports this change as an example of improved inference with more observations.
- Gradient verification: The Taylor test confirms first-order convergence for the zeroth-order residual expansion and second-order convergence for the first-order expansion.The reported residual scalings are rzeroth ∝ h and rfirst ∝ h^2.
- Topology optimization: Topology optimization examples use automatically computed sensitivities to optimize thin-plate and bracket structures under material constraints.The thin-plate example uses MMA with a constraint allowing only 50% of the material; the bracket uses 45% of its design region.
4 Integration with machine learning
JAX-FEM integrates RVE simulations, neural-network training, and macroscopic finite-element analysis in one data-driven multi-scale workflow. The trained surrogate reproduces DNS mechanical responses well for the composite tensile example.
- Unified workflow: JAX-FEM performs RVE data generation, constitutive-model training, and macroscopic FE deployment within one workflow.Conventional workflows transfer data among FEM and machine-learning tools, whereas JAX-FEM performs all three steps internally.
- RVE computation: RVE simulations impose macroscopic deformation gradients with periodic boundary conditions and generate a database of macroscopic tensor–energy pairs.The constitutive relation W(C) is approximated from repeated RVE computations using different C values.
- Model training: Three MLP architectures are evaluated on an 8:1:1 train-validation-test split, and MLP2 achieves the best validation SMSE.MLP2 has 8 hidden layers with width 64 and is selected for deployment.
- Macroscopic deployment: The deployed MLP2 represents the homogenized constitutive relationship for a 10 × 2 × 10 mm3 sample containing 10 × 2 × 10 RVEs.The macroscopic problem applies quasi-static uniaxial tension up to 1 mm, corresponding to 10% of the sample’s y-axis size.
- Validation: Good agreement between DNS and NN results is observed for total elastic energy and tensile force during loading up to 10% relative displacement.The comparison includes DNS with 200,000 FEM cells and reference responses for bulk hard and soft materials.
5 Conclusions and future work
The paper presents JAX-FEM as an open-source FEM library for forward and inverse problems and data-driven computational mechanics. Future work targets broader FEM capabilities, stronger linear solvers, uncertainty-aware inverse problems, and multi-GPU scalability.
- Conclusion: JAX-FEM is shared as an open-source library for efficiently solving forward and inverse problems and facilitating data-driven computational mechanics.The conclusion presents the software as a foundational tool for these research workflows.
- Future work: The basic FEM toolkit needs richer element types and support for triangular and tetrahedral meshes.This is identified as a future improvement to the current toolkit.
- Future work: The current linear solver uses stabilized biconjugate gradients with the simplest Jacobi preconditioner, while PETSc interfacing is in progress.The paper states that a better solver and preconditioner can improve JAX-FEM performance.
- Future work: The inverse problems considered are deterministic; Bayesian inverse problems and uncertainty quantification are planned for future work.The stated scope boundary concerns the treatment of uncertainty in inverse problems.
- Future work: A single 48 GB NVIDIA GPU currently supports problems of around 10 million DOF before memory becomes insufficient.A multi-GPU version is identified as a future research goal.