OverLordGoldDragon / ssqueezepy

Synchrosqueezing, wavelet transforms, and time-frequency analysis in Python
MIT License
599 stars 92 forks source link

Custom `ssq_freqs` with custom shape [workaround included] #107

Open ARBolton96 opened 1 month ago

ARBolton96 commented 1 month ago

Hello,

I have a known set of frequencies that I am most interested in evaluating the synchrosqueezed CWT of a signal for. Let's say this frequency vector f_vect is defined as:

f_vect = 10**(np.arange(np.log10(fmin), np.log10(fmax), log_fstep))

I was hoping to provide this f_vect as the ssq_freqs optional input to the ssq_cwt() function, so as to get an output CWT whose columns match the time samples of a given signal and whose rows correspond to the desired frequencies (or scales that map directly to my given frequency vector). In pseudo-code, this would look like:

SSQ_CWT, CWT, ssq_freqs, scales = ssq_cwt(data, fs=fs, wavelet='gmw', ssq_freqs=f_vect)

print(SSQ_CWT.shape[0] == ssq_freqs.shape[0])
> True

But instead, the shape of the SSQ_CWT array is the default shape of ( len(scales), len(data) ), rather than my expected ( len(f_vect), len(data) ). As I understood the documentation, passing a np.ndarray into the ssq_freqs option should map the CWT scales onto the desired synchrosqueeze frequencies. What am I missing here?

OverLordGoldDragon commented 2 days ago

Thanks for the report, I confirmed the bug.

I've not attended to this feature beyond the original implementation, and it appears fully custom ssq_freqs wasn't supported. Fixing this isn't trivial, I might get around to it in the future. It's also not overwhelming; I'll accept a PR for it.

In the meantime, one can do a workaround with minimal performance overhead:

WORKAROUND EXAMPLE ```python # -*- coding: utf-8 -*- import numpy as np from ssqueezepy import ssq_cwt from ssqueezepy.visuals import imshow def extend_log(vi, ssq_freqs0): """Assumes `ssq_freqs` and `vi` are log-spaced.""" vi_len = len(vi) if min(vi) >= min(ssq_freqs0): start_idx = 1 else: start_idx = 0 v = np.zeros_like(ssq_freqs0) dvil = np.log10(vi[1]) - np.log10(vi[0]) vi_first = vi[0] vi_last = vi[-1] v[start_idx:start_idx + vi_len] = vi if start_idx == 1: v[0] = 10**(np.log10(vi_first) - dvil) n_pts = len(v) - vi_len - 1 else: n_pts = len(v) - vi_len startl = np.log10(vi_last) + dvil endl = startl + (n_pts - 1) * dvil v_ext1 = np.logspace(startl, endl, n_pts) if start_idx == 1: v[vi_len + 1:] = v_ext1 else: v[vi_len:] = v_ext1 # make crop slice crop_slc = slice(len(v) - start_idx - vi_len, len(v) - start_idx) return v, crop_slc def energy(x): return np.sum(abs(x)**2) # make signal, embed sine np.random.seed(0) x = np.random.randn(2048) x += np.cos(2*np.pi * (.05 * len(x)) * np.linspace(0, 1, len(x))) # first transform Sx_full, _, ssq_freqs0, scales0 = ssq_cwt(x, scales='log') # your frequencies vi = np.logspace(np.log10(.001), np.log10(.1), 100) ssq_freqs, crop_slc = extend_log(vi, ssq_freqs0) # final transform Sx, Wx, *_ = ssq_cwt(x, scales=scales0, ssq_freqs=ssq_freqs) Sx = Sx[crop_slc] #%% Visualize print(Sx.shape, Wx.shape, ssq_freqs.shape, sep='\n') imshow(Sx_full, abs=1, yticks=ssq_freqs0, h=.8, title="original ssq_freqs") imshow(Sx, abs=1, yticks=vi[::-1], h=.8, title="custom ssq_freqs") #%% Confirm energies E_x0 = energy(Sx_full.real.sum(axis=0)) E_x = energy(Sx.real.sum(axis=0)) E_x0_ext = energy(Sx_full[np.argmin(abs(ssq_freqs0 - max(vi))):].real.sum(axis=0)) print("%-4.3g -- `x` energy per Sx_full" % E_x0) print("%-4.3g -- `x` energy per Sx" % E_x) print("%-4.3g -- `x` energy per Sx_full inside of `vi` freqs" % E_x0_ext) # NOTE: exact match not guaranteed ```
101  -- `x` energy per Sx_full
45.9 -- `x` energy per Sx
45.9 -- `x` energy per Sx_full inside of `vi` freqs

image

How it works: we make the desired ssq_freqs a subset of the full required format. Let the portion of interest be vi, and the full input be v:

Note, we can only validate energy via inversion, per energy of transform vs energy of signal.

As an aside, I'm not a fan of ssq_freqs docs saying Scale-frequency mapping is only approximate and wavelet-dependent, the "only approximate" part. I guess I mimicked/copied that in early days; it's not "wrong" per se but can be misleading, an explanation here.

OverLordGoldDragon commented 2 days ago

Note, original scales can be any (can remove scales='log' in my example).