JAX-CPFEM: an efficient GPU-computing simulation for differentiable crystal plasticity finite element method (CPFEM), built on top of JAX-FEM. Leveraging JAX, this tool addresses the need for inverse design for smart manufacturing.
:fire: JAX-PF, an efficient GPU-computing simulation for differentiable phase field models was released, with the ecosystem as JAX-CPFEM for doing multiscale simulation
The concept of differentiable CPFEM
We want to emphasize the following three features that differential JAX-CPFEM from other CPFEM software:
- GPU-acceleration: Efficient solution to forward CPFEM (involving complicated nonlinear relations) with GPU acceleration based on array programming style and matrix formulation.
- Automatic Constitutive Laws: Free to realize different deformation mechanics represented by constitutive materials laws by evaluating the case-by-case Jacobian matrix using automatic differentiation.
- Automatic Sensitivity: Differential simulation for sensitivity analysis with respect to design parameters used in single crystal or polycrystal, which can be seamlessly integrated with inverse design.
:fire: Join us for the development of JAX-CPFEM! This project is under active development!
Applications
Forward simulations
An efficient solution to forward CPFEM with Kalidindi’s hardening law.
:mega: Join us in considering more constitutive laws! No need for deriving the Jacobian matrix by hand!
Crystal plasticity: 304 steel grain structure (left) and von Mises stress (right).
Gradient-based calibration of CPFEM model parameters with AD‐based sensitivities
The calibration was performed on a RVE constructured based on EBSD measurement.
Inverse Design via AD-based sensitivities
Based on the pipeline combining differentiable JAX-CPFEM with gradient-based optimization, here is a demo of the initial microstructure design of a polycrystal metal featuring the targeted mechanical property after applied deformations.
:mega: Join us in structure-processing-microstructure design!
Inverse design of the crystal orientation of polycrystalline copper, involving three sequential rotations of Euler angles around the Z, Y, and X axes relative to their initial position.
Installation
JAX-CPFEM supports Linux and macOS, which depend on JAX-FEM.
Install JAX-FEM
JAX-FEM is a collection of several numerical tools, including the Finite Element Method (FEM). See JAX-FEM installation instructions. Depending on your hardware, you may install the CPU or GPU version of JAX. Both will work, while the GPU version usually gives better performance.
Install Neper
Neper is a free/open-source software package for polycrystal generation and meshing. It can be used to generate polycrystals with a wide variety of morphological properties. A good instruction video is on Youtube.
Install JAX-CPFEM
Place the downloaded file in the applications/ folder of JAX-FEM, and then you can run it.
Quick Tests
For example, you can download polycrystal_304steel/ folder and place it in the applications/ folder of JAX-FEM, run
python -m applications.polycrystal_304steel.polycrystal_304steel
from the root directory. This is a CPFEM example of polycrystal 304 stainless steel under tensile loading. Use Paraview for visualization.
Polycrystal 304 steel (FCC) under tensile loading. A comparison of simulation results between JAX-CPFEM (point) and MOOSE (line).
Tutorial
We shared the introduction of JAX-CPFEM at npj Computational Materials, which contains an inverse microstructure design case.
Citations
If you found this library useful in academic or industry work, we appreciate your support if you consider 1) starring the project on Github, and 2) citing relevant papers:
- Efficient GPU-computing simulation platform JAX-CPFEM for differentiable crystal plasticity finite element method.
DOI: https://doi.org/10.1038/s41524-025-01528-2
- An efficient GPU-accelerated calibration of crystal plasticity model parameters by multi-objective gradient-based optimization with AD‐based sensitivities.
DOI: https://doi.org/10.1115/1.4070536
评论