dorsal/arxiv
View SchemaEfficient GPU-computing simulation platform JAX-PF for differentiable phase field model
| Authors | Fanglei Hu, Jiachen Guo, Stephen Niezgoda, Wing Kam Liu, Jian Cao |
|---|---|
| Categories | |
| ArXiv ID | 2601.06079vv1 |
| URL | https://arxiv.org/abs/2601.06079 |
| License | http://creativecommons.org/licenses/by-nc-sa/4.0/ |
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.
{
"annotation_id": "03aa5747-0122-4f9d-ae92-c07842b58420",
"date_created": "2026-02-17T05:53:05.001000Z",
"date_modified": "2026-02-17T05:53:05.001000Z",
"file_hash": "2547c539ab24f4543136e64d4835f68722c87e2216bb000fef26bac5c97305a7",
"private": false,
"record": {
"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.",
"arxiv_id": "2601.06079",
"authors": [
"Fanglei Hu",
"Jiachen Guo",
"Stephen Niezgoda",
"Wing Kam Liu",
"Jian Cao"
],
"categories": [
"physics.comp-ph",
"cs.CE"
],
"license": "http://creativecommons.org/licenses/by-nc-sa/4.0/",
"title": "Efficient GPU-computing simulation platform JAX-PF for differentiable phase field model",
"url": "https://arxiv.org/abs/2601.06079",
"version": "v1"
},
"schema_id": "dorsal/arxiv",
"source": {
"execution_id": "57b93689-729c-41ea-8db3-162ca3bbfe76",
"id": "arXiv Dataset",
"type": "Model",
"variant": "snapshot-2026-01-17",
"version": "0.1.0"
},
"user_id": 1000002
}