def time_cwt(x, dtype, scales, cache_wavelet):
wavelet = Wavelet(dtype=dtype)
if cache_wavelet:
for _ in range(3): # warmup run
_ = cwt(x, wavelet, scales=scales, cache_wavelet=True)
del _
gc.collect()
return timeit(
lambda: cwt(x, wavelet, scales=scales, cache_wavelet=cache_wavelet))
def test_cwt_batched_for_loop():
"""Ensure basic batched cwt works with both `vectorized`."""
os.environ['SSQ_GPU'] = '0'
np.random.seed(0)
x = np.random.randn(4, 256)
for dtype in ('float64', 'float32'):
Wx0, *_ = cwt(x, _wavelet(dtype=dtype), vectorized=True)
Wx1, *_ = cwt(x, _wavelet(dtype=dtype), vectorized=False)
adiff_Wx01 = np.abs(Wx0 - Wx1)
assert np.allclose(Wx0, Wx1), (dtype, adiff_Wx01.mean())
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_cwt():
x = np.random.randn(64)
Wx, *_ = cwt(x, 'morlet', vectorized=True)
_ = icwt(Wx, 'morlet', one_int=True)
_ = icwt(Wx, 'morlet', one_int=False)
Wx2, *_ = cwt(x, 'morlet', vectorized=False)
mae = np.mean(np.abs(Wx - Wx2))
assert mae <= 1e-16, f"MAE = {mae} > 1e-16 for for-loop vs vectorized `cwt`"
_ = est_riskshrink_thresh(Wx, nv=32)
_ = _icwt_norm(scaletype='linear', l1_norm=False)
x[0] = np.nan
x[1] = np.inf
x[2] = -np.inf
_ = cwt(x, 'morlet', vectorized=False, derivative=True, l1_norm=False)
def test_cwt_for_loop():
"""Ensure `vectorized=False` runs on GPU and outputs match `=True`."""
if not CAN_GPU:
return
np.random.seed(0)
x = np.random.randn(256)
kw = dict(derivative=True, astensor=False)
os.environ['SSQ_GPU'] = '1'
for dtype in ('float64', 'float32'):
Wx0, _, dWx0 = cwt(x, _wavelet(dtype=dtype), vectorized=False, **kw)
Wx1, _, dWx1 = cwt(x, _wavelet(dtype=dtype), vectorized=True, **kw)
adiff_Wx = np.abs(Wx0 - Wx1)
adiff_dWx = np.abs(dWx0 - dWx1)
atol = 1e-12 if dtype == 'float64' else 1e-6
assert np.allclose(Wx0, Wx1, atol=atol), (dtype, adiff_Wx.mean())
assert np.allclose(dWx0, dWx1, atol=atol), (dtype, adiff_dWx.mean())
def test_parallel():
"""Ensure `parallel=True` output matches that of `=False`."""
for N in (255, 512):
x = np.random.randn(N)
Wx, scales = cwt(x)
out0 = extract_ridges(Wx, scales, parallel=False)
out1 = extract_ridges(Wx, scales, parallel=True)
adiff = np.abs(out0 - out1)
assert np.allclose(out0, out1), "N=%s, Max err: %s" % (N, adiff.max())
def test_trigdiff():
"""Ensure `trigdiff` matches `cwt(derivative=True)`."""
os.environ['SSQ_GPU'] = '0'
N = 256
x = np.random.randn(N)
Wx, _, dWx = cwt(x, derivative=True, rpadded=True)
_, n1, _ = utils.p2up(N)
dWx2 = utils.trigdiff(Wx, rpadded=True, N=N, n1=n1)
dWx = dWx[:, n1:n1 + N]
mae = np.mean(np.abs(dWx - dWx2))
th = 1e-15 if dWx.dtype == np.cfloat else 1e-7
assert mae < th, mae
def test_cwt_higher_order():
os.environ['SSQ_GPU'] = '0'
N = 256
tsigs = TestSignals()
x, t = tsigs.par_lchirp(N=N)
x += x[::-1]
for noise in (False, True):
if noise:
x += np.random.randn(len(x))
Wx_k, scales = cwt(x, 'gmw', order=range(3), average=False)
visuals.viz_cwt_higher_order(Wx_k, scales, 'gmw')
print("=" * 80)
_ = cwt(x, ('gmw', {
'norm': 'energy'
}),
order=(0, 1),
average=True,
l1_norm=False)
_ = cwt(x, 'gmw', order=1, average=False, derivative=True)
def test_cwt():
errs = []
for fn in test_fns:
x, ts = fn(2048)
for l1_norm in (True, False):
# 'linear' default can't handle low frequencies for large N
kw = dict(wavelet=wavelet, scales='log', l1_norm=l1_norm)
Wx, *_ = cwt(x, t=ts, **kw)
xrec = icwt(Wx, one_int=True, **kw)
errs.append(round(mad_rms(x, xrec), 5))
assert errs[-1] < th, (errs[-1], fn.__name__,
f"l1_norm: {l1_norm}")
print("\ncwt PASSED\nerrs:", ', '.join(map(str, errs)))
def test_cwt():
errs = []
for fn in test_fns:
x, ts = fn(2048)
for l1_norm in (True, False):
# 'linear' default can't handle low frequencies for large N
kw = dict(wavelet=wavelet, scales='log', l1_norm=l1_norm, nv=32)
Wx, *_ = cwt(x, t=ts, **kw)
xrec = icwt(Wx, one_int=True, **kw)
errs.append(round(mad_rms(x, xrec), 5))
title = f"abs(CWT) | l1_norm={l1_norm}"
title = "abs(SSQ_CWT) | {}, l1_norm={}".format(fn.__qualname__,
l1_norm)
_maybe_viz(Wx, x, xrec, title, errs[-1])
assert errs[-1] < th, (errs[-1], fn.__name__, f"l1_norm: {l1_norm}")
print("\ncwt PASSED\nerrs:", ', '.join(map(str, errs)))
def test_phase_cwt():
os.environ['SSQ_GPU'] = '0'
x = TestSignals(N=1000).par_lchirp()[0]
x += x[::-1]
wavelet = Wavelet()
scales = process_scales('log', len(x), wavelet, nv=32)[:240]
Wx, _, dWx = cwt(x,
wavelet,
scales=scales,
derivative=True,
cache_wavelet=1)
for dtype in ('complex128', 'complex64'):
# Wx = np.random.randn(100, 8192).astype(dtype) * (1 + 2j)
# dWx = np.random.randn(100, 8192).astype(dtype) * (2 - 1j)
Wx, dWx = Wx.astype(dtype), dWx.astype(dtype)
if CAN_GPU:
Wxt = torch.tensor(Wx, device='cuda')
dWxt = torch.tensor(dWx, device='cuda')
gamma = 1e-2
_out = (dWx / Wx).imag / (2 * np.pi)
_out[np.abs(Wx) < gamma] = np.inf
_out = np.abs(_out)
out0 = phase_cwt_cpu(Wx, dWx, gamma, parallel=False)
out1 = phase_cwt_cpu(Wx, dWx, gamma, parallel=True)
if CAN_GPU:
out2 = phase_cwt_gpu(Wxt, dWxt, gamma).cpu().numpy()
with np.errstate(invalid='ignore'):
mape0_ = _noninf_mean(np.abs(_out - out0) / np.abs(_out))
mape01 = _noninf_mean(np.abs(out0 - out1) / np.abs(out0))
if CAN_GPU:
mape02 = _noninf_mean(np.abs(out0 - out2) / np.abs(out0))
assert np.allclose(out0, _out), ("base", dtype, mape0_)
assert np.allclose(out0, out1), ("parallel", dtype, mape01)
if CAN_GPU:
assert np.allclose(out0, out2), ("gpu", dtype, mape02)
xo = x.copy()
np.random.seed(4)
x += np.sqrt(noise_var) * np.random.randn(len(x))
#### Show signal & its global spectrum #######################################
axf = np.abs(rfft(x))
plot(xo)
scat(xo, s=8, show=1)
plot(x)
scat(x, s=8, show=1)
plot(axf, show=1)
#%%# Synchrosqueeze ##########################################################
kw = dict(wavelet=('morlet', {'mu': 4.5}), nv=32, scales='log')
Tx, *_ = ssq_cwt(x, t=ts, **kw)
Wx, *_ = cwt(x, t=ts, **kw)
#%%# Visualize ###############################################################
pkw = dict(abs=1, w=.86, h=.9, aspect='auto', cmap='bone')
_Tx = np.pad(Tx, [[4, 4]]) # improve display of top- & bottom-most freqs
imshow(Wx, **pkw)
imshow(np.flipud(_Tx), norm=(0, 2e-1), **pkw)
#%%# Estimate inversion ridge ###############################################
bw, slope, offset = .035, .45, .45
Cs, freqband = lin_band(Tx, slope, offset, bw, norm=(0, 2e-1))
#%%###########################################################################
xrec = issq_cwt(Tx, kw['wavelet'], Cs, freqband)[0]
plot(xo)
plot(xrec, show=1)
axof = np.abs(rfft(xo))
axrecf = np.abs(rfft(xrec))
return np.linspace(min, max, N, False)
def cos_f(freqs, N=128, phi=0):
return np.concatenate([
np.cos(2 * np.pi * f * (_t(i, i + 1, N) + phi))
for i, f in enumerate(freqs)
])
#%%## Configure, compute, plot ###############################################
wavelet = ('morlet', {'mu': 5})
f, N = 12, 512
x = cos_f([f], N=N)
Wx, scales, *_ = cwt(x, wavelet, dt=1 / N)
#%%# Show, print max row
imshow(Wx, abs=1, yticks=scales, title="f=%d, N=%d" % (f, N), show=1)
mxidx = np.where(np.abs(Wx) == np.abs(Wx).max())[0][0]
print("Max row idx:", mxidx, flush=True)
#%%# Plot around max row
idxs = slice(mxidx - 30, mxidx + 20)
Wxz = Wx[idxs]
imshow(Wxz, abs=1, title="abs(CWT), zoomed", show=0)
plt.axhline(30, color='r')
plt.show()
#%%## Animate rows ###########################################################
return np.linspace(min, max, N, False)
def cos_f(freqs, N=128, phi=0):
return np.concatenate([
np.cos(2 * np.pi * f * (_t(i, i + 1, N) + phi))
for i, f in enumerate(freqs)
])
#%%## Configure, compute, plot ###############################################
wavelet = ('morlet', {'mu': 5})
f, N = 12, 512
x = cos_f([f], N=N)
Wx, scales, *_ = cwt(x, wavelet, fs=N)
#%%# Show, print max row
imshow(Wx,
abs=1,
yticks=scales,
title="f=%d, N=%d" % (f, N),
show=1,
cmap='bone')
mxidx = np.where(np.abs(Wx) == np.abs(Wx).max())[0][0]
print("Max row idx:", mxidx, flush=True)
#%%# Plot aroundI max row
idxs = slice(mxidx - 30, mxidx + 20)
Wxz = Wx[idxs]
imshow(Wxz, abs=1, title="abs(CWT), zoomed", show=0, cmap='bone')
def test_ridge_extraction():
"""For @jit coverage."""
Wx, scales = cwt(np.random.randn(128))
_ = extract_ridges(Wx, scales, transform='cwt', parallel=False)
_ = extract_ridges(Wx, scales, transform='cwt', parallel=True)
"""
if __name__ != '__main__':
raise Exception("ran example file as non-main")
import numpy as np
from ssqueezepy import cwt, TestSignals
from ssqueezepy.visuals import viz_cwt_higher_order, viz_gmw_orders
#%%# CWT with higher-order GMWs #############################################
N = 1024
order = 2
tsigs = TestSignals()
x, t = tsigs.par_lchirp(N=N)
x += x[::-1]
for noise in (False, True):
if noise:
x += np.random.randn(len(x))
Wx_k, scales = cwt(x, 'gmw', order=range(order + 1), average=False)
viz_cwt_higher_order(Wx_k, scales, 'gmw')
print("=" * 80)
#%%# Higher-order GMWs #######################################################
gamma, beta, norm = 3, 60, 'bandpass'
n_orders = 3
scale = 5
viz_gmw_orders(N, n_orders, scale, gamma, beta, norm)
