JAX-BEM: Gradient-Based Acoustic Shape Optimisation via a Differentiable Boundary Element Method

Abstract

Engineering structures are increasingly designed using numerical optimisation. However, traditional optimisation methods can be challenging with multiple objectives and many parameters. In machine learning, stable training of artificial neural networks with millions or billions of parameters is achieved using automatic differentiation frameworks such as JAX and Pytorch. Because these frameworks provide accelerated numerical linear algebra with automatic gradient tracking, they also enable differentiable implementations of numerical methods to be built. This facilitates faster gradient-based optimisation of geometry and materials, as well as solution of inverse problems. We demonstrate JAX-BEM, a differentiable Boundary Element Method (BEM) solver, showing that it matches the error of existing BEM codes for a benchmark problem and enables gradient-based geometry optimisation. Although the demonstrated examples are for acoustic simulations, the concept could be readily extended to electromagnetic waves.