2"""Generate the figures for spectrogram-phase-proposal.md.
4Run from anywhere: python3 docs/make_figures.py
5Figures are written to docs/figures/.
6"""
7from pathlib import Path
9import matplotlib
11matplotlib.use("Agg")
12import matplotlib.pyplot as plt
13import numpy as np
14from scipy.ndimage import gaussian_filter1d
15from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
17OUT = Path(__file__).parent / "figures"
18OUT.mkdir(exist_ok=True)
19rng = np.random.default_rng(3)
21FS = 16000
22T = 0.8
23N = int(FS * T)
24t = np.arange(N) / FS
25F0 = 220
26CLICK_T = 0.15
27TONE_T0, TONE_T1 = 0.35, 0.7
29COL_A = "#1f5fa8"
30COL_B = "#d9641e"
31GRAY = "#888888"
33plt.rcParams.update({"font.size": 9, "axes.titlesize": 10, "figure.dpi": 130})
36# ----------------------------------------------------------------------------
37# Test signal: a click followed by a harmonic tone
38# ----------------------------------------------------------------------------
39def test_signal():
40 x = np.zeros(N)
41 i0 = int(CLICK_T * FS)
42 x[i0 : i0 + 4] += np.hanning(6)[1:-1]
43 seg = (t >= TONE_T0) & (t < TONE_T1)
44 tt = t[seg] - TONE_T0
45 env = np.ones_like(tt)
46 fade = int(0.01 * FS)
47 env[:fade] = 0.5 - 0.5 * np.cos(np.pi * np.arange(fade) / fade)
48 env[-fade:] = env[:fade][::-1]
49 tone = sum(np.sin(2 * np.pi * k * F0 * tt) / k for k in range(1, 9))
50 x[seg] += 0.25 * env * tone
51 return x
54x = test_signal()
57# ----------------------------------------------------------------------------
58# Gaussian filter bank (the STFT with a hop of one sample)
59# ----------------------------------------------------------------------------
60class GaussBank:
61 """Complex Gaussian bandpass filters whose responses sum to one on f >= 0.
63 Band k has center f_k = k * df and response W(f - f_k) / S(f), where
64 W is a Gaussian of width sigma_f = 1 / (2 pi sigma_t) and
65 S(f) = sum_k W(f - f_k). The subband z_k = x * h_k is the STFT at bin k
66 evaluated at every sample. Because the responses sum to one,
67 x = Re sum_k z_k exactly (positive frequencies are doubled first).
68 """
70 def __init__(self, n, fs, sigma_t, spacing=1.0):
71 self.n, self.fs, self.sigma_t = n, fs, sigma_t
72 pad = int(8 * sigma_t * fs)
73 self.L = 1 << int(np.ceil(np.log2(n + 2 * pad)))
74 self.f = np.fft.fftfreq(self.L, 1 / fs)
75 self.sigma_f = 1 / (2 * np.pi * sigma_t)
76 self.df = spacing * self.sigma_f
77 self.centers = np.arange(0, fs / 2 + self.df / 2, self.df)
78 S = np.zeros(self.L)
79 for fc in self.centers:
80 S += self._W(fc)
81 self.S = np.where(S > 1e-300, S, np.inf) # far outside the covered range the sum underflows; those bins are never used
83 def _W(self, fc):
84 return np.exp(-((self.f - fc) ** 2) / (2 * self.sigma_f**2))
86 def response(self, k):
87 return self._W(self.centers[k]) / self.S
89 def analytic_spectrum(self, x):
90 X = np.fft.fft(x, self.L)
91 X[self.f < 0] = 0
92 X[self.f > 0] *= 2
93 return X
95 def subband(self, Xa, k):
96 return np.fft.ifft(Xa * self.response(k))[: self.n]
98 def apply(self, x, phase=None):
99 """y = Re sum_k exp(i phi_k(t)) z_k(t). phase(k) returns a scalar or an array over time."""
100 Xa = self.analytic_spectrum(x)
101 y = np.zeros(self.n)
102 for k in range(len(self.centers)):
103 z = self.subband(Xa, k)
104 if phase is not None:
105 z = z * np.exp(1j * phase(k))
106 y += z.real
107 return y
109 def spectrogram(self, x, hop, fmax):
110 Xa = self.analytic_spectrum(x)
111 ks = [k for k, fc in enumerate(self.centers) if fc <= fmax]
112 idx = np.arange(0, self.n, hop)
113 M = np.zeros((len(ks), len(idx)))
114 for i, k in enumerate(ks):
115 M[i] = np.abs(self.subband(Xa, k)[idx])
116 return M, self.centers[ks], idx / self.fs
118 def effective_filter(self, phi):
119 """Frequency response of the LTI filter obtained with constant per-band phases phi_k."""
120 G = np.zeros(self.L, dtype=complex)
121 for k, p in enumerate(phi):
122 G += np.exp(1j * p) * self.response(k)
123 return G
126def smooth_random_phases(nbands, smooth_bands, rng):
127 """Uniform random phases, smoothed across bands with a Gaussian of the given width, then rescaled."""
128 raw = rng.uniform(-np.pi, np.pi, nbands)
129 if smooth_bands <= 0:
130 return raw
131 sm = gaussian_filter1d(raw, smooth_bands, mode="reflect")
132 return sm * raw.std() / sm.std()
135def db(M):
136 return 20 * np.log10(np.maximum(M, 1e-12) / M.max())
139def show_spec(ax, M, freqs, times, title, vmin=-60):
140 ax.imshow(db(M), origin="lower", aspect="auto", extent=[times[0], times[-1], freqs[0], freqs[-1]], vmin=vmin, vmax=0, cmap="magma")
141 ax.set_title(title)
142 ax.set_xlabel("time (s)")
143 ax.set_ylabel("frequency (Hz)")
146# ----------------------------------------------------------------------------
147# Figure 1: the STFT as sliding windows, and the spectrogram
148# ----------------------------------------------------------------------------
149def fig_stft_basics():
150 sigma = 0.005
151 fig, axs = plt.subplots(2, 1, figsize=(8, 5.2), gridspec_kw={"height_ratios": [1, 1.4]})
152 ax = axs[0]
153 ax.plot(t, x, color="k", lw=0.6)
154 hop = 0.02
155 for i, c in enumerate(np.arange(0.06, 0.56, hop)):
156 w = np.exp(-((t - c) ** 2) / (2 * sigma**2))
157 ax.plot(t, 1.05 * w, color=COL_A if i % 2 else COL_B, lw=0.9, alpha=0.8)
158 ax.annotate("hop", xy=(0.30, 1.12), xytext=(0.30, 1.35), ha="center", arrowprops=dict(arrowstyle="-[, widthB=1.6", lw=0.8))
159 ax.set_xlim(0, T)
160 ax.set_ylim(-1.1, 1.5)
161 ax.set_ylabel("amplitude")
162 ax.set_title("Test signal (a click, then a harmonic tone) with Gaussian analysis windows, sigma = 5 ms")
163 bank = GaussBank(N, FS, sigma)
164 M, freqs, times = bank.spectrogram(x, 32, 4000)
165 show_spec(axs[1], M, freqs, times, "Spectrogram: magnitude of each windowed Fourier coefficient (dB)")
166 fig.tight_layout()
167 fig.savefig(OUT / "01-stft-basics.png")
168 plt.close(fig)
171# ----------------------------------------------------------------------------
172# Figure 2: the filter bank view
173# ----------------------------------------------------------------------------
174def fig_filterbank():
175 sigma = 0.005
176 bank = GaussBank(N, FS, sigma)
177 fig = plt.figure(figsize=(8, 7.5))
178 gs = fig.add_gridspec(4, 1, height_ratios=[1, 1, 1, 1])
180 ax = fig.add_subplot(gs[0])
181 pos = bank.f >= 0
182 f = bank.f[pos]
183 total = np.zeros(pos.sum())
184 for k in range(len(bank.centers)):
185 r = bank.response(k)[pos]
186 total += r
187 if bank.centers[k] <= 1500:
188 ax.plot(f, r, lw=0.8, color=COL_A if k % 2 else COL_B, alpha=0.8)
189 ax.plot(f, total, "k--", lw=1, label="sum over all bands")
190 ax.set_xlim(0, 1500)
191 ax.set_ylim(0, 1.15)
192 ax.set_xlabel("frequency (Hz)")
193 ax.set_ylabel("gain")
194 ax.legend(loc="lower right")
195 ax.set_title("Each STFT bin is a bandpass filter. The filters overlap and sum to one.")
197 Xa = bank.analytic_spectrum(x)
198 picks = [int(round(F0 * m / bank.df)) for m in (1, 3)] + [int(round(1200 / bank.df))]
199 for row, k in enumerate(picks, start=1):
200 ax = fig.add_subplot(gs[row])
201 z = bank.subband(Xa, k)
202 ax.plot(t, z.real, color=COL_A, lw=0.6, label="Re z_k(t)")
203 ax.plot(t, np.abs(z), color=COL_B, lw=1.0, label="|z_k(t)| (envelope)")
204 ax.set_xlim(0, T)
205 ax.set_ylabel("amplitude")
206 ax.set_title(f"Subband k = {k}, center {bank.centers[k]:.0f} Hz", loc="left")
207 if row == 1:
208 ax.legend(loc="upper right", ncol=2)
209 ax.set_xlabel("time (s)")
210 y = bank.apply(x)
211 fig.text(0.5, 0.005, f"Adding all subbands back together reproduces the signal: max |x - Re sum_k z_k| = {np.max(np.abs(x - y)):.1e}", ha="center", fontsize=9)
212 fig.tight_layout(rect=(0, 0.02, 1, 1))
213 fig.savefig(OUT / "02-filterbank.png")
214 plt.close(fig)
217# ----------------------------------------------------------------------------
218# Figure 3: rotating subband phases, for three window lengths
219# ----------------------------------------------------------------------------
220SMOOTH = 6 # bands
223def rotated_versions(sigmas, smooth_bands=SMOOTH, seed=5):
224 out = {}
225 for s in sigmas:
226 bank = GaussBank(N, FS, s)
227 phi = smooth_random_phases(len(bank.centers), smooth_bands, np.random.default_rng(seed))
228 y = bank.apply(x, lambda k: phi[k])
229 out[s] = (bank, phi, y)
230 return out
233ROT = rotated_versions([0.0005, 0.005, 0.02])
236def fig_phase_rotation():
237 fig, axs = plt.subplots(3, 2, figsize=(8, 7), gridspec_kw={"width_ratios": [1.4, 1]})
238 for row, (s, (bank, phi, y)) in enumerate(ROT.items()):
239 ax = axs[row, 0]
240 sel = (t > CLICK_T - 0.05) & (t < CLICK_T + 0.08)
241 ax.plot(t[sel] * 1000, x[sel], color=GRAY, lw=0.8, label="original")
242 ax.plot(t[sel] * 1000, y[sel], color=COL_B, lw=0.8, label="phases rotated")
243 ax.set_title(f"window sigma = {s*1000:g} ms: the click", loc="left")
244 ax.set_ylabel("amplitude")
245 if row == 0:
246 ax.legend(loc="upper right")
247 if row == 2:
248 ax.set_xlabel("time (ms)")
249 ax = axs[row, 1]
250 sel = (t > 0.5) & (t < 0.5 + 3 / F0)
251 ax.plot(t[sel] * 1000, x[sel], color=GRAY, lw=0.8)
252 ax.plot(t[sel] * 1000, y[sel], color=COL_B, lw=0.8)
253 ax.set_title("same map, steady tone (3 periods)", loc="left")
254 if row == 2:
255 ax.set_xlabel("time (ms)")
256 fig.suptitle("y = Re sum_k exp(i phi_k) z_k(t) with smooth random phi_k", fontsize=10)
257 fig.tight_layout(rect=(0, 0, 1, 0.96))
258 fig.savefig(OUT / "03-phase-rotation.png")
259 plt.close(fig)
262# ----------------------------------------------------------------------------
263# Figure 4: spectrograms before and after, at two analysis resolutions
264# ----------------------------------------------------------------------------
265def fig_spectrograms():
266 _, _, y = ROT[0.02]
267 fig, axs = plt.subplots(2, 2, figsize=(8, 5.6))
268 for col, s in enumerate([0.02, 0.002]):
269 bank = GaussBank(N, FS, s)
270 for row, (sig, name) in enumerate([(x, "original"), (y, "rotated (20 ms bank)")]):
271 M, freqs, times = bank.spectrogram(sig, 32, 4000)
272 show_spec(axs[row, col], M, freqs, times, f"{name}, analysed at sigma = {s*1000:g} ms")
273 fig.tight_layout()
274 fig.savefig(OUT / "04-spectrograms.png")
275 plt.close(fig)
278# ----------------------------------------------------------------------------
279# Figure 5: the equivalent all-pass filter and its group delay
280# ----------------------------------------------------------------------------
281def fig_group_delay():
282 fig, axs = plt.subplots(3, 1, figsize=(8, 7), sharex=True)
283 for s, (bank, phi, _) in ROT.items():
284 G = bank.effective_filter(phi)
285 pos = (bank.f > 0) & (bank.f <= 4000)
286 f = bank.f[pos]
287 ph = np.unwrap(np.angle(G[pos]))
288 tau = -np.gradient(ph, f) / (2 * np.pi)
289 lab = f"sigma = {s*1000:g} ms"
290 axs[0].plot(f, 20 * np.log10(np.abs(G[pos])), lw=0.8, label=lab)
291 axs[1].plot(f, ph, lw=0.8, label=lab)
292 axs[2].plot(f, tau * 1000, lw=0.8, label=lab)
293 axs[0].set_ylabel("magnitude (dB)")
294 axs[0].set_ylim(-3, 1)
295 axs[0].set_title("Constant per-band phases act as a filter G(f): magnitude (top), phase (middle), group delay (bottom)")
296 axs[0].legend(loc="lower left", ncol=3, fontsize=8)
297 axs[1].set_ylabel("phase (rad)")
298 axs[2].axhspan(-1.5, 1.5, color="green", alpha=0.12)
299 axs[2].text(0.99, 0.95, "green band: roughly the audibility threshold for group delay peaks", ha="right", va="top", transform=axs[2].transAxes, fontsize=8, color="green")
300 axs[2].set_ylabel("group delay (ms)")
301 axs[2].set_xlabel("frequency (Hz)")
302 axs[2].set_xlim(0, 4000)
303 fig.tight_layout()
304 fig.savefig(OUT / "05-group-delay.png")
305 plt.close(fig)
308# ----------------------------------------------------------------------------
309# Figure 6: why the phases must vary smoothly across neighbouring bands
310# ----------------------------------------------------------------------------
311def fig_notches():
312 bank = GaussBank(N, FS, 0.005)
313 fig, axs = plt.subplots(2, 2, figsize=(8, 5), sharex="col")
314 for col, (sb, name) in enumerate([(0, "independent random phase per band"), (SMOOTH, f"random phases smoothed over {SMOOTH} bands")]):
315 phi = smooth_random_phases(len(bank.centers), sb, np.random.default_rng(5))
316 G = bank.effective_filter(phi)
317 pos = (bank.f > 0) & (bank.f <= 2000)
318 f = bank.f[pos]
319 ax = axs[0, col]
320 ax.plot(bank.centers, np.angle(np.exp(1j * phi)), ".", ms=3, color=COL_B)
321 ax.set_ylabel("phi_k (rad)")
322 ax.set_title(name)
323 ax.set_xlim(0, 2000)
324 ax = axs[1, col]
325 ax.plot(f, 20 * np.log10(np.abs(G[pos])), color=COL_A, lw=0.8)
326 ax.set_ylim(-30, 3)
327 ax.set_ylabel("|G(f)| (dB)")
328 ax.set_xlabel("frequency (Hz)")
329 for m in range(1, 9):
330 ax.axvline(m * F0, color=GRAY, lw=0.5, ls=":")
331 axs[1, 0].set_title("|G(f)|: notches where neighbours disagree", fontsize=9)
332 axs[1, 1].set_title("|G(f)|: nearly flat (dotted: harmonics of the tone)", fontsize=9)
333 fig.tight_layout()
334 fig.savefig(OUT / "06-notches.png")
335 plt.close(fig)
338# ----------------------------------------------------------------------------
339# Figure 7: hop-based STFT, per-cell random phase, inverse STFT
340# ----------------------------------------------------------------------------
341def stft(x, n, hop):
342 w = np.hanning(n + 1)[:-1]
343 frames = [x[i : i + n] * w for i in range(0, len(x) - n + 1, hop)]
344 return np.fft.rfft(np.array(frames), axis=1), w
347def istft(S, w, hop, length):
348 n = len(w)
349 y = np.zeros(length)
350 norm = np.zeros(length)
351 for m, spec in enumerate(S):
352 i = m * hop
353 y[i : i + n] += np.fft.irfft(spec, n) * w
354 norm[i : i + n] += w * w
355 return y / np.maximum(norm, 1e-9)
358def fig_inconsistency():
359 n, hop = 512, 128
360 S, w = stft(x, n, hop)
361 freqs = np.fft.rfftfreq(n, 1 / FS)
362 keep = freqs <= 4000
363 times = np.arange(S.shape[0]) * hop / FS
365 S_cell = S * np.exp(1j * rng.uniform(-np.pi, np.pi, S.shape))
366 y_cell = istft(S_cell, w, hop, N)
367 phi_band = smooth_random_phases(S.shape[1], SMOOTH, np.random.default_rng(5))
368 S_band = S * np.exp(1j * phi_band)[None, :]
369 y_band = istft(S_band, w, hop, N)
371 fig, axs = plt.subplots(2, 3, figsize=(9.5, 5.4), gridspec_kw={"height_ratios": [1, 1.6]})
372 for col, (sig, title) in enumerate([(x, "original"), (y_cell, "after random phase in every cell"), (y_band, "after one phase per frequency bin")]):
373 axs[0, col].plot(t, sig, color="k", lw=0.5)
374 axs[0, col].set_xlim(0, T)
375 axs[0, col].set_ylim(-1.1, 1.1)
376 axs[0, col].set_title(title)
377 M = np.abs(stft(sig, n, hop)[0])[:, keep].T
378 show_spec(axs[1, col], M, freqs[keep], times, "spectrogram recomputed from the signal")
379 fig.suptitle("Hop-based STFT (Hann, 32 ms, 75% overlap), modify phases, inverse STFT by overlap-add", fontsize=10)
380 fig.tight_layout(rect=(0, 0, 1, 0.96))
381 fig.savefig(OUT / "07-inconsistency.png")
382 plt.close(fig)
385# ----------------------------------------------------------------------------
386# Figure 8: a time-varying phase offset is frequency modulation
387# ----------------------------------------------------------------------------
388def fig_time_varying():
389 fc = 0.5
390 delta = np.pi
391 tt = np.linspace(0, 2 / fc, 2000)
392 s = 0.5 - 0.5 * np.cos(2 * np.pi * fc * tt)
393 phi = delta * s
394 dphi = np.gradient(phi, tt)
395 fig, axs = plt.subplots(1, 3, figsize=(9.5, 2.9))
396 axs[0].plot(tt, phi, color=COL_A)
397 axs[0].set_title("phi(t): 0 to pi and back at 0.5 Hz")
398 axs[0].set_xlabel("time (s)")
399 axs[0].set_ylabel("rad")
400 axs[1].plot(tt, dphi / (2 * np.pi), color=COL_B)
401 axs[1].set_title("frequency shift phi'(t) / 2 pi")
402 axs[1].set_xlabel("time (s)")
403 axs[1].set_ylabel("Hz")
404 fcs = np.logspace(-2, 2, 100)
405 axs[2].loglog(fcs, delta * fcs / 2, color=COL_B)
406 axs[2].axhspan(0.5, 2, color="green", alpha=0.12)
407 axs[2].text(0.011, 2.4, "roughly the pitch JND near 220 Hz", fontsize=8, color="green")
408 axs[2].set_title("peak shift vs cycle rate (swing pi)")
409 axs[2].set_xlabel("cycle rate (Hz)")
410 axs[2].set_ylabel("peak shift (Hz)")
411 fig.tight_layout()
412 fig.savefig(OUT / "08-time-varying.png")
413 plt.close(fig)
416# ----------------------------------------------------------------------------
417# Figure 9: proposed processing pipeline
418# ----------------------------------------------------------------------------
419def fig_pipeline():
420 fig, ax = plt.subplots(figsize=(9.5, 3.4))
421 ax.set_xlim(0, 10)
422 ax.set_ylim(0, 3.4)
423 ax.axis("off")
425 def box(xc, yc, wdt, hgt, text, color="#eef3fa", ec="#333"):
426 ax.add_patch(FancyBboxPatch((xc - wdt / 2, yc - hgt / 2), wdt, hgt, boxstyle="round,pad=0.05", fc=color, ec=ec, lw=0.9))
427 ax.text(xc, yc, text, ha="center", va="center", fontsize=8.5)
429 def arrow(x0, y0, x1, y1):
430 ax.add_patch(FancyArrowPatch((x0, y0), (x1, y1), arrowstyle="-|>", mutation_scale=12, lw=0.9, color="#333"))
432 box(0.9, 1.7, 1.5, 1.2, "source sound x(t)\n(click, noise burst,\npluck, file, mic)")
433 box(3.1, 1.7, 1.7, 1.2, "Gaussian filter bank\nwindow sigma\n\nz_k(t) for each band k")
434 box(5.55, 2.6, 1.9, 0.75, "phase map A: phi_k = 0", color="#e5edf8", ec=COL_A)
435 box(5.55, 0.8, 1.9, 0.75, "phase map B: phi_k chosen", color="#fbeadf", ec=COL_B)
436 box(7.9, 2.6, 1.7, 0.75, "A = Re sum z_k = x", color="#e5edf8", ec=COL_A)
437 box(7.9, 0.8, 1.7, 0.75, "B = Re sum e^(i phi_k) z_k", color="#fbeadf", ec=COL_B)
438 box(9.4, 1.7, 1.1, 1.0, "blind test\nA/B, cycling")
439 arrow(1.65, 1.7, 2.25, 1.7)
440 arrow(3.95, 1.9, 4.6, 2.6)
441 arrow(3.95, 1.5, 4.6, 0.8)
442 arrow(6.5, 2.6, 7.05, 2.6)
443 arrow(6.5, 0.8, 7.05, 0.8)
444 arrow(8.75, 2.6, 9.0, 2.0)
445 arrow(8.75, 0.8, 9.0, 1.4)
446 ax.text(5.55, 1.7, "same subband envelopes |z_k(t)|,\ndifferent subband phases", ha="center", va="center", fontsize=8, color="#555")
447 fig.tight_layout()
448 fig.savefig(OUT / "09-pipeline.png")
449 plt.close(fig)
452if __name__ == "__main__":
453 fig_stft_basics()
454 fig_filterbank()
455 fig_phase_rotation()
456 fig_spectrograms()
457 fig_group_delay()
458 fig_notches()
459 fig_inconsistency()
460 fig_time_varying()
461 fig_pipeline()
462 print("wrote", sorted(p.name for p in OUT.iterdir()))