"""Tests for the wavelets module. The DWT claims are pinned to the two properties that make it useful: it is an orthonormal change of basis (energy-preserving) and perfectly invertible. The CWT claims are pinned to physics: a Morlet scalogram peaks at the scale whose mapped frequency matches a pure tone. Denoising is pinned to an SNR gain. """ import numpy as np import pytest from wavelets import ( wavelet_filters, wavedec, waverec, cwt, cwt_frequencies, morlet, ricker, ricker_cwt, denoise, ) WAVELETS = ["haar", "db4"] # --- filter relations ------------------------------------------------------- @pytest.mark.parametrize("name", WAVELETS) def test_lowpass_is_orthonormal(name): dec_lo, dec_hi = wavelet_filters(name) # unit energy and sum sqrt(2): the defining normalisation of an orthonormal # scaling filter assert np.isclose(np.sum(dec_lo ** 2), 1.0) assert np.isclose(np.sum(dec_lo), np.sqrt(2.0)) @pytest.mark.parametrize("name", WAVELETS) def test_highpass_is_orthogonal_to_lowpass(name): dec_lo, dec_hi = wavelet_filters(name) # low- and high-pass are orthogonal under even shifts (the QMF condition); # the zero-shift inner product is the basic case assert np.isclose(np.dot(dec_lo, dec_hi), 0.0) assert np.isclose(np.sum(dec_hi ** 2), 1.0) def test_db4_has_two_vanishing_moments(): # the high-pass annihilates constant and linear signals (zeroth and first # moments), the meaning of "two vanishing moments" _, dec_hi = wavelet_filters("db4") k = np.arange(len(dec_hi)) assert np.isclose(np.sum(dec_hi), 0.0) assert np.isclose(np.sum(k * dec_hi), 0.0) def test_unknown_wavelet_rejected(): with pytest.raises(ValueError): wavelet_filters("not-a-wavelet") # --- perfect reconstruction ------------------------------------------------- @pytest.mark.parametrize("name", WAVELETS) @pytest.mark.parametrize("level", [1, 2, 3]) def test_perfect_reconstruction(name, level): rng = np.random.default_rng(0) x = rng.standard_normal(64) # divisible by 2**3 coeffs = wavedec(x, name, level) rec = waverec(coeffs, name) assert np.allclose(rec, x, atol=1e-12) @pytest.mark.parametrize("name", WAVELETS) def test_transform_preserves_energy(name): # an orthonormal transform conserves the L2 norm (Parseval): the energy in # the coefficients equals the energy in the signal rng = np.random.default_rng(1) x = rng.standard_normal(128) coeffs = wavedec(x, name, level=3) coeff_energy = sum(np.sum(c ** 2) for c in coeffs) assert np.isclose(coeff_energy, np.sum(x ** 2)) def test_wavedec_lengths_halve_each_level(): x = np.zeros(64) coeffs = wavedec(x, "haar", level=3) # [cA_3, cD_3, cD_2, cD_1] -> lengths [8, 8, 16, 32] assert [len(c) for c in coeffs] == [8, 8, 16, 32] def test_haar_single_level_is_average_and_difference(): x = np.array([1.0, 3.0, 2.0, 6.0]) cA, cD = wavedec(x, "haar", level=1) # Haar approximation = (a+b)/sqrt2, detail = (a-b)/sqrt2 on each pair (the # sign of the detail is a convention; this QMF gives first-minus-second) assert np.allclose(cA, [(1 + 3) / np.sqrt(2), (2 + 6) / np.sqrt(2)]) assert np.allclose(cD, [(1 - 3) / np.sqrt(2), (2 - 6) / np.sqrt(2)]) def test_length_not_divisible_rejected(): with pytest.raises(ValueError): wavedec(np.zeros(30), "haar", level=2) # 30 not divisible by 4 # --- energy compaction (the wavelet's reason to exist) ---------------------- def test_localized_event_compacts_into_few_coefficients(): # a single localized bump: the DWT puts almost all its energy in a handful # of coefficients, whereas the DFT smears it across the whole spectrum. N = 256 x = np.zeros(N) x[100:104] = [2.0, 4.0, -4.0, -2.0] # the localized event coeffs = wavedec(x, "db4", level=4) flat = np.concatenate(coeffs) total_e = np.sum(flat ** 2) # how many of the largest coefficients hold 99% of the energy? order = np.sort(flat ** 2)[::-1] n_wave = np.searchsorted(np.cumsum(order), 0.99 * total_e) + 1 # the DFT for comparison X = np.abs(np.fft.rfft(x)) ** 2 n_fft = np.searchsorted(np.cumsum(np.sort(X)[::-1]), 0.99 * np.sum(X)) + 1 assert n_wave < n_fft # wavelet is far more compact assert n_wave <= 16 # --- continuous wavelet transform ------------------------------------------- def test_morlet_has_unit_energy(): w = morlet(512, scale=16.0) # the 1/sqrt(scale) normalisation makes the discrete energy ~1 (continuous) assert np.isclose(np.sum(np.abs(w) ** 2), 1.0, rtol=0.05) def test_cwt_peaks_at_the_tone_frequency(): fs = 1000.0 f0 = 50.0 t = np.arange(2048) / fs x = np.sin(2 * np.pi * f0 * t) scales = np.geomspace(4, 200, 60) power = np.abs(cwt(x, scales)) ** 2 # average over time, find the scale with the most energy, map it to Hz best = scales[np.argmax(power.mean(axis=1))] f_est = float(cwt_frequencies(best, fs)) assert abs(f_est - f0) < 0.15 * f0 # within 15% of the tone def test_cwt_localizes_a_transient_in_time(): fs = 1000.0 N = 2048 x = np.zeros(N) x[1000] = 1.0 # an impulse at sample 1000 scales = np.geomspace(4, 64, 40) power = np.abs(cwt(x, scales)) ** 2 # the high-energy column sits at the impulse, not smeared across all time peak_time = np.argmax(power.sum(axis=0)) assert abs(peak_time - 1000) < 20 def test_cwt_frequencies_inverse_of_scale(): # frequency is inversely proportional to scale: doubling the scale halves f fs = 8000.0 f = cwt_frequencies(np.array([10.0, 20.0]), fs) assert np.isclose(f[0], 2 * f[1]) # --- Ricker wavelet as a matched filter ------------------------------------- def test_ricker_is_zero_mean_and_symmetric(): w = ricker(201, scale=8.0) # admissibility: a wavelet integrates to zero; the discrete sum is ~0 assert abs(np.sum(w)) < 1e-3 * np.sum(np.abs(w)) # the Ricker is symmetric about its centre assert np.allclose(w, w[::-1]) def test_ricker_cwt_detects_r_peaks(): from scipy.signal import find_peaks fs = 500.0 N = int(6 * fs) t = np.arange(N) / fs # R-peaks as narrow Gaussians at 72 BPM (0.833 s apart), first at 0.5 s beat_times = np.arange(0.5, 6.0, 60.0 / 72.0) sigma = 0.010 # ~10 ms QRS-like width ecg = np.zeros(N) for tb in beat_times: ecg += np.exp(-((t - tb) ** 2) / (2 * sigma ** 2)) # Ricker scale matched to the spike width acts as a QRS matched filter scale = sigma * fs # 5 samples response = np.abs(ricker_cwt(ecg, [scale])[0]) peaks, _ = find_peaks(response, height=0.5 * response.max(), distance=int(0.3 * fs)) # 300 ms refractory assert len(peaks) == len(beat_times) errs = [np.min(np.abs(beat_times - p / fs)) for p in peaks] assert max(errs) < 0.015 # each R-peak within 15 ms of truth # --- denoising -------------------------------------------------------------- def test_denoise_improves_snr_on_a_piecewise_signal(): rng = np.random.default_rng(2) N = 1024 t = np.linspace(0, 1, N, endpoint=False) clean = np.where(t < 0.5, 1.0, -1.0) # a step (piecewise-constant) clean += np.where((t > 0.2) & (t < 0.3), 1.5, 0.0) noisy = clean + 0.4 * rng.standard_normal(N) out = denoise(noisy, "db4", level=4) def snr(est): return 10 * np.log10(np.sum(clean ** 2) / np.sum((est - clean) ** 2)) assert snr(out) > snr(noisy) + 3 # at least 3 dB cleaner def test_denoise_leaves_a_clean_signal_nearly_untouched(): # with no noise the universal threshold is ~0 (MAD of tiny details), so a # smooth signal survives essentially intact t = np.linspace(0, 1, 512, endpoint=False) clean = np.sin(2 * np.pi * 3 * t) out = denoise(clean, "db4", level=3) assert np.sqrt(np.mean((out - clean) ** 2)) < 0.05 def test_denoise_rejects_bad_mode(): with pytest.raises(ValueError): denoise(np.zeros(64), "haar", level=2, mode="medium")