LogPSplinePSD 0.0.14

LogPsplines in JAX

LogPSplinePSD

Log-spline representation of the power spectral density (PSD) in the frequency domain, using penalized B-splines with a discrete penalty on spline coefficients to prevent overfitting.

• GitHub Repository: https://github.com/nz-gravity/LogPSplinePSD

• Paper: Aimen et al (in preparation)

Overview

LogPSplinePSD implements a Bayesian model for PSD estimation by fitting a log-spline to the periodogram. Main features:

• Log-frequency representation: Works on the log-scale of frequencies for numerical stability and improved resolution.

• P-spline prior: Applies a discrete difference penalty to log B-spline coefficients, enforcing smoothness in the log-PSD domain [Eilers1996].

• Whittle likelihood: Employs Whittle’s approximation for fast likelihood evaluation on periodogram ordinates.

• CPU/GPU sampling: Uses NumPyro (JAX) to perform efficient Hamiltonian Monte Carlo inference.

Methodology

The approach follows the P-spline framework for spectral density estimation described by Maturana-Russel & Meyer (2021) [MaturanaRussel2021].

1. Basis construction Define order-r B-spline basis functions \(B_k(\omega)\), \(k=1,\dots,K+r\), on a grid of interior knots in the log-frequency domain.

2. Penalized prior Apply a discrete \(D\)th-order difference penalty to the spline coefficients \(\{\beta_k\}\), which induces smoothness in the estimated log-PSD.

3. Knot placement (optional) For spectra with sharp features, knot locations can be set based on quantiles of the raw periodogram values to allocate flexibility where needed.

4. Model and likelihood The log-PSD is modeled as:

\begin{equation*} \log f(\lambda_l) = \sum_k \beta_k \, B_k(\log \lambda_l) \end{equation*}

Whittle’s approximation for the periodogram \(I_n(\lambda_l)\) yields the log-likelihood:

\begin{equation*} \log L(\beta) \propto -\sum_{l=1}^{\nu} \left[ \log f(\lambda_l) + \frac{I_n(\lambda_l)}{f(\lambda_l)} \right] \end{equation*}

5. Inference We have two options for inference:

• Use Metropolis-Hastings to sample from the posterior distribution of the spline coefficients.

• Use NumPyro’s NUTS sampler to jointly sample the spline coefficients.

• When init_from_vi is enabled (default), the stochastic VI warm-start runs before NUTS and immediately saves convergence plots and PSD previews to <outdir>/diagnostics so you can inspect them before launching the expensive sampler.

This fixed-basis P-spline approach avoids reversible-jump MCMC over knot numbers and positions, reducing computational cost while retaining flexibility to capture complex spectral features.

Finally, one can also provide a ‘parametric model’ of the PSD as a function that can then be ‘corrected’ non-parametrically by the spline model. This is useful for cases where a known functional form (e.g., power-law) is expected, but additional flexibility is needed to account for deviations in the data.

Installation

pip install LogPSplinePSD

Basic Usage

See demo.py

Author

NZ Gravity

Acknowledgements

Part of the NZ-Gravity and International LISA Consortium efforts on gravitational-wave data analysis.

References

Eilers, P. H. C., & Marx, B. D. (1996). Flexible smoothing with B-splines and penalties. Statistical Science, 11(2), 89–121. DOI:10.1214/ss/1038425655.

Maturana-Russel, J., & Meyer, R. (2021). P-spline spectral density estimation with a discrete penalty. arXiv:1905.01832.