def test_ssq_cwt():
os.environ['SSQ_GPU'] = '0' # in case concurrent tests set it to '1'
np.random.seed(5)
x = np.random.randn(64)
for wavelet in ('morlet', ('morlet', {'mu': 20}), 'bump'):
Tx, *_ = ssq_cwt(x, wavelet)
issq_cwt(Tx, wavelet)
kw = dict(x=x, wavelet='morlet')
params = dict(
squeezing=('lebesgue', ),
scales=('linear', 'log:minimal', 'linear:naive',
np.power(2**(1 / 8), np.arange(1, 32))),
difftype=('phase', 'numeric'),
padtype=('zero', 'replicate'),
maprange=('maximal', 'energy', 'peak', (1, 32)),
)
for name in params:
for value in params[name]:
errored = True
try:
if name == 'maprange' and value in ('maximal', (1, 32)):
_ = ssq_cwt(**kw, **{name: value}, scales='log', get_w=1)
else:
_ = ssq_cwt(**kw, **{name: value}, get_w=1)
errored = False
finally:
if errored:
print(f"\n{name}={value} failed\n")
_ = ssq_cwt(x, wavelet, fs=2, difftype='numeric', difforder=2, get_w=1)
_ = ssq_cwt(x, wavelet, fs=2, difftype='numeric', difforder=1, get_w=1)
def test_ssq_cwt():
x = np.random.randn(64)
for wavelet in ('morlet', ('morlet', {'mu': 20}), 'bump'):
Tx, *_ = ssq_cwt(x, wavelet)
issq_cwt(Tx, wavelet)
kw = dict(x=x, wavelet='morlet')
params = dict(
squeezing=('lebesgue', ),
scales=('linear', 'log:minimal', 'linear:naive',
np.power(2**(1 / 16), np.arange(1, 32))),
difftype=('phase', 'numeric'),
padtype=('zero', 'replicate'),
mapkind=('energy', 'peak'),
)
for name in params:
for value in params[name]:
try:
_ = ssq_cwt(**kw, **{name: value})
except Exception as e:
raise Exception(f"{name}={value} failed with:\n{e}")
_ = ssq_cwt(x, wavelet, fs=2, difftype='numeric', difforder=2)
_ = ssq_cwt(x, wavelet, fs=2, difftype='numeric', difforder=1)
def test_component_inversion():
def echirp(N):
t = np.linspace(0, 10, N, False)
return np.cos(2 * np.pi * np.exp(t / 3)), t
N = 2048
noise_var = 6
x, ts = echirp(N)
x *= (1 + .3 * cos_f([1], N)) # amplitude modulation
xo = x.copy()
np.random.seed(4)
x += np.sqrt(noise_var) * np.random.randn(len(x))
wavelet = ('morlet', {'mu': 4.5})
Tx, *_ = ssq_cwt(x, wavelet, scales='log:maximal', nv=32, t=ts)
# hand-coded, subject to failure
bw, slope, offset = .035, .45, .45
Cs, freqband = lin_band(Tx, slope, offset, bw, norm=(0, 2e-1))
xrec = issq_cwt(Tx, wavelet, Cs, freqband)[0]
axof = np.abs(np.fft.rfft(xo))
axrecf = np.abs(np.fft.rfft(xrec))
err_sig = mad_rms(xo, xrec)
err_spc = mad_rms(axof, axrecf)
print("signal MAD/RMS: %.6f" % err_sig)
print("spectrum MAD/RMS: %.6f" % err_spc)
assert err_sig <= .42, f"{err_sig} > .42"
assert err_spc <= .14, f"{err_spc} > .14"
def test_ssq_cwt():
errs = []
for fn in test_fns:
x, ts = fn(2048)
for scales in ('log', 'linear'):
# 'linear' default can't handle low frequencies for large N
if scales == 'linear' and fn.__name__ == 'fast_transitions':
continue
Tx, *_ = ssq_cwt(x, wavelet, scales=scales, nv=32, t=ts)
xrec = issq_cwt(Tx, wavelet)
errs.append(round(mad_rms(x, xrec), 5))
assert errs[-1] < th, (errs[-1], fn.__name__, scales)
print("\nssq_cwt PASSED\nerrs:", ', '.join(map(str, errs)))
def test_cwt_log_piecewise():
x, ts = echirp(1024)
wavelet = 'gmw'
Tx, Wx, ssq_freqs, scales, *_ = ssq_cwt(x,
wavelet,
scales='log-piecewise',
t=ts,
preserve_transform=True)
xrec_ssq_cwt = issq_cwt(Tx, 'gmw')
xrec_cwt = icwt(Wx, wavelet, scales=scales)
err_ssq_cwt = round(mad_rms(x, xrec_ssq_cwt), 5)
err_cwt = round(mad_rms(x, xrec_cwt), 5)
assert err_ssq_cwt < .02, err_ssq_cwt
assert err_cwt < .02, err_cwt
def test_ssq_cwt():
errs = []
for fn in test_fns:
x, ts = fn(2048)
for scales in ('log', 'log-piecewise', 'linear'):
if fn.__name__ == 'low_freqs':
if scales == 'linear':
# 'linear' default can't handle low frequencies for large N
# 'log-piecewise' maps it too sparsely
continue
else:
scales = f'{scales}:maximal'
Tx, *_ = ssq_cwt(x, wavelet, scales=scales, nv=32, t=ts)
xrec = issq_cwt(Tx, wavelet)
errs.append(round(mad_rms(x, xrec), 5))
title = "abs(SSQ_CWT) | {}, scales='{}'".format(
fn.__qualname__, scales)
_maybe_viz(Tx, x, xrec, title, errs[-1])
assert errs[-1] < th, (errs[-1], fn.__name__, scales)
print("\nssq_cwt PASSED\nerrs:", ', '.join(map(str, errs)))
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))
plot(axof)
plot(axrecf, show=1)
print("signal MAD/RMS: %.6f" % mad_rms(xo, xrec))
print("spectrum MAD/RMS: %.6f" % mad_rms(axof, axrecf))
