Is ssq_cwt designed for negative frequency analysis?

[aspirin-1218]

Hello,

First of all, thank you for sharing the great work.

I am trying to do time-frequency analysis on radar doppler signals (complex valued) and thus being able to analyze negative frequency really matters. We've experimented STFT and would like to test with CWT. After checking the ssq_cwt, I am not sure if we are able to analyze negative frequency with that. Could you please give us a simple example if I misunderstood the usage of ssq_cwt?

[OverLordGoldDragon]

No support currently. Open to adding in the future.

The following is a WIP, but I've not fully tested so I won't yet commit it; you can try it by replacing experimental.py with it.

expand

# -*- coding: utf-8 -*-
import warnings
import numpy as np
from numpy.fft import fft, ifft
from .wavelets import Wavelet, center_frequency
from .utils import backend as S, cwt_scalebounds
from .utils.common import EPS32, EPS64, p2up, trigdiff
from .ssqueezing import ssqueeze
from ._ssq_cwt import phase_cwt, phase_cwt_num, ssq_cwt
from ._ssq_stft import phase_stft, _make_Sfs
from ._cwt import cwt


__all__ = ['freq_to_scale', 'scale_to_freq', 'phase_ssqueeze', 'phase_transform',
           'cwt_complex', 'ssq_cwt_complex']


def freq_to_scale(freqs, wavelet, N, fs=1, n_search_scales=None, kind='peak',
                  base=2):
    """Convert frequencies to scales.

    # Arguments:
        freqs: np.ndarray
            1D array of frequencies. Must range between 0 and `N/fs/2` (Nyquist).

        wavelet: wavelets.Wavelet / str / tupe[str, dict]
            Wavelet.

        N: int
            `len(x)` of interest.

        fs: int
            Sampling rate in Hz.

        n_search_scales: int / None
            This method approximates the conversion of scales. Higher = better
            accuracy, but takes longer. Defaults to `10 * len(freqs)`.

        kind: str
            Mapping to use, one of: 'peak', 'energy', 'peak-ct'.
            See `help(ssqueezepy.center_frequency)`.

        base: int
            Base of exponent of `freqs`. Defaults to 2.
            `freqs` can be any distributed in any way, including mix of log
            and linear, so the base only helps improve the search if it matches
            that of `freqs`. If `freqs` is purely exponential, then
            `base = np.diff(np.log(freqs))[0] * 2.718281828`.

    # Returns:
        scales: np.ndarray
            1D arrays of scales.
    """
    def logb(x, base=2):
        return np.log(x) / np.log(base)

    def log(x):
        return logb(x, base)

    freqs = freqs / fs  # convert to unitless, [0., 0.5)
    assert np.all(freqs >= 0),       "frequencies must be positive"
    assert freqs.max() <= 0.5,       "max frequency must be 0.5"
    assert freqs.max() == freqs[-1], "max frequency must be last sample"
    assert freqs.min() == freqs[0],  "min frequency must be first sample"

    M = len(freqs)
    if n_search_scales is None:
        n_search_scales = 10 * M
    smin, smax = cwt_scalebounds(wavelet, N, preset='maximal', use_padded_N=False)
    search_scales = np.logspace(log(smin), log(smax), n_search_scales, base=base)

    w_from_scales = []
    for scale in search_scales:
        w = center_frequency(wavelet, scale, N, kind=kind)
        w_from_scales.append(min(max(w, 0), np.pi))
    f_from_scales = np.array(w_from_scales) / (2*np.pi)

    # pick closest match
    fmin, fmax = freqs.min(), freqs.max()
    smax = search_scales[np.argmin(np.abs(f_from_scales - fmin))]
    smin = search_scales[np.argmin(np.abs(f_from_scales - fmax))]
    # make scales between found min and max
    scales = np.logspace(log(smax), log(smin), M, base=base)

    return scales


def scale_to_freq(scales, wavelet, N, fs=1, padtype='reflect'):
    """Convert scales to frequencies.

    # Arguments:
        freqs: np.ndarray
            1D array of frequencies. Must range between 0 and `N/fs/2` (Nyquist).

        wavelet: wavelets.Wavelet / str / tupe[str, dict]
            Wavelet.

        N: int
            `len(x)` of interest.

        fs: int
            Sampling rate in Hz.

        padtype: str / None
            `padtype` used in the transform. Used to determine the length
            of wavelets used in the transform: `None` uses half the length
            relative to `not None`.
            The exact value doesn't matter, only whether it's `None` or not.

    # Returns:
        freqs: np.ndarray
            1D arrays of frequencies.
    """
    # process args
    if isinstance(scales, float):
        scales = np.array([scales])
    wavelet = Wavelet._init_if_not_isinstance(wavelet)

    # evaluate wavelet at `scales`
    Npad = p2up(N)[0] if padtype is not None else N
    psis = wavelet(scale=scales, N=Npad)
    if hasattr(psis, 'cpu'):
        psis = psis.cpu().numpy()
    # find peak indices
    idxs = np.argmax(psis, axis=-1)

    # check
    # https://github.com/OverLordGoldDragon/ssqueezepy/issues/41
    if np.any(idxs > Npad//2) or 0 in idxs:
        warnings.warn("found potentially ill-behaved wavelets (peak indices at "
                      "negative freqs or at dc); will round idxs to 1 or N/2")
        n_psis = len(psis)
        for i, ix in enumerate(idxs):
            if ix > Npad//2 or ix == 0:
                if i > n_psis // 2:  # low freq
                    idxs[i] = 1
                else:  # high freq
                    idxs[i] = Npad//2
    # convert
    freqs = idxs / Npad  # [0, ..., .5]
    assert freqs.min() >= 0,   freqs.min()
    assert freqs.max() <= 0.5, freqs.max()

    freqs *= fs   # [0, ..., fs/2]
    return freqs


def phase_ssqueeze(Wx, dWx=None, ssq_freqs=None, scales=None, Sfs=None, fs=1.,
                   t=None, squeezing='sum', maprange=None, wavelet=None,
                   gamma=None, was_padded=True, flipud=False,
                   rpadded=False, padtype=None, N=None, n1=None,
                   difftype=None, difforder=None,
                   get_w=False, get_dWx=False, transform='cwt'):
    """Take `phase_transform` then `ssqueeze`. Can be used on an arbitrary
    CWT/STFT-like time-frequency transform `Wx`.
    Experimental; prefer `ssq_cwt` & `ssq_stft`.
    # Arguments:
        Wx, dWx (see w), ssq_freqs, scales, Sfs, fs, t, squeezing, maprange,
        wavelet, gamma, was_padded, flipud:
            See `help(ssqueezing.ssqueeze)`.
        rpadded: bool (default None) / None
            Whether `Wx` (and `dWx`) is passed in padded. `True` will unpad
            `Wx` and `dWx`  before SSQ. Also, if `dWx` is None:
                - `rpadded==False`: will pad `Wx` in computing `dWx` if
                `padtype!=None`, then unpad both before SSQ
                - `rpadded==True`: won't pad `Wx` regardless of `padtype`
        padtype: str / None
            Used if `rpadded==False`. See `help(utils.padsignal)`. Note that
            padding `Wx` isn't same as passing padded `Wx` from `cwt`, but it
            can get close.
        N, n1: int / None
            Needed if `rpadded==True` to unpad `Wx` & `dWx` as `Wx[:, n1:n1 + N]`.
        difftype, difforder: str
            Used if `dWx = None` and `transform == 'cwt'`; see `help(ssq_cwt)`.
        get_w, get_dWx: bool
            See `help(ssq_cwt)`.
    # Returns:
        Tx, Wx, ssq_freqs, scales, Sfs, w, dWx
    """
    w, Wx, dWx, Sfs, gamma = phase_transform(
        Wx, dWx, difftype, difforder=difforder, gamma=gamma, rpadded=rpadded,
        padtype=padtype, N=N, n1=n1, get_w=get_w, fs=fs, transform=transform)

    if w is not None and not get_dWx:
        dWx = None

    if maprange is None:
        maprange = 'peak' if transform == 'cwt' else 'maximal'
    Tx, ssq_freqs = ssqueeze(Wx, w, ssq_freqs, scales, Sfs, fs=fs, t=t,
                             squeezing=squeezing, maprange=maprange,
                             wavelet=wavelet, gamma=gamma, was_padded=was_padded,
                             flipud=flipud, dWx=dWx, transform=transform)
    return Tx, Wx, ssq_freqs, scales, Sfs, w, dWx


def phase_transform(Wx, dWx=None, difftype='trig', difforder=4, gamma=None,
                    fs=1., Sfs=None, rpadded=False, padtype='reflect', N=None,
                    n1=None, get_w=False, transform='cwt'):
    """Unified method for CWT & STFT SSQ phase transforms.
    See `help(_ssq_cwt.phase_cwt)` and `help(_ssq_stft.phase_stft)`.
    """
    def _cwt(Wx, dWx, fs, gamma, N, n1, difftype, difforder, rpadded, padtype,
             get_w):
        # infer `N` and/or `n1`
        if N is None and not rpadded:
            N = Wx.shape[-1]
        if n1 is None:
            _, n1, _ = p2up(N)
        # compute `dWx` if not supplied
        if dWx is None:
            dWx = trigdiff(Wx, fs, padtype, rpadded, N=N, n1=n1, transform='cwt')

        if get_w:
            if difftype == 'trig':
                # calculate instantaneous frequency directly from the
                # frequency-domain derivative
                w = phase_cwt(Wx, dWx, difftype, gamma)
            elif difftype == 'phase':
                # !!! bad; yields negatives, and forcing abs(w) doesn't help
                # calculate inst. freq. from unwrapped phase of CWT
                w = phase_cwt(Wx, None, difftype, gamma)
            elif difftype == 'numeric':
                # !!! tested to be very inaccurate for small scales
                # calculate derivative numericly
                Wx = Wx[:, (n1 - 4):(n1 + N + 4)]
                dt = 1 / fs
                w = phase_cwt_num(Wx, dt, difforder, gamma)
        else:
            w = None
        return w, Wx, dWx

    def _stft(Wx, dWx, fs, gamma, Sfs, get_w):
        if Sfs is None:
            Sfs = _make_Sfs(Wx, fs)
        if get_w:
            w = phase_stft(Wx, dWx, Sfs, gamma)
        else:
            w = None
        return w, Wx, dWx, Sfs

    # validate args
    if transform == 'stft' and dWx is None:
        raise NotImplementedError("`phase_transform` without `dWx` for "
                                  "STFT is not currently supported.")
    if rpadded and N is None:
        raise ValueError("`rpadded=True` requires `N`")
    if Wx.ndim > 2 and get_w:
        raise NotImplementedError("`get_w=True` unsupported with batched input.")

    # gamma
    if gamma is None:
        gamma = 10 * (EPS64 if S.is_dtype(Wx, 'complex128') else EPS32)

    # take phase transform if `get_w` else only compute `dWx` (if None)
    if transform == 'cwt':
        w, Wx, dWx = _cwt(Wx, dWx, fs, gamma, N, n1, difftype, difforder,
                          rpadded, padtype, get_w)
        Sfs = None
    elif transform == 'stft':
        w, Wx, dWx, Sfs = _stft(Wx, dWx, fs, gamma, Sfs, get_w)

    return w, Wx, dWx, Sfs, gamma


def cwt_complex(x, wavelet, *a, **k):
    """https://github.com/OverLordGoldDragon/ssqueezepy/issues/117"""
    return _transform_complex('cwt', x, wavelet, *a, **k)


def ssq_cwt_complex(x, wavelet, *a, **k):
    """https://github.com/OverLordGoldDragon/ssqueezepy/issues/117"""
    return _transform_complex('ssq_cwt', x, wavelet, *a, **k)


def _transform_complex(fn_name, x, wavelet, *a, **k):
    xr_pos, xr_neg = _complex_to_real_pair(x)

    # generate wavelet array if it wasn't already
    if not hasattr(wavelet, '_Psih'):
        _ = cwt(x.real, wavelet)

    # fetch transform
    fn = {'cwt': cwt, 'ssq_cwt': ssq_cwt}[fn_name]

    # take positive transform
    Wx_pos, *_ = fn(xr_pos, wavelet, *a, **k)

    # make negative wavelet
    wavelet._Psih[:, 1:] = wavelet._Psih[:, 1:][:, ::-1]

    # take negative transform; guard with try-finally to restore orig wavelet
    try:
        Wx_neg, *_ = fn(xr_neg, wavelet, *a, **k)
    except:
        pass
    finally:
        # restore original wavelet
        wavelet._Psih[:, 1:] = wavelet._Psih[:, 1:][:, ::-1]

    # concatenate
    Wx = np.vstack([Wx_pos, Wx_neg[::-1]])

    return Wx


def _complex_to_real_pair(x):
    N = len(x)

    xf = fft(x)
    if not np.allclose(xf[0], 0) or not np.allclose(xf[N//2].imag, 0):
        # TODO why, this should just work.
        warnings.warn("Detected non-zero mean or non-zero imaginary Nyquist "
                      "in signal; will zero both.")
    xf[0] = 0
    xf[N//2] = 0

    xrf_pos = xf.copy()
    xrf_neg = xf.copy()

    # conjugate-mirror positives onto negatives:
    # x-real's positive freqs now represent original's positive portion
    xrf_pos[N//2 + 1:] = xf[1:N//2][::-1]
    xrf_pos[N//2 + 1:].imag *= -1
    # conjugate-mirror negatives onto positives:
    # x-real's positive freqs now represent original's negative portion
    xrf_neg[1:N//2] = xf[N//2 + 1:][::-1]
    xrf_neg[1:N//2].imag *= -1

    # to time
    xr_pos, xr_neg = map(ifft, (xrf_pos, xrf_neg))

    # assert real
    # for g in (xr_pos, xr_neg):
    #     assert np.abs(g.imag).max() < 1e-15
    # cast to real type
    xr_pos, xr_neg = xr_pos.real, xr_neg.real

    # assert portions match as intended
    # xrpf, xrnf = map(fft, (xr_pos, xr_neg))
    # assert np.allclose(xrpf[1:N//2], xf[1:N//2])
    # assert np.allclose(xrnf[N//2 + 1:], xf[N//2 + 1:])

    return xr_pos, xr_neg

from ssqueezepy.experimental import cwt_complex, ssq_cwt_complex
out = cwt_complex(x, wavelet)

Caveats:

• assumes wavelets are symmetric in positive and negative frequencies (which Morlet and GMW are, I've not checked others)
• only works on even-length x
• padding isn't handled optimally
• imaginary DC and Nyquist are zeroed (tiny difference for most x)
• not fully tested