Minimization of Nonlinear Energies in Python Using FEM and Automatic Differentiation Tools

TL;DR

This work investigates solving nonlinear energy minimization problems arising from PDEs via the finite element method by porting a MATLAB workflow to Python and exploiting automatic differentiation with $\text{JAX}$, graph coloring, and algebraic multigrid. It demonstrates three benchmarks—the $p$-Laplace equation, the Ginzburg–Landau energy, and a compressible Neo-Hookean hyperelasticity model—where the Python implementation computes $\nabla J$ and $\nabla^2 J$ exactly using AD and JVP, and employs JIT-compiled operators and AMG solvers. The results show substantial speedups (often $>10\times$) for large-scale problems and highlight the flexibility to add new nonlinear functionals with minimal code, contrasting with MATLAB’s limitations in supporting custom linear solvers. Overall, the study argues that Python provides a scalable, adaptable environment for high-fidelity variational simulations in physics and engineering, especially as problem sizes grow and solver sophistication increases.

Abstract

This contribution examines the capabilities of the Python ecosystem to solve nonlinear energy minimization problems, with a particular focus on transitioning from traditional MATLAB methods to Python's advanced computational tools, such as automatic differentiation. We demonstrate Python's streamlined approach to minimizing nonlinear energies by analyzing three problem benchmarks - the p-Laplacian, the Ginzburg-Landau model, and the Neo-Hookean hyperelasticity. This approach merely requires the provision of the energy functional itself, making it a simple and efficient way to solve this category of problems. The results show that the implementation is about ten times faster than the MATLAB implementation for large-scale problems. Our findings highlight Python's efficiency and ease of use in scientific computing, establishing it as a preferable choice for implementing sophisticated mathematical models and accelerating the development of numerical simulations.

Figures (7)

• Figure 1: Discretization of a rectangular domain (left) and the corresponding sparsity pattern (right).
• Figure 1: Performance comparison for the p-Laplace benchmark
• Figure 2: Numerical approximation of the solution $u$ of the p-Laplace benchmark.
• Figure 2: Performance comparison for the G-L benchmark
• Figure 3: Performance comparison for the hyperelasticity benchmark