A differentiable finite element analysis (FEA) solver for structural optimization, enabled by JAX.
JaxSSO is a differentiable finite element analysis (FEA) solver for structural optimization, enabled by JAX.
Features
Automatic differentiation (AD): an easy and accurate way for gradient evaluation. The implementation of AD avoids deriving derivatives manually or truncation errors from numerical differentiation.
Accelerated linear algebra (XLA) and just-in-time compilation: these features in JAX boost the gradient evaluation.
Hardware acceleration: run on GPUs and TPUs for a faster experience.
Support for beam-column elements and MITC-4 quadrilateral shell elements.
Shape optimization, size optimization, and topology optimization.
Seamless integration with machine learning (ML) libraries.
Here is an implementation of JaxSSO to form-find a structure inspired by Mannheim Multihalle using simple gradient descent. (First photo credit to Daniel Lukac.)
References
2024
JAX-SSO: Differentiable Finite Element Analysis Solver for Structural Optimization and Seamless Integration with Neural Networks
Shape optimization is of great significance in structural engineering, as an efficient geometry leads to better performance of structures. However, the application of gradient-based shape optimization for structural and architectural design is limited, which is partly due to the difficulty and the complexity in gradient evaluation. In this work, an efficient framework based on automatic differentiation (AD), the adjoint method and accelerated linear algebra (XLA) is proposed to promote the implementation of gradient-based shape optimization. The framework is realized by the implementation of the high-performance computing (HPC) library JAX. We leverage AD for gradient evaluation in the sensitivity analysis stage. Compared to numerical differentiation, AD is more accurate; compared to analytical and symbolic differentiation, AD is more efficient and easier to apply. In addition, the adjoint method is used to reduce the complexity of computation of the sensitivity. The XLA feature is exploited by an efficient programming architecture that we proposed, which can boost gradient evaluation. The proposed framework also supports hardware acceleration such as GPUs. The framework is applied to the form finding of arches and different free-form gridshells: gridshell inspired by Mannheim Multihalle, four-point supported gridshell, and canopy-like structures. Two geometric descriptive methods are used: non-parametric and parametric description via Bézier surface. Non-constrained and constrained shape optimization problems are considered, where the former is solved by gradient descent and the latter is solved by sequential quadratic programming (SQP). Through these examples, the proposed framework is shown to be able to provide structural engineers with a more efficient tool for shape optimization, enabling better design for the built environment.
