Search arXivSearch

arXiv · 2601.06079

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

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.

Explore related subjects

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Fanglei Hu, Jiachen Guo, Stephen Niezgoda, Wing Kam Liu, Jian Cao. 2025-12-29. Efficient GPU-computing simulation platform JAX-PF for differentiable phase field model. https://arxiv.org/abs/2601.06079

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

A subcell-refined entropy-residual-driven limiting strategy for high-order discontinuous Galerkin methods

Fine-grained, subcell-level dissipation control is essential for achieving robust high-order discontinuous Galerkin (DG) simulations of nonlinear hyperbolic systems in under-resolved regimes while preserving accuracy. This paper proposes a subcell-refined entropy-residual-driven limiting strategy for DG on Legendre-Gauss-Lobatto nodes. The limiter introduces only nearest-neighbor pairwise dissipation within each element, with closed-form coefficients that supply the minimal dissipation required to restore the element entropy inequality. The strategy is a diagonal, locally stable approximation of classical entropy-stable methods, and a generalized subcell framework reveals split-form DG and residual-distribution-based entropy correction schemes as particular choices of the limiting coefficients. For the Euler equations, a physically consistent jump operator separately models thermal and shear entropy production while preserving velocity and pressure equilibrium; a subcell refinement of the Zhang-Shu positivity limiter ensures pointwise positivity. Extensive numerical tests confirm that the scheme maintains optimal high-order accuracy, strictly enforces entropy dissipation, and significantly reduces the difficulty of a posteriori positivity-preserving procedures.

physics.comp-ph

VNS Tokamak for Medical Isotope Production

The Volumetric Neutron Source (VNS) tokamak is a proposed fusion reactor for testing components under fusion neutron irradiation, and has potential use for radioisotope production. The VNS geometry is modeled in the Serpent 2.2.2 and OpenMC 0.15.2 neutronics codes. Coupled neutron-photon simulations compared fluxes, spectra, and selected reaction rates in the blanket and vacuum vessel. Good agreement was found overall, with the largest difference found in (n, 2n) reactions. On an HPC cluster, Serpent 2 was found to have shorter computation time in coupled simulations, while OpenMC was faster in neutron only simulations. Radioisotope production yields were simulated in Serpent 2.2.2 for capsule and Cobalt plate irradiation facilities. Results indicate potential for large volume production of 99Mo, 131I, 225Ac, 177Lu, 192Ir, 64Cu, 67Cu, 161Tb, and 153Sm while 203Pb indicates lower potential. 100Mo and LEU target heating was calculated, suggesting the LEU target mass or the cooling may need adjustment. Optimized 60Co production yielded 1.2 GBq/mg and 100,000 TBq after a 3-year irradiation period. Sensitivity to plant outage for 99Mo, 131I, 177Lu, and 60Co was simulated, suggesting irradiation can be restarted for the same isotope loading and demonstrated long-lived 60Co to be robust to long plant dwell-time.

physics.comp-ph

MadVfold: accelerating NLO event generation and reducing negative weights with SIMD vectorization and GPUs

NLO simulations are essential for LHC physics analyses but are expensive, as they are not only slow but also lead to negative weights, which imply the need to simulate much larger samples of events. Folding is a powerful technique to reduce negative weights but is itself expensive. In this paper I propose ``vectorized folding'' as a new idea to speed up these calculations using SIMD and GPUs, and I present its CUDACPP-based implementation for MG5aMC in MadVfold, including its extension for unfolded NLO event generation. Preliminary results show overall speedups around 6x to 9x with folding and 3x without it. This work is based on a test-centric, LLM-assisted software development process.

physics.comp-ph