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.

翻译：本文介绍JAX-PF——一个开源、GPU加速且可微分的相场（PF）软件包，支持显式和隐式时间步进方案。该平台基于现代计算架构JAX，通过数组编程和GPU加速实现高性能：在约419万自由度的系统中，采用显式方案时相比PRISMS-PF配合MPI（24个CPU核心）可获得约5倍加速；对于大规模问题，其隐式方案亦具备高效扩展能力。此外，JAX-PF的核心特性在于自动微分（AD）功能，无需手动推导自由能泛函和雅可比矩阵。除正向模拟外，JAX-PF通过提供基于梯度的优化灵敏度，展示了在逆向设计方面的潜力。我们首次实现了基于AD灵敏度的PF材料参数标定，彰显了其处理高维逆问题的能力。通过融合高效性、灵活性与完全可微性，JAX-PF为正向模拟和逆向设计提供了快速、实用的一体化工具，推动材料与制造工艺的协同设计，支撑材料基因组计划目标的实现。