WDM.code.discrete_wavelet_transform package

Submodules

WDM.code.discrete_wavelet_transform.WDM module

class WDM.code.discrete_wavelet_transform.WDM.WDM_transform(dt: float, Nf: int, N: int, q: int = 16, d: int = 4, A_frac: float = 0.25, calc_m0: bool = False)

Bases: object

This class implements the WDM discrete wavelet transform.

dt

The cadence, or time step, of the time series, \(\delta t\). Equal to inverse of the sampling frequency.

Type:

float

Nf

Number of wavelet frequency bands, \(N_f\). Must be even. This controls the time/frequency resolution of the wavelets.

Type:

int

N

Length of the input time series, \(N\). Must be an even multiple of \(N_f\).

Type:

int

Nt

Number of wavelet time bands, \(N_t\). Equal to \(N/N_f\). Must be even.

Type:

int

q

Truncation parameter, \(q\). Formally the time domain wavelets have infinite extent but in practice are truncated at \(\pm q \Delta T\). This must be an integer in the range \(1 \leq q \leq N_t/2\).

Type:

int

d

Steepness parameter for the Meyer window transition. Must be a positive integer, \(d\geq 1\).

Type:

int

A_frac

Fraction of total bandwidth used for the flat-top response region. Must be in the range [0, 1].

Type:

float

B_frac

Fraction of total bandwidth used for the transition region. This is set based on A_frac so \(2A_{\mathrm{frac}}+B_{\mathrm{frac}}=1\).

Type:

float

A

Half-width of the flat-top response region in angular frequency (radians per unit time), \(A\). Satisfies \(\Delta \Omega = 2A + B\).

Type:

float

B

Width of the transition region in angular frequency (radians per unit time), \(B\). Satisfies \(\Delta \Omega = 2A + B\).

Type:

float

dF

Frequency resolution of the wavelets, or the total wavelet frequency bandwidth \(\Delta F = \frac{\Delta \Omega}{2 \pi}\).

Type:

float

dT

Time resolution of the wavelets. Related to the frequency resolution by \(\Delta F \Delta T = \frac{1}{2}\).

Type:

float

dOmega

Angular Frequency resolution of the wavelets (radians per unit time), or total wavelet angular frequency bandwidth, \(\Delta \Omega = 2A+B\).

Type:

float

T

Total duration of the time series. Related to \(N\) and \(\delta t\) by \(T = N \delta t\).

Type:

float

df

The frequency resolution of the time series, \(\delta f = 1/T\).

Type:

float

f_s

Sampling frequency of the time series, \(f_s = \frac{1}{\delta t}\).

Type:

float

f_Ny

Nyquist frequency (i.e. maximum frequency) of the time series, \(f_{\rm Ny} = \frac{1}{2 \delta t}\).

Type:

float

K

Window length in samples, \(K = 2 q N_f\). By definition, this is always an even integer.

Type:

int

times

The sample times of the time series, \(t_k = k \delta t\) for \(k\in\{0,1,\ldots,N-1\}\). Array shape=(N,).

Type:

jnp.ndarray

freqs

The sample frequencies of the time series, \(f_k = k \delta f\) for \(k\in\{-N/2,N/2+1,\ldots,N/2-1\}\). Array shape=(N,). Note, the zero-frequency component is in the center of the spectrum.

Type:

jnp.ndarray

Cnm

Coefficients \(C_{nm}\) used for the wavelet transform. Equal to 1 if \(n+m\) is even or \(i\) if it’s odd. Array shape=(N_t, N_f).

Type:

jnp.ndarray

calc_m0

If this is set to False (default value) then the wavelet coefficients with \(m=0\) are handled INCORRECTLY. This is faster. If these coefficients are needed the initialise the class with calc_m0=True.

Type:

bool

window_TD

The time-domain Meyer window function, \(\phi(t)\). Array shape=(N,).

Type:

jnp.ndarray

window_FD

The frequency-domain Meyer window function, \(\tilde{\Phi}(f)\). Array shape=(N,), dtype=complex.

Type:

jnp.ndarray

cached_Gnm_basis

The frequency-domain wavelet basis \(\tilde{G}_{nm}(f)\). Array shape=(N, Nt, Nf).

Type:

jnp.ndarray

cached_gnm_basis

The time-domain wavelet basis \(g_{nm}(t)\). Array shape=(N, Nt, Nf).

Type:

jnp.ndarray

jax_dtype

Use jax.config.update(“jax_enable_x64”, True).

Type:

jnp.float64

jax_dtype_int

Use jax.config.update(“jax_enable_x64”, True).

Type:

jnp.int64

Gnm(n: int, m: int, freq: Array = None) → Array

Compute the frequency-domain representation of the wavelets, \(\tilde{G}_{nm}(f)\).

This method computes the frequency-domain wavelet for a single choice of \(n\) and \(m\) using the expressions below. If you instead want to compute the full wavelet basis for all \(n\) and \(m\) efficiently, use the Gnm_basis method.

For \(m>0\), the wavelet is given by

(37)\[\tilde{G}_{nm}(f) = \frac{\exp(-2\pi i n f \Delta T)}{\sqrt{2}} \left( C_{nm}\tilde{\Phi}(f+m\Delta F) + C^*_{nm}\tilde{\Phi}(f-m\Delta F) \right) .\]

For the special case \(m=0\), the wavelet is given by

(38)\[\begin{split}\tilde{G}_{n0}(f) = \begin{cases} \exp(-4\pi i n f \Delta T) \tilde{\Phi}(f) & \mathrm{if}\; n<N_t/2 \\ \frac{1}{2} \exp(-4\pi i n f \Delta T) \left( \tilde{\Phi}(f-f_{\rm Ny}) + \tilde{\Phi}(f+f_{\rm Ny}) \right) & \mathrm{if}\; n\geq N_t/2 \end{cases}\end{split}\]

Parameters:

• n (int) – Wavelet time index.

• m (int) – Wavelet frequency index.

• freq (jnp.ndarray) – Frequencies at which to evaluate the wavelet. If None, then defaults to self.freqs. Optional

Returns:

Gnm – Complex array shaped like freq. The frequency-domain wavelet.

Return type:

jnp.ndarray

Gnm_basis() → Array

Efficient computation of frequency-domain wavelet basis \(\tilde{G}_{nm}(f)\). Instead of calling the functions for \(\tilde{G}_{nm}(f)\) explicilty as is done in the Gnm method, this function shifts indices of window_FD.

The result is cached to speed up subsequent calls.

Returns:

basis – Array of shape (N, Nt, Nf). The time-domain wavelet basis. The first axis is frequency, the second is the wavelet time index, and the third is the wavelet frequency index.

Return type:

jnp.ndarray

build_frequency_domain_window() → Array

Construct the frequency-domain window function \(\tilde{\Phi}(f)\).

Note, the zero-frequency component is in the center of the spectrum.

Returns:

Phi – Array of shape (N,). Complex-valued frequency-domain window.

Return type:

jnp.ndarray

build_time_domain_window() → Array

Construct the time-domain window function \(\phi(t)\).

This method builds the Meyer window in the frequency domain and applies an inverse FFT to obtain the corresponding time-domain window.

Returns:

phi – Array of shape (N,). Real-valued time-domain window.

Return type:

jnp.ndarray

check_indices(n: Array, m: Array) → bool

Check if the wavelet indices are in the valid range.

The n indices must satisfy \(0 \leq n < N_t\) and the m indices must satisfy \(0 \leq m < N_f\).

Parameters:

• n (jnp.ndarray) – Array of n indices, dtype=int. Wavelet time index.

• m (jnp.ndarray) – Array of m indices, dtype=int. Wavelet frequency index.

Returns:

check – True if the all indices are valid, otherwise False.

Return type:

bool

dwt(x: Array) → Array

Forward discrete wavelet transform.

Calls self.fast_forward_transform. Vectorised to allow for transforming multiple time series at once.

Parameters:

x (jnp.ndarray) – Input time series. Array shape=(N,) or (…, N).

Returns:

w – Wavelet coefficients. Array shape=(Nt, Nf) or (…, Nt, Nf).

Return type:

jnp.ndarray

forward_transform_exact(x: Array) → Array

Perform the forward discrete wavelet transform. Transforms the input signal from the time domain into the time-frequency domain.

This method computes the wavelet coefficients using the exact expression

(39)\[w_{nm} = \delta t \sum_{k=0}^{N-1} g_{nm}[k] x[k] ,\]

where the sum is over the whole time-domain signal (no truncation).

This method is slow but exact.

Parameters:

x (jnp.ndarray) – Array shape shape (N,). Input time-domain signal to be transformed.

Returns:

w – Array shape shape (Nt, Nf). WDM time-frequency-domain wavelet coefficients.

Return type:

jnp.ndarray

forward_transform_fft(x: Array) → Array

Perform the forward discrete wavelet transform. Transforms the input signal from the time domain into the time-frequency domain.

For the \(m>0\) terms, the wavelet coefficients are calculated using the following expression,

(40)\[w_{nm} = \frac{\sqrt{2}\delta t}{N} (-1)^{nm} \,\mathrm{Re}\, \Big( C_{nm}^* x_m[n] \Big)\]

where

(41)\[x_m[n] = \sum_{l=-N_t/2}^{N_t/2-1} \exp\left(\frac{2\pi i nl}{N_t} \right) \Phi[l] X[l-mN_t/2] .\]

The \(m=0\) terms, if required, are calculated using the same method as in forward_transform_truncated_window.

This is vectorised to allow for batch jobs computing the dwt for multiple time series at once; note the shapes of the input and output arrays.

Parameters:

x (jnp.ndarray) – The time-domain signal. Array shape (…, N).

Returns:

w – Wavelet coefficients. Array shape (…, Nt, Nf).

Return type:

jnp.ndarray

Notes

This method is fast. Use this to perform discrete wavelet transforms for production analysis. This method is called by self.dwt.

forward_transform_short_fft(x: Array) → Array

Perform the forward discrete wavelet transform. Transforms the input signal from the time domain into the time-frequency domain.

For the \(m>0\) terms, the wavelet coefficients are calculated using the following expression,

(42)\[w_{nm} = \sqrt{2} \delta t \, \mathrm{Re}\, C_{nm}^* X_n[mq] ,\]

where the short FFT is defined as

(43)\[X_n[j] = \sum_{k=-K/2}^{K/2-1} \exp(2\pi i kj/K) x[nN_f+k] \phi[k].\]

The \(m=0\) terms, if required, are calculated using the same method as in forward_transform_truncated_window.

Parameters:

x (jnp.ndarray) – Array shape (N,). Input time-domain signal to be transformed.

Returns:

w – WDM time-frequency-domain wavelet coefficients.

Return type:

jnp.ndarray of shape (Nt, Nf)

Notes

This method is fairly fast. But forward_transform_fft is usually faster. This is included for testing and debugging purposes.

forward_transform_truncated(x: Array) → Array

Perform the forward discrete wavelet transform. Transforms the input signal from the time domain into the time-frequency domain.

This method computes the wavelet coefficients using the truncated expressions

(44)\[w_{n0} = \delta t\sum_{k=-K/2}^{K/2-1} g_{nm}[k + 2 n N_f] x[k + 2 n N_f] ,\]

(45)\[w_{nm} = \delta t\sum_{k=-K/2}^{K/2-1} g_{nm}[k + n N_f] x[k + n N_f] \quad \mathrm{for} \; m>0 ,\]

where the sum is over the truncated window of length \(K=2qN_f\).

In the above expressions, indices out of bounds of the array are to be understood as wrapping around circularly.

Parameters:

x (jnp.ndarray) – Array shape (N,). Input time-domain signal to be transformed.

Returns:

w – Array shape (Nt, Nf). WDM time-frequency-domain wavelet coefficients.

Return type:

jnp.ndarray

Notes

This method is slow. It is only intended to be used for testing and debugging purposes.

forward_transform_truncated_window(x: Array) → Array

Perform the forward discrete wavelet transform. Transforms the input signal from the time domain into the time-frequency domain.

This method computes the wavelet coefficients using the truncated expressions using the window function:

(46)\[\begin{split}w_{n0} = \delta t \begin{cases} \sum_{k=-K/2}^{K/2-1} x[k+2nN_f]\phi[k] & \mathrm{if}\;n<N_t/2 \\ \sum_{k=-K/2}^{K/2-1} (-1)^k x[k+2nN_f]\phi[k] & \mathrm{if}\;n\geq N_t/2 \\ \end{cases} ,\end{split}\]

(47)\[w_{nm} = \sqrt{2}\delta t \, \mathrm{Re} \sum_{k=-K/2}^{K/2-1} C^*_{nm} \exp\left(\frac{i\pi km}{N_f}\right) x[k+nN_f] \phi[k] \quad \mathrm{for}\; m>0.\]

Parameters:

x (jnp.ndarray) – Array shape (N,). Input time-domain signal to be transformed.

Returns:

w – Array shape (Nt, Nf). WDM time-frequency-domain wavelet coefficients.

Return type:

jnp.ndarray

Notes

This method is slow. It is only intended to be used for testing and debugging purposes.

gnm(n: int, m: int) → Array

Compute the time-domain representation of the wavelets, \(g_{nm}(t)\).

This method computes the frequency-domain wavelets for a single choice of \(n\) and \(m\) and performs and inverse Fourier transform. If you instead want to compute the full wavelet basis for all \(n\) and \(m\) efficiently, use the gnm_basis method.

Parameters:

• n (int) – Wavelet time index.

• m (int) – Wavelet frequency index.

Returns:

gnm – Array shape (N,). The time-domain wavelet.

Return type:

jnp.ndarray

gnm_basis() → Array

Efficient computation of time-domain wavelet basis \(g_{nm}(f)\). Instead of calling the functions for \(\tilde{G}_{nm}(f)\) and performing an inverse Fourier transform, as is done in the gnm method, this function shifts indices of window_TD.

For \(m>0\), the wavelet is given by

(48)\[\begin{split}g_{nm}(t) = \begin{cases} \sqrt{2} (-1)^{mn} \cos\left(\frac{\pi m t}{\Delta T}\right) \phi(t-n\Delta T) & \mathrm{if}\;n+m\;\mathrm{even} \\ \sqrt{2} \sin\left(\frac{\pi m t}{\Delta T}\right) \phi(t-n\Delta T) & \mathrm{if}\;n+m\;\mathrm{odd} \end{cases} .\end{split}\]

For the special case \(m=0\), the wavelet is given by

(49)\[\begin{split}g_{n0}(t) = \begin{cases} \phi(t-2n\Delta T) & \mathrm{if}\;n<N_t/2 \\ \frac{1}{2} \exp(-4\pi i n f \Delta T) \left( \tilde{\Phi}(f-f_{\rm Ny}) + \tilde{\Phi}(f+f_{\rm Ny}) \right) & \mathrm{if}\; n\geq N_t/2 \end{cases}.\end{split}\]

The result is cached to speed up subsequent calls.

Returns:

basis – Array of shape (N, Nt, Nf). The time-domain wavelet basis.

Return type:

jnp.ndarray

idwt(w: Array) → Array

Inverse discrete wavelet transform.

Calls self.inverse_transform. Vectorised to allow for transforming multiple time series at once.

Parameters:

w (jnp.ndarray) – Wavelet coefficients. Array shape=(Nt, Nf) or (…, Nt, Nf).

Returns:

x – Input time series. Array shape=(N,) or (…, N).

Return type:

jnp.ndarray

inverse_transform(w: Array) → Array

Perform the inverse discrete wavelet transform. Transforms the wavelet coefficients from the time-frequency domain into the time domain.

This method computes the inverse dwt using the truncated wavelets. This is also vectorised to allow for batch jobs computing the idwt for multiple sets of wavelet coefficients at once; note the shapes of the input and output arrays.

Parameters:

w (jnp.ndarray) – Wavelet coefficients. Array shape (…, Nt, Nf).

Returns:

x – The time-domain signal. Array shape (…, N).

Return type:

jnp.ndarray

inverse_transform_exact(w: Array) → Array

Perform the inverse discrete wavelet transform. Transforms the wavelet coefficients from the time-frequency domain into the time domain.

This method computes the inverse dwt direcrtly using the expression

(50)\[x[k] = \sum_{n=0}^{N_t-1} \sum_{m=0}^{N_f-1} w_{nm} g_{nm}[k] .\]

This method is slow and very memory inefficient. It is here mainly for testing. Consider using inverse_transform instead.

Parameters:

w (jnp.ndarray) – Array shape (Nt, Nf). WDM time-frequency-domain wavelet coefficients.

Returns:

x – Array shape (N,). The time-domain signal.

Return type:

jnp.ndarray

short_fft(x: Array) → Array

The windowed short FFT of the input.

The input time series is split into \(N_t\) overlapping segments each of length \(K\) and with a hop interval of \(N_f\) between their centres. Each of these segments is then windowed and FFT’d.

(51)\[X_n[j] = \sum_{k=-K/2}^{K/2-1} \exp(2\pi i kj/K) x[nN_f+k] \phi[k]\]

Parameters:

x (jnp.ndarray) – Array shape (N,). Input time series signal to be transformed.

Returns:

windowed_fft – Array shape shape (Nt, K). Short FFT of the input, \(X_n[j]\).

Return type:

jnp.ndarray

validate_parameters() → None

Validate the parameters provided to the WDM_transform __init__ method. Raises an AssertionError if any parameters are invalid.

Return type:

None

wavelet_central_time_frequency(n: Array, m: Array) → Tuple[Array, Array]

Compute the central time \(t_{nm}= n \Delta t\) and the central frequency \(f_{nm} = m \Delta f\) of the wavelet \(g_{nm}(t)\).

The case \(m=0\) is special and is handled separately using

(52)\[t_{n0} = 2n \Delta t ,\]

(53)\[\begin{split}f_{n0} = \begin{cases} 0 & \mathrm{if}\; n<N_t/2 \\ f_{\mathrm{Ny}} & \mathrm{if}\; n\geq N_t/2 \end{cases} .\end{split}\]

Parameters:

• n (jnp.ndarray) – Wavelet time index, dtype=int, shape=(num_n,).

• m (jnp.ndarray) – Wavelet frequency index, dtype=int, shape=(num_m,).

Returns:

• t_nm (jnp.ndarray) – Array of times, shape=(num_n, num_m). The wavelet central times.

• f_nm (jnp.ndarray) – Array of frequencies, shape=(num_n, num_m). The wavelet central frequencies.

wavelet_central_time_frequency_compiled(n: Array, m: Array) → Tuple[Array, Array]

Compiled part of wavelet_central_time_frequency method.

Parameters:

• n (jnp.ndarray) – Wavelet time index, dtype=int, shape=(num_n,).

• m (jnp.ndarray) – Wavelet frequency index, dtype=int, shape=(num_m,).

Returns:

• t_nm (jnp.ndarray) – Array of times, shape=(num_n, num_m). The wavelet central times.

• f_nm (jnp.ndarray) – Array of frequencies, shape=(num_n, num_m). The wavelet central frequencies.