WDM.code.utils package

Submodules

WDM.code.utils.Meyer module

WDM.code.utils.Meyer.Meyer(omega: Array, d: int, A: float, B: float) → Array[source]

This function evaluates the Meyer-type frequency-domain window function:

(54)\[\begin{split}\Phi(\omega) = \begin{cases} \frac{1}{\sqrt{\Delta \Omega}}, & |\omega| < A \\ \frac{1}{\sqrt{\Delta \Omega}} \cos\left( \frac{\pi}{2} \, \nu_d\left( \frac{|\omega| - A}{B} \right) \right), & A \leq |\omega| \leq A + B \\ 0, & \text{otherwise} \end{cases}\end{split}\]

where \(\\nu_d(x)\) is a smooth transition function defined using the normalized incomplete beta function, and \(\\Delta \\Omega = 2A + B\) is the total frequency support of the window.

Parameters:

omega (jnp.ndarray) – Input angular frequency values (radians/second).
d (int) – Steepness parameter controls smoothness of the transition region.
A (float) – Half-width of the flat-top region of the window (radians/second).
B (float) – Width of the transition (roll-off) region (radians/second).

Returns:

phi_w – The window function \(\Phi(\omega)\).

Return type:

jnp.ndarray

Notes

The window is flat in the region \(|\omega| < A\), rolls off smoothly for \(A \leq |\omega| \leq A + B\), and is zero outside that range. The function is symmetric and real-valued.

WDM.code.utils.Meyer.nu_d(x: Array, d: int) → Array[source]

This function evaluates the normalised incomplete beta function.

(55)\[\nu_d(x) = \frac{ \int_0^x \mathrm{d}t \, t^{d-1} (1 - t)^{d-1} } { \int_0^1 \mathrm{d}t \, t^{d-1} (1 - t)^{d-1} }\]

Both the input x and output \(\nu_d(x)\) are in the range [0, 1].

Parameters:

x (jnp.ndarray) – Input array of values in the interval \(0 \leq x \leq 1\).
d (int) – Steepness parameter controls smoothness of the transition region.

Returns:

nu_d – The normalized beta function \(\nu_d(x)\).

Return type:

jnp.ndarray

Notes

This function calls jax.scipy.special.betainc under the hood.

This function returns nan if \(x\) is outside the range [0, 1].

WDM.code.utils.utils module

WDM.code.utils.utils.C_nm(n: int, m: int) → complex[source]

Compute the complex-valued modulation coefficient \(C_{nm}\).

This coefficient alternates between 1 and \(i\) to apply modulation in the WDM transform.

Parameters:

n (int) – Time index.
m (int) – Frequency index.

Returns:

Coefficient \(C_{nm}\), equal to 1 or \(i\) depending on parity of \(n+m\).

Return type:

complex

WDM.code.utils.utils.next_multiple(i: int, N: int) → int[source]

Return smallest integer multiple of N greater than or equal to integer i.

Parameters:

i (int) – The input number.
N (int) – The multiple to align to.

Returns:

j – The next multiple of N.

Return type:

int

Notes

Example with N = 3:

>>> for i in [-4, -3, -2, -1, 0, 1, 2, 3, 4]:
...     print(f"{i} -> {next_multiple(i, 4)}")
-4 -> -3
-3 -> -3
-2 ->  0
-1 ->  0
 0 ->  0
 1 ->  3
 2 ->  3
 3 ->  3
 4 ->  6

WDM.code.utils.utils.overlapping_windows(x: Array, K: int, Nt: int, Nf: int) → Array[source]

Extract overlapping, wrapped windows from input array x.

Parameters:

x (jnp.ndarray, shape (N,)) – Input array to extract windows from.
K (int) – Window length (must be even).
Nt (int) – Number of windows (time steps).
Nf (int) – Hop size between window centers.

Returns:

windows – Array of overlapping windows with wraparound indexing.

Return type:

jnp.ndarray, shape (Nt, K)

Examples

>>> import jax.numpy as jnp
>>> Nt = 4
>>> Nf = 4
>>> K = 8
>>> x = jnp.arange(Nt*Nf)
>>> overlapping_windows(x, K, Nt, Nf)
array([ [12, 13, 14, 15,  0,  1,  2,  3],
        [ 0,  1,  2,  3,  4,  5,  6,  7],
        [ 4,  5,  6,  7,  8,  9, 10, 11],
        [ 8,  9, 10, 11, 12, 13, 14, 15] ])

WDM.code.utils.utils.pad_signal(x: Array, N: int, where: str = 'end') → Array[source]

The transform method requires the input time series signal to have a specific length \(N\). This method can be used to zero-pad any signal to the desired length.

This function also returns a Boolean mask that can be used later to recover arrays of the original length.

Parameters:

x (jnp.ndarray) – Input signal to be padded.
N (int) – Desired length of the output signal.
where (str) – Where to add the padding. Options are ‘end’, ‘start’, or ‘equal’ which puts the zero padding at the end of the signal, the start of the signal, or equally at both ends respectively. Optional.

Returns:

x_padded (jnp.ndarray) – Padded signal to length N, with zeros added at the end.
mask (jnp.ndarray) – Boolean mask indicating the valid part of the padded signal.

Notes

The Boolean mask can be used to get back to the original signal; i.e. x_padded[mask] will recover the original signal, x.