ripple: Differentiable and Hardware-Accelerated Waveforms for Gravitational Wave Data Analysis

TL;DR

This paper introduces ripple, a framework using JAX for differentiable, GPU-accelerated gravitational waveforms, enabling faster analysis, improved calibration, and more efficient statistical inference in GW science.

Contribution

The paper presents ripple, a novel JAX-based framework for differentiable gravitational waveforms that accelerates analysis tasks and enhances waveform calibration and inference methods.

Findings

Waveforms in ripple achieve similar speed to C implementation.

GPU acceleration yields over 10x speedup in waveform evaluation.

Differentiable waveforms improve calibration and inference efficiency.

Abstract

We propose the use of automatic differentiation through the programming framework jax for accelerating a variety of analysis tasks throughout gravitational wave (GW) science. Firstly, we demonstrate that complete waveforms which cover the inspiral, merger, and ringdown of binary black holes (i.e. IMRPhenomD) can be written in jax and demonstrate that the serial evaluation speed of the waveform (and its derivative) is similar to the lalsuite implementation in C. Moreover, jax allows for GPU-accelerated waveform calls which can be over an order of magnitude faster than serial evaluation on a CPU. We then focus on three applications where efficient and differentiable waveforms are essential. Firstly, we demonstrate how gradient descent can be used to optimize the $\sim 200$ coefficients that are used to calibrate the waveform model. In particular, we demonstrate that the typical match with numerical relativity waveforms can be improved by more than 50% without any additional overhead. Secondly, we show that Fisher forecasting calculations can be sped up by $\sim 100\times$ (on a CPU) with no loss in accuracy. This increased speed makes population forecasting substantially simpler. Finally, we show that gradient-based samplers like Hamiltonian Monte Carlo lead to significantly reduced autocorrelation values when compared to traditional Monte Carlo methods. Since differentiable waveforms have substantial advantages for a variety of tasks throughout GW science, we propose that waveform developers use jax to build new waveforms moving forward. Our waveform code, ripple, can be found at https://github.com/tedwards2412/ripple, and will continue to be updated with new waveforms as they are implemented.