Efficient GPU-computing simulation platform JAX-CPFEM for differentiable crystal plasticity finite element method

We present the formulation and applications of JAX-CPFEM, an open-source, GPU-accelerated, and differentiable 3-D crystal plasticity finite element method (CPFEM) software package. Leveraging the modern computing architecture JAX, JAX-CPFEM features high performance through array programming and GPU acceleration, achieving a 39× speedup in a polycrystal case with ~52,000 degrees of freedom compared to MOOSE with MPI (8 cores). Furthermore, JAX-CPFEM utilizes the automatic differentiation technique, enabling users to handle complex, non-linear constitutive materials laws without manually deriving the case-specific Jacobian matrix. Beyond solving forward problems, JAX-CPFEM demonstrates its potential in an inverse design pipeline, where initial crystallographic orientations of polycrystal copper are optimized to achieve targeted mechanical properties under deformations. The end-to-end differentiability of JAX-CPFEM allows automatic sensitivity calculations and high-dimensional inverse design using gradient-based optimization. The concept of differentiable JAX-CPFEM provides an affordable, flexible, and multi-purpose tool, advancing efficient and accessible computational tools for inverse design in smart manufacturing.