Argus: JAX state-space filtering for gravitational wave detection with a pulsar timing array

TL;DR

Argus is a high-performance Python package for detecting and characterising nanohertz gravitational waves in pulsar timing array data that leverages JAX for just-in-time compilation, GPU acceleration, and automatic differentiation, facilitating rapid Bayesian inference with gradient-based samplers.

Abstract

Argus is a high-performance Python package for detecting and characterising nanohertz gravitational waves in pulsar timing array data. The package provides a complete Bayesian inference framework based on state-space models, using Kalman filtering for efficient likelihood evaluation. Argus leverages JAX for just-in-time compilation, GPU acceleration, and automatic differentiation, facilitating rapid Bayesian inference with gradient-based samplers. The state-space approach provides a computationally efficient alternative to traditional frequency-domain methods, offering linear scaling with the number of pulse times-of-arrival, and natural handling of non-stationary processes.