braintools.analysis.spectral.morlet_transform

braintools.analysis.spectral.morlet_transform#

braintools.analysis.spectral.morlet_transform(x, dt, freqs, n_cycles=7.0, *, zero_mean=True)#

Compute the complex Morlet wavelet transform along time axis 0.

For each frequency \(f\) with \(\sigma_t = n_{cycles} / (2\pi f)\), the wavelet [1]

\[W_f(t) = c \, \left(e^{2\pi i f t} - \kappa\right) e^{-t^2 / (2\sigma_t^2)}, \qquad |t| < 5\sigma_t,\]

is sampled at dt (with a sample at exactly t = 0), where \(\kappa = e^{-2(\pi f \sigma_t)^2}\) if zero_mean else 0, and \(c\) makes \(\|W_f\|_2 = \sqrt{2}\) (MNE’s normalisation). The signal is convolved with each wavelet ('same' mode, the output aligned with the input samples, zero padding outside the signal) via the FFT.

Parameters:

  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Real signal of shape (T, ...); time is axis 0. A brainunit.Quantity is accepted.

  • dt (float | Quantity) – Sampling interval (seconds if a plain number). Static.

  • freqs (Sequence[float]) – Wavelet frequencies in Hz, 0 < f <= fs / 2. Static.

  • n_cycles (float | Sequence[float]) – Number of cycles of each wavelet (one value, or one per frequency). Static.

  • zero_mean (bool) – Subtract the Gaussian-weighted offset so each wavelet has zero mean (the default of mne.time_frequency.tfr_array_morlet).

Returns:

tfr – Complex coefficients of shape (n_freqs, T, ...); abs(tfr)**2 is the power and angle(tfr) the phase. Carries the unit of x.

Return type:

Array | Quantity

Raises:

ValueError – If a frequency is not in (0, fs / 2], n_cycles has the wrong length or is not positive, or the longest wavelet is longer than T (as in MNE).

See also

spectrogram

Short-time Fourier power spectrogram.

analytic_signal

Hilbert-transform analytic signal.

Notes

The result equals mne.time_frequency.tfr_array_morlet(x.T[None], sfreq=fs, freqs=freqs, n_cycles=n_cycles, zero_mean=zero_mean, output='complex')[0] with the channel axis moved behind the time axis. The wavelets are designed on the host and their FFTs (length next_fast_len(T + L_max - 1), which avoids circular wrap-around) are constants; the convolution runs in JAX.

The function is jittable and vmappable and differentiable with respect to x.

References

Examples

>>> import jax.numpy as jnp
>>> from braintools.analysis.spectral import morlet_transform
>>> t = jnp.arange(2000) * 1e-3
>>> x = jnp.sin(2 * jnp.pi * 20 * t)
>>> tfr = morlet_transform(x, dt=1e-3, freqs=[10.0, 20.0, 40.0])
>>> tfr.shape
(3, 2000)
>>> power = jnp.mean(jnp.abs(tfr[:, 500:1500]) ** 2, axis=1)
>>> int(jnp.argmax(power))
1