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

TL;DR

JAX-BEM introduces a differentiable Boundary Element Method leveraging automatic differentiation frameworks like JAX, enabling efficient gradient-based acoustic shape optimization and inverse problem solving.

Contribution

It presents a novel differentiable BEM solver built with JAX, facilitating faster and more flexible gradient-based optimization of geometries and materials.

Findings

JAX-BEM matches existing BEM codes in accuracy for benchmark problems.

It enables effective gradient-based geometry optimization.

The approach can be extended to electromagnetic wave simulations.

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.