1#!/usr/bin/env python3
2"""
3Time-Frequency Analysis: STFT Spectrogram and Wavelet Transform
4================================================================
5
6Classical Fourier analysis reveals *which* frequencies are present in a
7signal, but not *when* they occur. Time-frequency analysis resolves both
8dimensions simultaneously.
9
10Short-Time Fourier Transform (STFT)
11-------------------------------------
12The STFT divides the signal into overlapping short frames, applies a window
13function to each frame, and computes the DFT:
14
15 STFT{x}(m, k) = sum_{n} x[n] * w[n - mH] * exp(-j*2*pi*k*n/N)
16
17where H is the hop size and N is the DFT length (window length).
18
19Resolution trade-off (Heisenberg uncertainty principle):
20 - Long window โ good frequency resolution, poor time resolution
21 - Short window โ good time resolution, poor frequency resolution
22
23Continuous Wavelet Transform (CWT)
24------------------------------------
25The CWT uses a scaled and translated 'mother wavelet' ฯ(t):
26
27 CWT{x}(a, b) = (1/sqrt(a)) * integral x(t) * ฯ*((t-b)/a) dt
28
29where a is the scale (~ 1/frequency) and b is the translation (time).
30
31The CWT provides adaptive resolution:
32 - High-frequency components: good time resolution (narrow ฯ)
33 - Low-frequency components: good frequency resolution (wide ฯ)
34
35This is a key advantage over the fixed-resolution STFT.
36
37Author: Educational example for Signal Processing
38License: MIT
39"""
40
41import numpy as np
42import matplotlib.pyplot as plt
43from scipy import signal
44
45
46# ============================================================================
47# SIGNAL GENERATORS
48# ============================================================================
49
50def make_chirp(fs=1000.0, duration=2.0, f0=10.0, f1=200.0):
51 """
52 Generate a linear chirp: instantaneous frequency increases linearly from
53 f0 to f1 over the specified duration.
54
55 Args:
56 fs (float): Sample rate (Hz)
57 duration (float): Signal length (seconds)
58 f0 (float): Start frequency (Hz)
59 f1 (float): End frequency (Hz)
60
61 Returns:
62 t (ndarray): Time axis
63 x (ndarray): Chirp signal
64 """
65 t = np.arange(int(fs * duration)) / fs
66 x = signal.chirp(t, f0=f0, f1=f1, t1=duration, method='linear')
67 return t, x
68
69
70def make_multicomponent(fs=1000.0, duration=2.0):
71 """
72 Construct a signal with both transient and stationary components:
73 - Stationary 80 Hz tone (full duration)
74 - Gaussian-modulated transient at t = 0.5 s (burst)
75 - Stationary 200 Hz tone in second half only
76
77 This tests whether the time-frequency method can separate
78 persistent tones from brief events.
79
80 Returns:
81 t (ndarray): Time axis
82 x (ndarray): Composite signal
83 """
84 t = np.arange(int(fs * duration)) / fs
85
86 # Persistent low-frequency tone
87 tone_80 = np.sin(2 * np.pi * 80 * t)
88
89 # Short Gaussian burst at t=0.5 s, centre frequency 300 Hz
90 burst_env = np.exp(-((t - 0.5) ** 2) / (2 * 0.01 ** 2))
91 burst = burst_env * np.sin(2 * np.pi * 300 * t)
92
93 # High-frequency tone that starts at t=1 s
94 tone_200 = np.where(t >= 1.0, np.sin(2 * np.pi * 200 * t), 0.0)
95
96 rng = np.random.default_rng(0)
97 noise = 0.05 * rng.standard_normal(len(t))
98
99 return t, tone_80 + burst + tone_200 + noise
100
101
102# ============================================================================
103# SECTION 1: STFT SPECTROGRAM AND WINDOW LENGTH TRADE-OFF
104# ============================================================================
105
106def demo_stft_resolution_tradeoff(fs, t, x):
107 """
108 Compute STFT spectrograms with three different window lengths and display
109 the time-frequency resolution trade-off for a chirp signal.
110 """
111 print("=" * 60)
112 print("SECTION 1: STFT Resolution Trade-off (chirp signal)")
113 print("=" * 60)
114
115 window_lengths = [32, 128, 512] # short, medium, long
116 labels = [
117 "Short window (32)\nGood time, poor freq",
118 "Medium window (128)\nBalanced",
119 "Long window (512)\nPoor time, good freq",
120 ]
121
122 fig, axes = plt.subplots(1, 3, figsize=(14, 5))
123 fig.suptitle("STFT Resolution Trade-off for a Linear Chirp\n"
124 "Heisenberg uncertainty: ฮt ยท ฮf โฅ 1/(4ฯ)",
125 fontsize=12, fontweight='bold')
126
127 for ax, nperseg, label in zip(axes, window_lengths, labels):
128 f, t_stft, Zxx = signal.stft(x, fs=fs, nperseg=nperseg,
129 noverlap=nperseg // 2,
130 window='hann')
131 power_db = 20 * np.log10(np.abs(Zxx) + 1e-12)
132 ax.pcolormesh(t_stft, f, power_db, vmin=-80, vmax=0,
133 cmap='inferno', shading='gouraud')
134 ax.set_ylim(0, fs / 2)
135 ax.set_xlabel("Time (s)")
136 ax.set_ylabel("Frequency (Hz)")
137 ax.set_title(label)
138 print(f" Window={nperseg:4d} "
139 f"Time res โ {nperseg/fs*1000:.0f} ms "
140 f"Freq res โ {fs/nperseg:.1f} Hz")
141
142 plt.tight_layout()
143 plt.savefig("14_stft_resolution.png", dpi=120)
144 print(" Saved: 14_stft_resolution.png")
145 plt.show()
146
147
148# ============================================================================
149# SECTION 2: CWT WITH MORLET WAVELET
150# ============================================================================
151
152def morlet_wavelet(M, w0=6.0):
153 """
154 Return a complex Morlet wavelet of length M.
155
156 The Morlet wavelet is a Gaussian-modulated complex sinusoid:
157 ฯ(t) = pi^(-1/4) * exp(j*w0*t) * exp(-t^2 / 2)
158
159 At scale a, the peak frequency is approximately w0 / (2*pi*a).
160
161 Args:
162 M (int) : Wavelet length in samples.
163 w0 (float): Angular frequency of the carrier (controls freq resolution).
164
165 Returns:
166 ndarray: Complex Morlet wavelet, normalised.
167 """
168 x = np.linspace(-w0 / 2, w0 / 2, M)
169 wavelet = np.exp(1j * w0 * x) * np.exp(-0.5 * x ** 2) * np.pi ** (-0.25)
170 return wavelet
171
172
173def demo_cwt(fs, t, x, signal_name=""):
174 """
175 Compute and display the Continuous Wavelet Transform scalogram using
176 scipy.signal.cwt with a Morlet wavelet.
177
178 scipy.signal.cwt uses the *real* part of the Morlet (ricker-like) by
179 default, so we directly call signal.cwt with morlet2 via a wrapper, or
180 equivalently compute it manually using convolution.
181 """
182 print("\n" + "=" * 60)
183 print(f"SECTION 2: Continuous Wavelet Transform [{signal_name}]")
184 print("=" * 60)
185
186 # Define scales: relate scale to pseudo-frequency via f โ w0 / (2*pi*a/fs)
187 w0 = 6.0
188 # Frequencies we want to resolve (Hz)
189 freqs_hz = np.logspace(np.log10(5), np.log10(fs / 2 - 1), 80)
190 # Corresponding scales (in samples): a = w0 * fs / (2*pi*f)
191 scales = w0 * fs / (2 * np.pi * freqs_hz)
192
193 # CWT via scipy.signal.cwt (uses ricker/morlet based on the function passed)
194 # We use scipy's built-in morlet2 wavelet
195 cwt_matrix = signal.cwt(x, signal.morlet2, scales, w=w0)
196 scalogram = np.abs(cwt_matrix)
197
198 fig, axes = plt.subplots(2, 1, figsize=(11, 8))
199 fig.suptitle(f"Continuous Wavelet Transform (Morlet) โ {signal_name}",
200 fontsize=12, fontweight='bold')
201
202 # Scalogram
203 ax = axes[0]
204 im = ax.pcolormesh(t, freqs_hz, scalogram, cmap='hot_r', shading='gouraud')
205 ax.set_yscale('log')
206 ax.set_ylabel("Pseudo-frequency (Hz)")
207 ax.set_xlabel("Time (s)")
208 ax.set_title("CWT Scalogram (log frequency axis)\n"
209 "High freq: narrow wavelet (good time res) | "
210 "Low freq: wide wavelet (good freq res)")
211 fig.colorbar(im, ax=ax, label="Magnitude")
212
213 # Morlet wavelet at two scales (visualise what the filter looks like)
214 ax2 = axes[1]
215 for scale, color in [(5, 'C0'), (50, 'C1')]:
216 M = int(10 * scale) | 1 # odd length
217 wav = morlet_wavelet(M, w0=w0)
218 t_wav = np.linspace(-M / (2 * fs), M / (2 * fs), M)
219 ax2.plot(t_wav * 1000, wav.real, color=color,
220 label=f'scale={scale} (fโ{w0*fs/(2*np.pi*scale):.0f} Hz)')
221 ax2.plot(t_wav * 1000, wav.imag, color=color, linestyle='--', alpha=0.5)
222 ax2.set_xlabel("Time (ms)")
223 ax2.set_ylabel("Amplitude")
224 ax2.set_title("Morlet Wavelets at Two Scales (solid=real, dashed=imag)\n"
225 "Smaller scale โ shorter wavelet โ better time resolution")
226 ax2.legend()
227 ax2.grid(True, alpha=0.3)
228
229 plt.tight_layout()
230 fname = f"14_cwt_{signal_name.replace(' ', '_').lower()}.png"
231 plt.savefig(fname, dpi=120)
232 print(f" Saved: {fname}")
233 plt.show()
234
235
236# ============================================================================
237# SECTION 3: SPECTROGRAM VS CWT โ SIDE-BY-SIDE COMPARISON
238# ============================================================================
239
240def demo_stft_vs_cwt(fs, t, x, signal_name=""):
241 """
242 Direct visual comparison of STFT spectrogram and CWT scalogram for the
243 same signal, illustrating their complementary strengths.
244 """
245 print("\n" + "=" * 60)
246 print(f"SECTION 3: Spectrogram vs CWT โ {signal_name}")
247 print("=" * 60)
248
249 # STFT (medium window)
250 nperseg = 128
251 f_stft, t_stft, Zxx = signal.stft(x, fs=fs, nperseg=nperseg,
252 noverlap=nperseg * 3 // 4,
253 window='hann')
254 power_stft = 20 * np.log10(np.abs(Zxx) + 1e-12)
255
256 # CWT
257 w0 = 6.0
258 freqs_hz = np.logspace(np.log10(5), np.log10(fs / 2 - 1), 80)
259 scales = w0 * fs / (2 * np.pi * freqs_hz)
260 cwt_matrix = signal.cwt(x, signal.morlet2, scales, w=w0)
261 scalogram = np.abs(cwt_matrix)
262
263 fig, axes = plt.subplots(3, 1, figsize=(12, 11))
264 fig.suptitle(f"STFT vs CWT for '{signal_name}'\n"
265 "STFT: fixed resolution | CWT: adaptive resolution",
266 fontsize=12, fontweight='bold')
267
268 # Raw signal
269 axes[0].plot(t, x, color='C0', linewidth=0.8)
270 axes[0].set_title("(a) Input Signal")
271 axes[0].set_xlabel("Time (s)")
272 axes[0].set_ylabel("Amplitude")
273 axes[0].grid(True, alpha=0.3)
274
275 # STFT spectrogram
276 axes[1].pcolormesh(t_stft, f_stft, power_stft, vmin=-80, vmax=0,
277 cmap='inferno', shading='gouraud')
278 axes[1].set_title(f"(b) STFT Spectrogram (window={nperseg} samples, "
279 f"ฮtโ{nperseg/fs*1000:.0f} ms, ฮfโ{fs/nperseg:.1f} Hz)")
280 axes[1].set_xlabel("Time (s)")
281 axes[1].set_ylabel("Frequency (Hz)")
282
283 # CWT scalogram
284 im = axes[2].pcolormesh(t, freqs_hz, scalogram, cmap='hot_r', shading='gouraud')
285 axes[2].set_yscale('log')
286 axes[2].set_title("(c) CWT Scalogram (Morlet, log frequency)\n"
287 "High-freq transients resolve sharply in time; "
288 "low-freq tones resolve sharply in frequency")
289 axes[2].set_xlabel("Time (s)")
290 axes[2].set_ylabel("Pseudo-frequency (Hz)")
291 fig.colorbar(im, ax=axes[2], label="Magnitude")
292
293 plt.tight_layout()
294 fname = f"14_stft_vs_cwt_{signal_name.replace(' ', '_').lower()}.png"
295 plt.savefig(fname, dpi=120)
296 print(f" Saved: {fname}")
297 plt.show()
298
299
300# ============================================================================
301# MAIN
302# ============================================================================
303
304if __name__ == "__main__":
305 print("Time-Frequency Analysis: STFT Spectrogram and Wavelet Transform")
306 print("=" * 60)
307
308 fs = 1000.0 # sample rate
309
310 # --- Chirp signal ---------------------------------------------------------
311 t_chirp, x_chirp = make_chirp(fs=fs, duration=2.0, f0=10.0, f1=200.0)
312 print(f"\nChirp signal: {len(t_chirp)} samples, "
313 f"f: 10โ200 Hz over {len(t_chirp)/fs:.1f} s")
314
315 demo_stft_resolution_tradeoff(fs, t_chirp, x_chirp)
316 demo_cwt(fs, t_chirp, x_chirp, signal_name="chirp")
317 demo_stft_vs_cwt(fs, t_chirp, x_chirp, signal_name="chirp")
318
319 # --- Multi-component signal -----------------------------------------------
320 t_mc, x_mc = make_multicomponent(fs=fs, duration=2.0)
321 print(f"\nMulti-component signal: {len(t_mc)} samples")
322 print(" Components: 80 Hz tone (full) + 300 Hz burst at t=0.5 s "
323 "+ 200 Hz tone (tโฅ1 s)")
324
325 demo_stft_vs_cwt(fs, t_mc, x_mc, signal_name="multicomponent")
326
327 print("\nDone. PNG files saved.")
328 print("\nKey takeaways:")
329 print(" - STFT: fixed ฮtยทฮf product; choose window for your application")
330 print(" - CWT : adaptive resolution; better for multi-scale signals")
331 print(" - Heisenberg uncertainty: you cannot have perfect ฮt AND ฮf")