A high-level interface for transient multiphysics modelling in JAX

  • Bode, Tobias (Federal Institute for Geosciences and Natural)

Please login to view abstract download link

The implementation of nonlinear, transient multiphysics problems often suffers from long development cycles due to the manual derivation of residuals and Jacobians. In this talk, a new high-level interface for the open-source library AutoPDEx [1] is presented, which allows a rapid prototyping of multiphysical models without significant implementation effort. The framework is built on JAX [2], a code-to-code transformation library that enables automatic differentiation and just-in-time compilation, including support for GPU hardware. By leveraging these features, AutoPDEx allows users to define the physics of a problem directly through its weak form as a JAX-traceable function. For spatial discretization, finite-element-based approaches are employed. For temporal discretization, the user can choose, e.g., from Newmark, Adams-Moulton or Backward Differentiation Formula schemes. The functionality of the interface is demonstrated through academic benchmarks and geomechanical applications.