pespace: A new tool of GPU-accelerated and auto-differentiable response generation and likelihood evaluation for space-borne gravitational wave detectors

TL;DR

Pespace presents a GPU-accelerated, auto-differentiable framework for space-borne GW data analysis, enabling full Bayesian parameter estimation of MBHB signals across LISA, Taiji, and Tianqin. By reimplementing PhenomXAS and PhenomXHM in the tiwave module and wrapping waveform generation within taichi-lang, the approach achieves efficient likelihood evaluations and exact Fisher matrix calculations via automatic differentiation. Results show that higher-mode content and detector networks substantially improve parameter measurements, and AD provides accurate, machine-precision derivatives for forecasting and initialization. This infrastructure lowers barriers to high-performance GW data analysis and sets the stage for future global-fitting and more realistic noise modeling in space-borne missions.

Abstract

Space-borne gravitational wave detectors will expand the scope of gravitational wave astronomy to the milli-Hertz band in the near future. The development of data analysis software infrastructure at the current stage is crucial. In this paper, we introduce \texttt{pespace} which can be used for the full Bayesian parameter estimation of massive black hole binaries with detectors including LISA, Taiji, and Tianqin. The core computations are implemented using the high-performance parallel programming framework \texttt{taichi-lang} which enables automatic differentiation and hardware acceleration across different architectures. We also reimplement the waveform models \texttt{PhenomXAS} and \texttt{PhenomXHM} in the separate package \texttt{tiwave} to integrate waveform generation within the \texttt{taichi-lang} scope, making the entire computation accelerated and differentiable. To demonstrate the functionality of the tool, we use a typical signal from a massive black hole binary to perform the full Bayesian parameter estimation with the complete likelihood function for three scenarios: including a single detector using the waveform with only the dominant mode; a single detector using the waveform including higher modes; and a detector network with higher modes included. The results demonstrate that higher modes are essential in breaking degeneracies, and coincident observations by the detector network can significantly improve the measurement of source properties. Additionally, automatic differentiation provides an accurate way to obtain the Fisher matrix without manual fine-tuning of the finite difference step size. Using a subset of extrinsic parameters, we show that the approximated posteriors obtained by the Fisher matrix agree well with those derived from Bayesian parameter estimation.

Figures (10)

• Figure 1: Comparison of waveforms PhenomXHM generated by lalsimulation and tiwave for an example signal. The left panels show the amplitude of waveforms, and the right panels show the real part of waveforms. The blue solid lines and orange dashed lines are used to represent waveforms from tiwave and lalsimulation, respectively.
• Figure 2: Comparison of numerical errors of derivatives obtained by automatic differentiation (AD) and numerical differentiation (ND). We use the derivative of the optimal SNR $\rho$, where $\rho^2\equiv\langle s|s\rangle$, with respect to the inclination $\iota$, which can be computed analytically, to illustrate the accuracy of derivative computation. The vertical axis shows the relative difference between symbolic derivatives (SD) and derivatives obtained from automatic or numerical differentiation, which is given by $\mathrm{abs}\left[ (\frac{\partial \rho}{\partial \iota}|_{\rm SD} - \frac{\partial \rho}{\partial \iota}|_{\rm AD, or ND})/{\frac{\partial \rho}{\partial \iota}|_{\rm SD}}\right]$. The results of numerical differentiation are obtained by the central finite difference scheme with different steps indicated by dashed lines in different colors. As discussed in Sec. \ref{['subsec_implementation']}, too small (or large) steps can induce significant round-off (or truncation) errors. By contrast, automatic differentiation can avoid manually tuning step sizes, and offer accurate computation of derivatives, as shown by the gray line in the above figure.
• Figure 3: Computational cost of waveform generation and likelihood evaluation. The blue and orange lines denote results of tests performed using CPU only and using GPU acceleration, respectively, and the solid and dashed lines are used to represent the different settings of floating-point precision. The vertical dotted lines denote the corresponding time domain duration with the sampling interval of 10 seconds. The number of frequency samples is determined by the length of time samples through numpy.fft.rfftfreq with the truncation of the range $[10^{-4}, 5\times10^{-2}]$. The computational times are obtained by averaging over 10 independent runs with parameters randomly sampled in the parameter space. These tests are performed on a computing platform equipped with a CPU of AMD Threadripper 7955WX and a GPU of NVIDIA RTX 5000 Ada Generation. The CPU is forced to use only one core in tests. Computational cost of waveform generation using lalsimulation is also shown in the upper panel, where we use the interface SimInspiralChooseFDWaveformSequence to get waveforms on specific frequency grids, all waveform flags use default settings except that the multibanding is switched off, and OpenMP parallelization is disabled. The likelihood evaluation incorporates one LISA-like detector with three orthogonal TDI channels.
• Figure 4: Posteriors obtained from the full Bayesian parameter estimation. The blue, orange, and green colors denote results of three scenarios: the single LISA detector observation and using the waveform with only the dominant 22 mode; the single LISA detector observation and using the waveform with higher modes (HM) including 21, 33, 32, 44 modes; the joint observation by the LISA-Taiji-Tianqin detector network and using the waveform with higher modes. The dashed lines in marginalized posterior histograms mark the $90\%$ credible interval, defined by the 5th and 95th percentiles.
• Figure 5: Heatmap of the covariance matrix obtained from the Fisher matrix computed via automatic differentiation. Warm colors indicate positive correlations and cool colors indicate negative correlations for the off-diagonal elements. Since a logarithmic color scale is used, the sign convention in the colorbar appears inverted.