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 graphics processing unit–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 ∼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%. Secondly, we show that Fisher forecasting calculations can be sped up by ∼3–5× (on a CPU) with no loss in accuracy. This increased speed makes Fisher forecasting for a population of events 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 on GitHub website and will continue to be updated with new waveforms as they are implemented. https://doi.org/10.1103/PhysRevD.110.064028 Published by the American Physical Society under the terms of the Creative Commons Attribution 4.0 International license. Further distribution of this work must maintain attribution to the author(s) and the published article's title, journal citation, and DOI. Funded by Bibsam. Published by the American Physical Society

Differentiable and accelerated spherical harmonic and Wigner transforms

Differentiable and accelerated spherical harmonic and Wigner transforms

Differentiable and accelerated wavelet transforms on the sphere and ball

Parallel Spherical Harmonic Transforms on heterogeneous architectures (GPUs/multi-core CPUs)

Fast and Accurate Spherical Harmonic Transform for Spatio-Temporal Regular Grid Data

Flexible Multi-Dimensional FFTs for Plane Wave Density Functional Theory Codes

Fast and backward stable transforms between spherical harmonic expansions and bivariate Fourier series

Differentiable and hardware-accelerated waveforms for gravitational wave data analysis

Spherical Harmonics Scaling

Numerical evolutions of fields on the 2-sphere using a spectral method based on spin-weighted spherical harmonics

Spherical framelets from spherical designs

High Performance Computing for Spherical Conformal and Riemann Mappings

Analytic spherical harmonic gradients for real-time rendering with many polygonal area lights

Ameliorating the Courant-Friedrichs-Lewy condition in spherical coordinates: A double FFT filter method for general relativistic MHD in dynamical spacetimes

On the design of scalable, high-precision spherical-radial Fourier features

On the computation of directional scale-discretized wavelet transforms on the sphere

Technologies for supporting high-order geodesic mesh frameworks for computational astrophysics and space sciences

Differential Walk on Spheres

HVOX: Scalable Interferometric Synthesis and Analysis of Spherical Sky Maps

Improved Algorithm for Efficient Computation of Slepian Functions Over Arbitrary Regions on the Sphere

Technical Note: An Efficient Implementation of the Spherical Radon Transform with Cylindrical Apertures