An explicit and differentiable Wilson-Daubechies-Meyer transform for gravitational-wave data analysis

Author(s)

Vajpeyi, Avi, Mentasti, Giorgio, Baghi, Quentin, Burke, Ollie, Speri, Lorenzo

Abstract

The Wilson-Daubechies-Meyer (WDM) time-frequency transform has been widely used in gravitational-wave astronomy, yet a self-contained, mathematically explicit reference for practitioners remains lacking. This is especially true for those wishing to adopt the transform in modern Python and JAX inference workflows. We present wdm_transform, an open-source Python package implementing the WDM wavelet-packet time-frequency transform, and document its mathematical foundations, statistical properties, and practical implementation for gravitational-wave data analysis. The package supplies NumPy and JAX backends, both transforms (forward and inverse) validated to floating-point precision, with the JAX backend enabling GPU-accelerated transforms of million-point data streams in tens of milliseconds. As a worked example, we verify that the WDM-domain likelihood reproduces frequency-domain posteriors for a resolved LISA galactic binary under a shared stationary noise model, confirming numerical equivalence of the two representations in that controlled setting. This work paves the way for systematic optimisation of WDM tilings, a particularly promising direction for the non-stationary noise, stochastic backgrounds, and data gaps anticipated in future detectors, and for direct comparisons with alternative time-frequency representations needed to meet the challenges of future gravitational-wave data analysis.

Figures

Caption

Schematic illustration of the WDM forward transform. (a)~Time-domain atoms $g_{nm}(t)$ for two representative WDM cells, labelled $\alpha$ and $\beta$. The higher-frequency cell $\beta$ has a more rapidly oscillating atom. (b)~Frequency-domain magnitudes $|\tilde{g}_{nm}(f)|$ for the same two cells, showing their localization around different frequency channels. (c)~The packed coefficient array $W_{nm}$, with time bins $n$ along the horizontal axis and frequency channels $m$ along the vertical axis. The highlighted cells correspond to the atoms shown in panels (a) and (b). Each interior cell stores $w_{nm} = \sqrt{2}\,(-1)^{nm}\,\mathrm{Re}[C^*_{nm}\,\tilde x_m[n]]$ (Eq.~\protect\eqref{eq:wnm_compact}); the phase factor $C_{nm}$ alternates between $1$ and $i$ on a checkerboard pattern, while the DC and Nyquist edge channels are stored separately.

Caption

\textbf{Orthogonality structure and empirical decorrelation of the sampled WDM transform.} Both panels use the small tiling $N_t = N_f = 8$, $N = N_t N_f = 64$, with packed atoms flattened by $\alpha = n(N_f + 1) + m$. Panel (a) lives in the $N_t(N_f+1) = 72$-dimensional packed coefficient space, while panel (b) lives in the $N$-dimensional signal space. The asymmetry $\boldsymbol{G}\boldsymbol{G}^\dagger \in \mathbb{C}^{72\times 72}$ versus $\boldsymbol{G}^\dagger \boldsymbol{G} \in \mathbb{C}^{64\times 64}$ is intrinsic to the rectangular embedding $G : \mathbb{C}^N \to \mathbb{C}^{N_t(N_f+1)}$ and is the matrix-level expression of the fact that $\boldsymbol{G}^\dagger \boldsymbol{G} = \mathbb{I}_N \neq \boldsymbol{G}\boldsymbol{G}^\dagger$. Faint grid lines in the panels mark the block boundaries $\alpha = n(N_f+1)$ for $n = 0, 1, \ldots, N_t$, separating successive time slices and making the $\pm N_t/2$ offset structure easy to read off by eye. \textbf{(a)} The packed-space Gram matrix $(\boldsymbol{G}\boldsymbol{G}^\dagger)_{(n,m),(p,q)} = \sum_\ell \tilde g_{nm}[\ell]\,\tilde g^*_{pq}[\ell]$. Interior channels ($m = 1,\ldots,N_f-1$) form an identity block, while the DC and Nyquist edge channels ($m \in \{0, N_f\}$) carry the off-diagonal entries at $p = n \pm N_t/2$ predicted by the first case of Eq.~\eqref{eq:orthonormality_like_cases}. These off-diagonals span the $N_t$-dimensional redundant subspace by which $GG^\dagger$ fails to be the identity. \textbf{(b)} The signal-space product $(\boldsymbol{G}^\dagger \boldsymbol{G})_{\ell\ell'} = \sum_{n,m} \tilde g_{nm}[\ell]\,\tilde g^*_{nm}[\ell']$, indexed by sample indices $(\ell,\ell')$ rather than the packed pair, equals $\mathbb{I}_N$ to machine precision ($\max|\boldsymbol{G}^\dagger \boldsymbol{G} - \mathbb{I}| \sim 10^{-15}$), demonstrating that the forward-then-inverse round trip is exact. The exact $\pm 1$ entries at the same $\pm N_t/2$ offsets seen in panel (a) are the empirical signature of the edge-channel redundancy: $w_{n,0}$ and $w_{n + N_t/2,\,0}$ (and likewise for $m = N_f$) are linearly dependent copies of the same underlying real degree of freedom, so their Pearson correlation is identically $\pm 1$ for any input.

Caption

\textbf{Architecture of the \package\ package.} The public API exposes the \texttt{TimeSeries}, \texttt{FrequencySeries}, and \texttt{WDM} datatypes with reversible conversions between them. The forward and inverse WDM transforms, together with the shared window machinery (\texttt{windows.py}), are dispatched to interchangeable NumPy, CuPy, and JAX backends, so the same high-level code runs on CPU, GPU, or TPU.

Caption

\textbf{Forward WDM runtime versus input length.} All transforms use a fixed tiling with $N_t=1024$ time bins and $N_f=N/N_t$ frequency channels. \emph{Top}: median single-transform wall-clock runtime over 7 runs after one warmup call. The thin solid black line shows the $N\log_2 N$ reference scaling. \emph{Bottom}: batched-execution speedup (WDM transform only) $t_\mathrm{serial}/t_\mathrm{batch}$ for a batch size of $B=3$, where the serial baseline applies the kernel independently to each of the three channels and the batched call processes all three simultaneously. Reported timings exclude JAX JIT compilation and host--device data transfer.

Caption

\textbf{Injected resolved galactic binary in both frequency and WDM domains.} (a): channel-A frequency-domain data, the instrumental noise PSD, and the injected source, with the analysis band shaded. Bottom: zooms of the source band in the frequency domain (left (b)) and in the whitened WDM time-frequency plane (right (c)), where amplitude is shown in units of the per-pixel noise standard deviation.

Caption

\textbf{Frequency-domain versus WDM-domain posterior.} Overlaid corner plot for the two analyses of the same realization, with frequency in blue and WDM in orange, the injected truth marked as black solid line, and the prior shown as dotted line in the posterior marginals. The two posteriors overlap closely in $(f_0,\dot f,A,\phi_0)$, and both are far narrower than the prior.

Caption

\textbf{Population PP plot.} Empirical distribution of the injected truth's posterior quantile across the 100 seeds, for each parameter, in the WDM domain (solid, lower opacity) and frequency domain (dashed). Grey shading shows the nested $1/2/3\sigma$ pointwise confidence bands expected under the uniform null. The two domains track each other closely and every parameter remains within the $3\sigma$ band. The smallest per-parameter Kolmogorov--Smirnov $p$-value across both domains is $\approx0.25$ (for $\phi_0$), consistent with uniformity.