JAX-SSO: Differentiable Finite Element Analysis Solver for Structural Optimization and Seamless Integration with Neural Networks
<|reference_start|>JAX-SSO: Differentiable Finite Element Analysis Solver for Structural Optimization and Seamless Integration with Neural Networks: 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.<|reference_end|>
@article{wu2024jax-sso:,
    title={JAX-SSO: Differentiable Finite Element Analysis Solver for Structural
  Optimization and Seamless Integration with Neural Networks},
    author={Gaoyuan Wu},
    journal={arXiv preprint arXiv:2407.20026},
    year={2024},
    archivePrefix={arXiv},
    eprint={2407.20026},
    primaryClass={cs.MS math.OC}
}