Efficient GPU-computing simulation platform JAX-PF for differentiable phase field model

TL;DR

JAX-PF is a GPU-accelerated, differentiable phase field simulation platform that significantly improves performance and enables inverse design through automatic differentiation, facilitating advanced materials research.

Contribution

The paper introduces JAX-PF, a novel open-source software that combines high-performance GPU computing with automatic differentiation for phase field modeling and inverse design.

Findings

Achieves ~5x speedup over existing CPU-based tools.

Supports automatic differentiation for free-energy functionals.

Demonstrates successful calibration of material parameters using sensitivities.

Abstract

We present JAX-PF, an open-source, GPU-accelerated, and differentiable Phase Field (PF) software package, supporting both explicit and implicit time stepping schemes. Leveraging the modern computing architecture JAX, JAX-PF achieves high performance through array programming and GPU acceleration, delivering ~5x speedup over PRISMS-PF with MPI (24 CPU cores) for systems with ~4.19 million degrees of freedom using explicit schemes, and scaling efficiently with implicit schemes for large-size problems. Furthermore, a key feature of JAX-PF is automatic differentiation (AD), eliminating manual derivations of free-energy functionals and Jacobians. Beyond forward simulations, JAX-PF demonstrates its potential in inverse design by providing sensitivities for gradient-based optimization. We demonstrate, for the first time, the calibration of PF material parameters using AD-based sensitivities, highlighting its capability for high-dimensional inverse problems. By combining efficiency, flexibility, and full differentiability, JAX-PF offers a fast, practical, and integrated tool for forward simulation and inverse design, advancing co-designing of material and manufacturing processes and supporting the goals of the Materials Genome Initiative.