Sound quilts

import os
import urllib.request

import matplotlib.pyplot as plt
import numpy as np

import sonore as so

plt.rcParams.update({"font.size": 9, "axes.titlesize": 10, "figure.dpi": 100})


def fetch(path):
    """A file from the sonore repository, by its path there: the local copy when this runs from
    the repository root, otherwise downloaded from GitHub to the same relative path."""
    if not os.path.exists(path):
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        urllib.request.urlretrieve("https://raw.githubusercontent.com/choyun1/sonore/main/" + path, path)
    return path


def finish(snd):
    """How every sound in the gallery is played: 5 ms ramps, RMS 0.1, peak at most 0.95."""
    snd = snd.ramp(5e-3).normalize(rms=0.1)
    return snd.normalize(peak=0.95) if snd.peak > 0.95 else snd

def finish_together(sounds):
    """How a recording and its quilts are played: 5 ms ramps, and one gain for all of them, the
    largest that keeps every RMS at most 0.1 and every peak at most 0.95. A quilt keeps the level
    of its recording, and one gain keeps that visible; scaling each sound on its own would let a
    single loud event in one of them turn that one down."""
    ramped = {label: snd.ramp(5e-3) for label, snd in sounds.items()}
    gain = min(min(0.1 / snd.rms, 0.95 / snd.peak) for snd in ramped.values())
    return {label: so.Sound(snd.data * gain, snd.fs) for label, snd in ramped.items()}

# A male and a female LibriSpeech reader (docs/speech/SOURCES.md), at 16 kHz.
READERS = {"Male reader": "3000_15664", "Female reader": "1462_170142"}


def without_pauses(snd, window=0.03, below_db=35, min_windows=8):
    """``snd`` with every run of at least ``min_windows`` windows more than ``below_db`` below the
    loudest window cut out, the rest joined by ``so.splice`` with 10 ms cross-fades."""
    n = int(round(window * snd.fs))
    frames = snd.data[: len(snd) // n * n, 0].reshape(-1, n)
    level = 10 * np.log10(np.mean(frames**2, axis=1) + 1e-20)  # dB per window
    quiet = np.concatenate([[False], level < level.max() - below_db, [False]])
    edges = np.flatnonzero(np.diff(quiet.astype(int)))  # starts and ends of quiet runs
    keep = np.ones(len(frames), dtype=bool)
    for start, end in zip(edges[::2], edges[1::2], strict=True):
        if end - start >= min_windows:
            keep[start:end] = False
    return so.splice(snd, np.flatnonzero(keep) * n / snd.fs, n / snd.fs, crossfade=0.01, max_shift=0)


passages = {
    label: without_pauses(so.load(fetch(f"docs/speech/librispeech_{reader}.flac")))
    for label, reader in READERS.items()
}
for label, passage in passages.items():
    print(f"{label}: {passage.duration:.1f} s without pauses")
SEGMENTS = [0.03, 0.12, 0.48, 0.96]  # s
QUILT_DUR = 6.0  # s
quilt_bank = so.cosine_filterbank(30, 20, 8000)  # the paper's bank, up to Nyquist at 16 kHz


def quilt_demo(passage):
    """The passage's first 6 s and a 6 s quilt at each segment length, with their cochleagrams
    on one color scale."""
    sounds = {"Original": so.Sound(passage.data[: int(QUILT_DUR * passage.fs)], passage.fs)}
    for seg in SEGMENTS:
        sounds[f"{seg * 1000:.0f} ms"] = so.quilt(passage, seg, QUILT_DUR, rng=1)
    sounds = finish_together(sounds)
    fig, axes = plt.subplots(len(sounds), 1, figsize=(7, 5), sharex=True, layout="constrained")
    envelopes = {label: quilt_bank.analyze(snd).envelopes() for label, snd in sounds.items()}
    top = envelopes["Original"].db[:, 1:-1].max()  # all on the original's color scale
    for ax, (label, env) in zip(axes, envelopes.items(), strict=True):
        env.plot(ax, colorbar=False, db_max=top)
        title = label if label == "Original" else f"Quilt, {label} segments"
        ax.set(title=title, xlabel="", ylabel="", xlim=(0, QUILT_DUR))
    axes[-1].set_xlabel("Time [s]")
    fig.supylabel("Frequency [Hz]", fontsize=9)
    return fig, sounds, {label: [ax] for label, ax in zip(sounds, axes, strict=True)}

texture_bank = so.cosine_filterbank(30, 20, 10000)


def texture_quilt_demo(recording, duration, segments=(0.03, 0.12, 0.96)):
    """The recording's first ``duration`` seconds and a quilt of the whole recording at each
    segment length, played at one gain, with their cochleagrams on one color scale."""
    sounds = {"Original": so.Sound(recording.data[: int(duration * recording.fs)], recording.fs)}
    for seg in segments:
        sounds[f"{seg * 1000:.0f} ms"] = so.quilt(recording, seg, duration, rng=1)
    sounds = finish_together(sounds)
    fig, axes = plt.subplots(len(sounds), 1, figsize=(7, 4), sharex=True, layout="constrained")
    envelopes = {label: texture_bank.analyze(snd).envelopes() for label, snd in sounds.items()}
    top = envelopes["Original"].db[:, 1:-1].max()  # all on the original's color scale
    for ax, (label, env) in zip(axes, envelopes.items(), strict=True):
        env.plot(ax, colorbar=False, db_max=top)
        title = label if label == "Original" else f"Quilt, {label} segments"
        ax.set(title=title, xlabel="", ylabel="", xlim=(0, duration))
    axes[-1].set_xlabel("Time [s]")
    fig.supylabel("Frequency [Hz]", fontsize=9)
    return fig, sounds, {label: [ax] for label, ax in zip(sounds, axes, strict=True)}