DOCSEARCHbyAlgolia
Submit Docs

ssqueezepy

repository·master·Indexed 21 days ago

https://github.com/overlordgolddragon/ssqueezepy

A high-performance Python library for synchrosqueezing, a reassignment method used to focus time-frequency representations. It provides implementations for both Continuous Wavelet Transform (CWT) and Short-Time Fourier Transform (STFT) to enable precise extraction of instantaneous amplitudes and frequencies. The library supports GPU acceleration via CuPy and PyTorch, multi-threaded CPU execution, and includes tools for frequency ridge extraction using the `extract_ridges` function.

Tokens
2.7K
Snippets
7
Records
12
Agent score
74%

What's inside ssqueezepy

  1. Tuning the penalty term in ridge extraction

    master

    The penalty parameter in extract_ridges is critical for controlling the stability of the extracted ridges. A low or zero penalty allows the ridge to jump frequently between different frequency components, which can lead to unstable results. Increasing the penalty term helps deter these frequency jumps, forcing the algorithm to find more continuous and physically plausible paths through the time-frequency plane.

    If your extracted ridges show erratic jumps in frequency, consider increasing the penalty value.

  2. Configure execution modes via environment variables

    master

    You can control whether ssqueezepy runs on a single CPU thread, multiple CPU threads, or the GPU by setting environment variables. Note that setting SSQ_GPU will override SSQ_PARALLEL.

    • CPU, single thread: Set os.environ['SSQ_PARALLEL'] = '0'
    • CPU, multi-threaded (default): Set os.environ['SSQ_PARALLEL'] = '1'. This uses all available CPU threads via Numba and SciPy.
    • GPU: Set os.environ['SSQ_GPU'] = '1'. This requires CuPy >= 8.0.0 and PyTorch >= 1.8.0. Support for AMD is experimental; NVIDIA is the primary supported hardware.

    Execution modes can be changed without restarting the Python kernel, but if you change the mode, you must re-instantiate any Wavelet objects to ensure they use the new mode.

    import os
# For GPU execution
os.environ['SSQ_GPU'] = '1'

# If you previously created a wavelet, recreate it to apply the new mode
wavelet = Wavelet()

  3. Enable GPU or CPU acceleration in ssqueezepy

    master

    By default, ssqueezepy uses multi-threaded execution. You can control acceleration using environment variables:

    • Disable multi-threading: Set os.environ['SSQ_PARALLEL'] = '0'.
    • Enable GPU acceleration: Set os.environ['SSQ_GPU'] = '1'. This requires CuPy >= 8.0.0 and PyTorch >= 1.8.0 to be installed.
    • Maximum CPU FFT speed: Install pyfftw to optionally speed up FFT operations on the CPU.
    import os
# To disable parallel execution
os.environ['SSQ_PARALLEL'] = '0'

# To enable GPU acceleration
os.environ['SSQ_GPU'] = '1'

  5. Handling edge effects in ridge extraction

    master

    Edge effects in time-frequency maps can make ridge extraction unstable, especially when using the synchrosqueezed transform (ssq_cwt). To mitigate these effects, consider the following:

    1. Signal Padding: Choose an appropriate padtype (e.g., 'wrap', 'reflect', or 'symmetric') when calling ssq_cwt or cwt. The choice of padding should be suited to the nature of your time signal.
    2. Wavelet Selection: The choice of wavelet type affects the localization and the resulting time-frequency representation.
    3. Parameter Tuning: Ridge extraction on syncrosqueezed transforms may require more expertise in tuning parameters like penalty and bw compared to standard CWT to achieve stable results.

  6. Optimize performance with batching and precision

    master

    To maximize throughput, use the following patterns:

    Batched Execution

    All forward transforms (cwt, stft, ssq_cwt, ssq_stft) support batched execution. Instead of looping over individual signals, pass a 2D array where the shape is (n_signals, signal_len). This is more memory-efficient and faster, especially for machine learning workflows.

    Using float32 for speed

    Using float32 instead of the default float64 can nearly double speed and halve memory usage with negligible accuracy loss.

    • For cwt and ssq_cwt: Pass a Wavelet object with the desired dtype, or use a config dictionary: ('gmw', {'dtype': 'float64'}).
    • For stft and ssq_stft: Pass the dtype='float32' keyword argument directly.

  7. Manage memory and caching

    master

    Caching

    ssqueezepy uses caching to speed up repeated computations.

    • Wavelet Caching: cwt and ssq_cwt support cache_wavelet=True (enabled by default if a Wavelet is passed and vectorized=True). This stores and reuses computed wavelets.
    • Numba Caching: The library uses @numba.jit(cache=True) extensively, so methods will run significantly faster on subsequent executions.

    Freeing Memory

    If you are using PyTorch and need to free GPU memory, do not use torch.cuda.empty_cache(). Instead, use standard Python garbage collection:

    import gc
Tx = [] # Clear your large tensors/objects
gc.collect()

  9. Configure CWT scales selection

    master

    When performing a Continuous Wavelet Transform (CWT), you can configure how scales are selected using several parameters. This is useful for controlling the resolution and frequency coverage of your analysis.

    Key parameters:

    • scaletype: Determines the scale distribution. Options are 'log', 'log-piecewise', or 'linear'.
    • preset: Determines the scale range. Options are 'minimal', 'maximal', or 'naive' (not recommended).
    • nv: Number of voices (wavelets per octave). Increasing this results in more scales.
    • downsample: A downsampling factor for higher scales. This is only used if scaletype='log-piecewise'.

    Example configuration:

    import numpy as np

# signal length
N = 2048
t = np.linspace(0, 1, N, endpoint=False)
x = np.cos(2*np.pi * 16 * t) + np.sin(2*np.pi * 64 * t)

# CWT configuration
wavelet = 'gmw'
padtype = 'reflect'
scaletype = 'log-piecewise'
preset = 'maximal'
nv = 32
downsample = 4

  10. Visualize wavelet transforms with units

    master

    To visualize transforms with physical units (like seconds and Hertz), use the cwt, stft, and imshow functions along with scale_to_freq from ssqueezepy.experimental.

    1. Compute the transform using cwt or stft.
    2. Convert wavelet scales to frequencies using scale_to_freq(scales, wavelet, len(x), fs=fs).
    3. Use imshow to plot, passing the calculated frequencies to the yticks parameter.
    import numpy as np
from ssqueezepy import Wavelet, cwt, stft, imshow
from ssqueezepy.experimental import scale_to_freq

N = 2048
fs = 400
t = np.linspace(0, N/fs, N)
x = np.random.randn(N)

# CWT with units
wavelet = Wavelet()
Wx, scales = cwt(x, wavelet)
freqs_cwt = scale_to_freq(scales, wavelet, len(x), fs=fs)

# STFT with units
Sx_full = stft(x)
Sx = Sx_full[::-1]
freqs_stft = np.linspace(1, 0, len(Sx)) * fs/2

# Visualization
ikw = dict(abs=1, xticks=t, xlabel="Time [sec]", ylabel="Frequency [Hz]")
imshow(Wx, **ikw, yticks=freqs_cwt)
imshow(Sx, **ikw, yticks=freqs_stft)

  11. Perform Synchrosqueezing CWT and STFT

    master

    The library provides synchrosqueezing implementations for both the Continuous Wavelet Transform (CWT) and the Short-Time Fourier Transform (STFT). These methods allow for better time-frequency localization and extraction of instantaneous amplitudes and frequencies.

    • ssq_cwt(x): Returns the synchrosqueezed CWT. Returns (Twx, Wx, ...) where Twx is the time-frequency representation and Wx is the wavelet transform.
    • ssq_stft(x): Returns the synchrosqueezed STFT. Returns (Tsx, Sx, ...) where Tsx is the time-frequency representation and Sx is the STFT.
    import numpy as np
from ssqueezepy import ssq_cwt, ssq_stft

# Example signal
x = np.random.randn(2048)

# Synchrosqueezed CWT
Twx, Wx, *_ = ssq_cwt(x)

# Synchrosqueezed STFT
Tsx, Sx, *_ = ssq_stft(x)

  12. Extract frequency ridges using `extract_ridges`

    master

    The extract_ridges function extracts the n_ridges (an integer specified by the user) most prominent frequency ridges from a time-frequency representation. It uses a forward-backward greedy path optimization algorithm that penalizes frequency jumps, similar to the MATLAB function tfridge.

    Key Parameters

    • n_ridges: The number of most prominent ridges to extract.
    • penalty: A term used to penalize frequency jumps during the optimization process. Increasing this value helps deter unstable frequency jumps.
    • bw: Bandwidth parameter (used in the examples).
    • scales: The scales used in the time-frequency representation.

    Usage with CWT vs. SSQ_CWT

    You can apply ridge extraction to either a standard Continuous Wavelet Transform (CWT) or a Synchrosqueezed CWT (SSQ_CWT). Note that ridge extraction on syncrosqueezed transforms may be more sensitive to edge effects and requires careful parameter tuning for stability.

    # Example: Extracting 2 ridges from a CWT representation
# Wx: CWT coefficients, scales: wavelet scales, penalty: jump penalty
ridge_idxs = extract_ridges(Wx, scales, penalty, n_ridges=2, bw=25)

# Example: Extracting 2 ridges from an SSQ_CWT representation
# Tx: Synchrosqueezed coefficients
ridge_idxs = extract_ridges(Tx, scales, penalty, n_ridges=2, bw=4)