JAX-PF: an efficient GPU-computing simulation for differentiable phase field (PF) simulaiton, built on top of JAX-FEM. Leveraging JAX, this tool addresses the need for co-designing of material and manufacturing processes, supporting the goals of the Materials Genome Initiative, and as a part of the Integrated Computational Materials Engineering.
We want to emphasize the following four features that differential JAX-PF from other PF software:
- Ease-of-use: Leveraging capability of automatic differentiation (AD) in JAX, JAX-PF automatically generates Jacobians and derivatives of different energy terms with machine precision, enabling the realization of multi-physics and multi-variable PF models.
- Automatic Sensitivity: Implicit time integration with customized adjoint-based AD enables efficient gradient-based optimization and inverse design of strongly nonlinear PF systems.
- High-performance GPU-acceleration: Through the XLA backend and vectorized operations, JAX-PF delivers competitive GPU performance, drastically reducing computational time relative to CPU-based solvers.
- Unified multiscale ecosystem with JAX-CPFEM: Built on the same JAX-FEM foundation, JAX-PF integrates seamlessly with JAX-CPFEM to enable coupled process–structure–property simulations (e.g., dynamic recrystallization), while preserving full differentiability for optimization and design.
🔥 Join us for the development of JAX-PF! This project is under active development!
Four benchmark problems are provided, including Allen–Cahn, Cahn–Hilliard, coupled Allen–Cahn and Cahn–Hilliard, each implemented with both explicit and implicit time integration, and Eshelby inclusion for lattice misfit in solid-state phase transformations.
📣 JAX-PF reproduces the benchmark results with good agreement to PRISMS-PF.
Validation of benchmark problems in JAX-PF. While PRISMS-PF implements explicit time integration, JAX-PF provides both explicit and implicit schemes, with results showing agreement between the two frameworks.
Several representative applications are included, ranging from solidification to solid-state transformations: grain growth, static recrystallization, spinodal decomposition, precipitation with the Wheeler–Boettinger–McFadden (WBM) model, and precipitation with the Kim–Kim–Suzuki (KKS) model.
🔥 For each case, both explicit and implicit time stepping schemes are provided
The initial (left) and final (right) grain structure for a 2D grain growth simulation.
The distribution of composition during a simulation of spinodal decomposition from initial fluctuations (left) to final two distinct phases (right).
A 2D simulations of the multi-variants precipitate in an Mg-Nd alloy were simulated using different frameworks, based on Wheeler-Boettinger-McFadden (WBM) model.
A 3D simulations of the single-variants precipitate in an Mg-Nd alloy were simulated using different frameworks, based on Kim-Kim-Suzuki (KKS) model.
📣 JAX-PF and JAX-CPFEM are built on top of the same underlying JAX-FEM ecosystem, the two can be tightly coupled within a single GPU-accelerated differentiable framework. This integration allows simultaneous treatment of process–structure and structure–property relationships, supporting multiscale simulations (PF-CPFEM) under complex internal and/or external applied fields.
Coupled JAX-PF and JAX-CPFEM framework for process–structure–property integration.
📣 A demos: calibration of material parameters using differentiable phase-field simulations.
Calibration of material parameters using differentiable PF simulations for multi-variants in Mg-Nd alloys. Subfigure (a) shows the evolution of the objective function with the number of optimization iterations, based on the synthetic imaging of the microstructure as the ground truth, as shown in the bottom subfigure in Fig. 6. Subfigure (b) shows the comparison between reference and calibrated material parameters.Subfigure (c) shows the comparison between the reference (red) and calibrated (blue) microstructural morphology of the different variants, with precipitate variants identified by concentrations in the range 0.12 to 0.16. Subfigure (d) and (e) shows the Concentration distributions for the reference and calibrated results across the domain.
JAX-PF supports Linux and macOS, which depend on 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.
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.
Place the downloaded phaseField/ file in the applications/ folder of JAX-FEM, and then you can run it.
For example, you can download phaseField/allenCahn/explicit_fem folder and place it in the applications/ folder of JAX-FEM, run
python -m applications.phaseField.allenCahn.explicit_fem.explicit_ACfrom the root directory. Use Paraview for visualization.
📣 Comming soon!
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