JZ-FMM: GPU-native differentiable N-body simulations with the Fast Multipole Method

Abstract

N-body simulations are an essential tool for modelling the evolution of gravitating systems. Differentiable N-body simulations allow solving complicated reconstruction problems by connecting observations of evolved systems to simple informative priors on their initial conditions. While so far most reconstruction efforts focus on mildly non-linear large scale dynamics, the internal dynamics of galaxies and clusters should exhibit rich information about their formation. The exploration of reconstruction possibilities in this highly non-linear regime naturally requires highly performant differentiable N-body simulations with small scale forces. Here, we present JZ-FMM, a GPU-native differentiable implementation of the fast multipole method for N-body simulations. We show that the GPU oriented design of tree structure, dual tree traversal and multipole translation operators improves performance over established CPU-native and hybrid N-body codes by notably more than an order of magnitude. Further, we show that gradients can be evaluated efficiently and accurately so that JZ-FMM can be used as an essential building block in future field-level reconstruction efforts. We apply the new code to the reconstruction problem of a tidally stripped satellite to show that deeply non-linear problems can indeed be solved efficiently with simulation gradients, although appropriate care must be taken to navigate the complicated optimization landscape.