JAX-FEM-ANISO: Differentiable GPU-Accelerated Finite Element Framework for Inverse Identification of Finite-Strain Anisotropic Plasticity

Abstract

We present a fully differentiable, GPU-accelerated finite element framework JAX-FEM-ANISO, for forward simulation and inverse parameter identification of finite-strain anisotropic plasticity. Built on JAX-FEM, the framework exploits modern accelerator architectures by parallelizing the three major computational bottlenecks in nonlinear FEM: elemental weak-form and tangent-stiffness evaluation, global sparse matrix assembly, and sparse linear solution. For a large-scale forward problem with 3 million degrees of freedom, JAX-FEM-ANISO on a single NVIDIA H100 GPU achieves up to 9.4$\times$ speed-up over a 24-core CPU Abaqus baseline. Automatic differentiation is applied through the constitutive update and solver workflow, providing consistent Jacobians for complex constitutive models without manual derivation and accurate gradients for PDE-constrained inverse analysis. Compared with finite differences, the JAX-AD gradients avoid step-size sensitivity and provide the required sensitivities at substantially lower computational cost. For inverse characterization, we combine information-rich, topology-optimized heterogeneous specimens with full-field displacement data to identify advanced constitutive model parameters from a single test, replacing what would otherwise require many conventional experiments. We demonstrate accurate recovery of anisotropic yield and hardening parameters in progressively challenging settings, including uniform and spatially varying material properties. The resulting AD-based formulation enables efficient optimization in high-dimensional parameter spaces where finite-difference approaches are computationally infeasible. These results establish differentiable, GPU-accelerated FEM as a practical high-throughput engine for simulation, characterization, and optimization workflows in advanced manufacturing.

Topics & keywords