def test_ssq_stft_batched():
"""Ensure batched (2D `x`) inputs output same as if samples fed separately,
and agreement between CPU & GPU.
"""
np.random.seed(0)
x = np.random.randn(4, 256)
for dtype in ('float64', 'float32'):
os.environ['SSQ_GPU'] = '0'
kw = dict(astensor=False, dtype=dtype)
Tx0, Sx0, *_ = ssq_stft(x, **kw)
Tx00 = np.zeros(Tx0.shape, dtype=Tx0.dtype)
Sx00 = Tx00.copy()
for i, _x in enumerate(x):
out = ssq_stft(_x, **kw)
Tx00[i], Sx00[i] = out[0], out[1]
if CAN_GPU:
os.environ['SSQ_GPU'] = '1'
Tx1, Sx1, *_ = ssq_stft(x, **kw)
atol = 1e-12 if dtype == 'float64' else 1e-6
adiff_Tx000 = np.abs(Tx00 - Tx0).mean()
adiff_Sx000 = np.abs(Sx00 - Sx0).mean()
assert np.allclose(Sx00, Sx0), (dtype, adiff_Sx000)
assert np.allclose(Tx00, Tx0), (dtype, adiff_Tx000)
if CAN_GPU:
adiff_Tx01 = np.abs(Tx0 - Tx1)
adiff_Sx01 = np.abs(Sx0 - Sx1)
assert np.allclose(Sx0, Sx1, atol=atol), (dtype, adiff_Sx01)
assert np.allclose(Tx0, Tx1, atol=atol), (dtype, adiff_Tx01)
def tf_transforms(x,
t,
wavelet='morlet',
window=None,
padtype='wrap',
penalty=.5,
n_ridges=2,
cwt_bw=15,
stft_bw=15,
ssq_cwt_bw=4,
ssq_stft_bw=4):
kw_cwt = dict(t=t, padtype=padtype)
kw_stft = dict(fs=1 / (t[1] - t[0]), padtype=padtype, flipud=1)
Twx, Wx, ssq_freqs_c, scales, *_ = ssq_cwt(x, wavelet, **kw_cwt)
Tsx, Sx, ssq_freqs_s, Sfs, *_ = ssq_stft(x, window, **kw_stft)
Sx, Sfs = Sx[::-1], Sfs[::-1]
ckw = dict(penalty=penalty, n_ridges=n_ridges, transform='cwt')
skw = dict(penalty=penalty, n_ridges=n_ridges, transform='stft')
cwt_ridges = extract_ridges(Wx, scales, bw=cwt_bw, **ckw)
ssq_cwt_ridges = extract_ridges(Twx, ssq_freqs_c, bw=ssq_cwt_bw, **ckw)
stft_ridges = extract_ridges(Sx, Sfs, bw=stft_bw, **skw)
ssq_stft_ridges = extract_ridges(Tsx, ssq_freqs_s, bw=ssq_stft_bw, **skw)
viz(x, Wx, cwt_ridges, scales, ssq=0, transform='cwt', show_x=1)
viz(x, Twx, ssq_cwt_ridges, ssq_freqs_c, ssq=1, transform='cwt', show_x=0)
viz(x, Sx, stft_ridges, Sfs, ssq=0, transform='stft', show_x=0)
viz(x,
Tsx,
ssq_stft_ridges,
ssq_freqs_s,
ssq=1,
transform='stft',
show_x=0)
def test_ssq_stft():
"""Same as `test_stft` except don't test `hop_len` or `modulated` since
only `1` and `True` are invertible (by the library, and maybe theoretically).
`window_scaling=.5` has >x2 greater MAE for some reason. May look into.
"""
th = 1e-1
for N in (128, 129):
x = np.random.randn(N)
for n_fft in (120, 121):
for window_scaling in (1., .5):
if window_scaling == 1:
window = None
else:
window = get_window(window, win_len=n_fft//1, n_fft=n_fft)
window *= window_scaling
Sx, *_ = ssq_stft(x, window=window, n_fft=n_fft)
xr = issq_stft(Sx, window=window, n_fft=n_fft)
txt = ("\nSSQ_STFT: (N, n_fft, window_scaling) = ({}, {}, {})"
).format(N, n_fft, window_scaling)
assert len(x) == len(xr), "%s != %s %s" % (N, len(xr), txt)
mae = np.abs(x - xr).mean()
assert mae < th, "MAE = %.2e > %.2e %s" % (mae, th, txt)
def test_misc():
_ = cwt(np.random.randn(128), 'gmw', cache_wavelet=True)
_ = cwt(np.random.randn(128),
Wavelet(),
cache_wavelet=True,
vectorized=False)
_ = ssq_stft(np.random.randn(100), get_w=1, get_dWx=1)
pass_on_error(cwt, np.random.randn(2, 2, 2))
pass_on_error(cwt, 5)
pass_on_error(ssq_stft, np.random.randn(2, 2, 2), get_w=1)
def test_ssq_stft():
N = 256
tsigs = TestSignals(N=N)
gpu_atol = 1e-5
for dtype in ('float64', 'float32'):
x = tsigs.par_lchirp()[0].astype(dtype)
kw = dict(modulated=1, n_fft=128, dtype=dtype, astensor=False)
os.environ['SSQ_GPU'] = '0'
Tx00 = ssq_stft(x, **kw, get_w=1)[0]
Tx01 = ssq_stft(x, **kw, get_w=0)[0]
if CAN_GPU:
os.environ['SSQ_GPU'] = '1'
Tx10 = ssq_stft(x, **kw, get_w=1)[0]
Tx11 = ssq_stft(x, **kw, get_w=0)[0]
adiff0001 = np.abs(Tx00 - Tx01).mean()
assert np.allclose(Tx00, Tx01), (dtype, adiff0001)
if CAN_GPU:
adiff0010 = np.abs(Tx00 - Tx10).mean()
adiff0011 = np.abs(Tx00 - Tx11).mean()
assert np.allclose(Tx00, Tx10, atol=gpu_atol), (dtype, adiff0010)
assert np.allclose(Tx00, Tx11, atol=gpu_atol), (dtype, adiff0011)
def time_ssq_stft(x, dtype, n_fft):
for _ in range(3):
_ = ssq_stft(x, dtype=dtype, n_fft=n_fft)
del _
gc.collect()
return timeit(lambda: ssq_stft(x, dtype=dtype, n_fft=n_fft))
num: '',
f'{num}-cwt': time_cwt(x, dtype, scales, cache_wavelet),
f'{num}-stft': time_stft(x, dtype, n_fft),
f'{num}-ssq_cwt': time_ssq_cwt(x, dtype, scales, cache_wavelet,
ssq_freqs),
f'{num}-ssq_stft': time_ssq_stft(x, dtype, n_fft)
}
#%%# Setup ###################################################################
# warmup
x = np.random.randn(1000)
for dtype in ('float32', 'float64'):
wavelet = Wavelet(dtype=dtype)
_ = ssq_cwt(x, wavelet, cache_wavelet=False)
_ = ssq_stft(x, dtype=dtype)
del _, wavelet
#%%# Prepare reusable parameters such that STFT & CWT output shapes match ####
N0, N1 = 10000, 160000 # selected such that CWT pad length ratios are same
n_rows = 300
n_fft = n_rows * 2 - 2
wavelet = Wavelet()
scales = process_scales('log-piecewise', N1, wavelet=wavelet)[:n_rows]
ssq_freqs = _compute_associated_frequencies(scales,
N1,
wavelet,
'log-piecewise',
maprange='peak',
was_padded=True,