JAX-SSO: Differentiable Finite Element Analysis Solver for Structural Optimization and Seamless Integration with Neural Networks
Gaoyuan Wu
TL;DR
JAX-SSO delivers a differentiable finite element analysis solver built on Google's JAX to enable gradient-based structural optimization and seamless physics-informed integration with neural networks. By leveraging automatic differentiation and an adjoint sensitivity framework, it computes design-gradients for problems where $\mathbf{K}(\mathbf{p})\mathbf{u}=\mathbf{f}(\mathbf{p})$ holds, and supports efficient forward/gradient evaluation on CPU and GPU through vectorization and JIT. The work validates accuracy against SAP2000 for 2D and 3D shell/beam cases, and demonstrates performance trade-offs across dense and sparse solvers for large-scale problems, as well as shape, size, topology optimization and NN-assisted design. It highlights the potential of differentiable physics in structural and architectural design, offering a practical tool for integrating FE analysis with neural networks, including physics-informed training. Future directions include expanding element types, addressing nonlinearity, and incorporating dynamic analysis for broader applicability.
Abstract
Differentiable numerical simulations of physical systems have gained rising attention in the past few years with the development of automatic differentiation tools. This paper presents JAX-SSO, a differentiable finite element analysis solver built with JAX, Google's high-performance computing library, to assist efficient structural design in the built environment. With the adjoint method and automatic differentiation feature, JAX-SSO can efficiently evaluate gradients of physical quantities in an automatic way, enabling accurate sensitivity calculation in structural optimization problems. Written in Python and JAX, JAX-SSO is naturally within the machine learning ecosystem so it can be seamlessly integrated with neural networks to train machine learning models with inclusion of physics. Moreover, JAX-SSO supports GPU acceleration to further boost finite element analysis. Several examples are presented to showcase the capabilities and efficiency of JAX-SSO: i) shape optimization of grid-shells and continuous shells; ii) size (thickness) optimization of continuous shells; iii) simultaneous shape and topology optimization of continuous shells; and iv) training of physics-informed neural networks for structural optimization. We believe that JAX-SSO can facilitate research related to differentiable physics and machine learning to further address problems in structural and architectural design.