A richer source-filter model

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

# The two talkers, by the CMU ARCTIC speaker names. Sources: docs/speech/SOURCES.md.
SPEAKERS = {"Male talker": "bdl", "Female talker": "slt"}
talkers = {
    label: finish(so.load(fetch(f"docs/speech/{speaker}_arctic_a0131.flac")))
    for label, speaker in SPEAKERS.items()
}
tracks = {label: so.f0_track(snd) for label, snd in talkers.items()}
medians = {label: float(np.median(track.f0[0][track.voiced[0]])) for label, track in tracks.items()}
harvest = {}
for label, speaker in SPEAKERS.items():
    harvest_times, harvest_f0 = np.loadtxt(
        fetch(f"docs/speech/{speaker}_arctic_a0131_f0.csv"), delimiter=",", skiprows=2
    ).T
    harvest[label] = (harvest_times, harvest_f0)
for label, f0 in medians.items():
    print(f"{label}: median F0 {f0:.1f} Hz")

FS = 16000
DUR = 1.0
# "Hod" for each talker's pitch, (frequency, bandwidth) in Hz: Peterson and Barney's male and
# female F1 to F3, and F4 and F5 as on the source-filter page.
HOD = {
    "Male talker": [(730, 60), (1090, 90), (2440, 150), (3500, 200), (4500, 250)],
    "Female talker": [(850, 60), (1220, 90), (2810, 150), (4100, 200), (4900, 250)],
}


def contour_of(track):
    """The times and F0 values of a track (an F0Track or a (times, values) pair), NaN where
    unvoiced, for drawing."""
    times, values = (track.t, track.f0[0]) if hasattr(track, "f0") else track
    return times, np.where(np.asarray(values) > 0, values, np.nan)


def draw(container, snd, title, fmax, contour=None):
    """Waveform and narrowband spectrogram (Hann 33 ms, harmonics resolved), on one time axis, in
    a figure or subfigure, with an F0 contour over the spectrogram if one is given. Returns the two
    panels, for the playhead."""
    ax0, ax1 = container.subplots(2, 1, sharex=True, height_ratios=[0.45, 1])
    snd.plot(ax0, color="k", lw=0.4)
    ax0.set(title=title, xlabel="")
    so.STFT(snd, win_dur=0.0333, hop_dur=0.002).plot(ax1, db_range=70, colorbar=False, fmax=fmax)
    ax1.set_title("Spectrogram (Hann 33 ms)")
    if contour is not None:
        times, values = contour_of(contour)
        ax1.plot(times, values, color="c", lw=1.2)
    for ax in (ax0, ax1):
        ax.set_xlim(0, snd.duration)
    return [ax0, ax1]


def show(snd, title, fmax=5000):
    fig = plt.figure(figsize=(10, 4.2), layout="constrained")
    return fig, draw(fig, snd, f"Waveform: {title}", fmax)


def show_pair(sounds, titles, fmax=5000, contours=None):
    """One column per sound; `titles` maps each sound's label to its column's title, and
    `contours`, if given, to an F0 track drawn over its spectrogram."""
    fig = plt.figure(figsize=(10, 4.4), layout="constrained")
    columns = fig.subfigures(1, 2)
    playhead = {
        label: draw(column, snd, titles[label], fmax, None if contours is None else contours[label])
        for column, (label, snd) in zip(columns, sounds.items(), strict=True)
    }
    return fig, playhead


def titled(what):
    return {label: f"{label}: {what}" for label in talkers}


def formants(snd, table):
    for freq, bandwidth in table:
        snd = so.resonator(snd, freq, bandwidth)
    return snd


# Klatt's glottal source and the radiation from the lips, and the formants, as gains at any
# frequency: the spectra of their impulse responses.
impulse = so.Sound(np.r_[1.0, np.zeros(2**14 - 1)], FS)
glottal_gain = np.abs(np.fft.rfft(so.resonator(impulse, 0, 100).data[:, 0]))
gain_freqs = np.fft.rfftfreq(2**14, 1 / FS)
tract_gains = {
    label: np.abs(np.fft.rfft(formants(impulse, table).data[:, 0])) for label, table in HOD.items()
}


def source_gain(f):
    return np.interp(f, gain_freqs, glottal_gain) * np.abs(2 * np.sin(np.pi * np.asarray(f) / FS))


def true_envelope_db(f, label):
    """The vowel's envelope in dB, up to a constant: source times radiation times formants."""
    return 20 * np.log10(source_gain(f) * np.interp(f, gain_freqs, tract_gains[label]) + 1e-12)


def breathy_source(f0):
    """Harmonics of f0 with Klatt's source spectrum, up to 7.2 kHz, and white noise with the
    harmonics' power per hertz at 4 kHz."""
    harmonic_numbers = np.arange(1, int(0.45 * FS / f0) + 1)
    buzz = so.harmonic_complex(
        DUR, FS, f0, harmonics=harmonic_numbers, amplitudes=lambda t, f: source_gain(f)
    )
    # harmonic_complex scales its output to RMS 1; the harmonics' power per hertz near f is then
    # (scale * gain(f))^2 / 2 per harmonic, one harmonic every f0 hertz
    scale = 1 / np.sqrt(np.sum(source_gain(harmonic_numbers * f0) ** 2) / 2)
    noise_density = (scale * source_gain(4000.0)) ** 2 / 2 / f0
    hiss = so.gaussian_noise(DUR, FS, rng=1) * np.sqrt(noise_density * FS / 2)  # white: flat density
    return buzz, hiss


def true_aperiodicity(f):
    """The share of noise at f, from the source alone. With the noise as strong as the harmonics
    at 4 kHz it is the same at any pitch."""
    return 1 / (1 + (source_gain(f) / source_gain(4000.0)) ** 2)


sources = {label: breathy_source(f0) for label, f0 in medians.items()}
buzz_vowels = {label: formants(buzz, HOD[label]) for label, (buzz, _) in sources.items()}
vowels = {label: formants(buzz + hiss, HOD[label]) for label, (buzz, hiss) in sources.items()}
window_times = np.arange(0, DUR, 0.005)
exact = {label: (window_times, np.full(len(window_times), f0)) for label, f0 in medians.items()}  # every 5 ms
inside = (window_times > 0.1) & (window_times < 0.9)  # time windows away from the ends
pitch_names = {label: f"the {label.lower()}'s pitch, {f0:.0f} Hz" for label, f0 in medians.items()}

N_FFT = 1024
spectrum_freqs = np.fft.rfftfreq(N_FFT, 1 / FS)
band = (spectrum_freqs > 200) & (spectrum_freqs < 4000)


def hann_power(samples, center, length, shape="hann"):
    """P(t, f) at sample `center` under a window `length` samples long (may be fractional): Hann,
    or Blackman, scaled to unit energy."""
    offsets = np.arange(-int(np.ceil(length / 2)), int(np.ceil(length / 2)) + 1)
    phase = 2 * np.pi * offsets / length
    weights = (
        0.5 + 0.5 * np.cos(phase)
        if shape == "hann"
        else 0.42 + 0.5 * np.cos(phase) + 0.08 * np.cos(2 * phase)
    )
    weights[np.abs(offsets) >= length / 2] = 0
    segment = samples[int(round(center)) + offsets] * weights
    return np.abs(np.fft.rfft(segment / np.sqrt(np.sum(weights**2)), N_FFT)) ** 2


def ripples(label, periods, shape="hann", pair=False):
    """The ripple across frequency and over time (dB) of windows `periods` long, over 16
    positions within one period at 0.5 s; with `pair`, each position's power is TANDEM's
    average of two windows a quarter period before and after it."""
    samples, period = buzz_vowels[label].data[:, 0], FS / medians[label]
    levels = []
    for step in range(16):
        center = 0.5 * FS + step / 16 * period
        if pair:
            power = (
                hann_power(samples, center - period / 4, periods * period, shape)
                + hann_power(samples, center + period / 4, periods * period, shape)
            ) / 2
        else:
            power = hann_power(samples, center, periods * period, shape)
        levels.append(10 * np.log10(power + 1e-12))
    levels = np.array(levels)
    truth = true_envelope_db(spectrum_freqs, label)
    across = np.median([np.std((level - truth)[band]) for level in levels])
    over_time = np.median(np.std(levels[:, band], axis=0))
    return levels, across, over_time


print(f"{'[dB]':<34}{'off the envelope':>18}{'flicker':>12}")
ripple = {}
for label in talkers:
    for periods in (1, 2, 3, 6):
        ripple[label, periods] = ripples(label, periods)
        _, across, over_time = ripple[label, periods]
        print(
            f"{label + f', Hann {periods} period' + 's' * (periods > 1):<34}{across:>18.1f}{over_time:>12.2f}"
        )

print(f"{'Ripple over time [dB]':<34}{'one window':>12}{'TANDEM pair':>13}")
for label in talkers:
    single, pair = ripples(label, 2.5, "blackman")[2], ripples(label, 2.5, "blackman", pair=True)[2]
    print(f"{label + ', Blackman 2.5 periods':<34}{single:>12.2f}{pair:>13.2f}")

label = "Female talker"
F0, samples = medians[label], buzz_vowels[label].data[:, 0]
power = hann_power(samples, 0.5 * FS, 3 * FS / F0)  # step 1

# step 2: the moving average over 2 F0 / 3, on a grid eight times finer than the FFT's, with the
# spectrum mirrored about 0 Hz and Nyquist, as a real signal's spectrum continues there
mirrored = np.concatenate([power[:0:-1], power, power[-2::-1]])
fine = np.interp(np.arange(0, len(mirrored) - 1, 1 / 8), np.arange(len(mirrored)), mirrored)
half_width = int(round(F0 / 3 / (FS / N_FFT) * 8))
fine_smoothed = np.convolve(fine, np.ones(2 * half_width + 1) / (2 * half_width + 1), mode="same")
smoothed = fine_smoothed[::8][len(power) - 1 : 2 * len(power) - 1]

# step 3: the cepstrum of the log power, the two lifters, and back
cepstrum = np.fft.irfft(np.log(smoothed), N_FFT)
quefrency = np.minimum(np.arange(N_FFT), N_FFT - np.arange(N_FFT)) / FS


def liftered(q1):
    lifters = np.sinc(F0 * quefrency) * ((1 - 2 * q1) + 2 * q1 * np.cos(2 * np.pi * F0 * quefrency))
    return np.exp(np.fft.rfft(lifters * cepstrum, N_FFT).real)


steps = {
    "1. Hann, three periods": power,
    "2. smoothed over 2 F0 / 3": smoothed,
    "3. liftered, q1 = −0.09": liftered(-0.09),
}
library = so.pitch_adaptive_envelope(buzz_vowels[label], exact[label])
library_db = 10 * np.log10(library(np.array([0.5]), spectrum_freqs)[0, :, 0])
below_7k = spectrum_freqs < 7000
print(
    f"largest difference from so.pitch_adaptive_envelope below 7 kHz: "
    f"{np.max(np.abs(10 * np.log10(steps['3. liftered, q1 = −0.09']) - library_db)[below_7k]):.2f} dB"
)
truth = true_envelope_db(spectrum_freqs, label)
for name, step in steps.items():
    print(
        f"{name:<28} distance from the true envelope (SD, 200 Hz to 4 kHz): "
        f"{np.std((10 * np.log10(step + 1e-12) - truth)[band]):.1f} dB"
    )

Q1 = (0.0, -0.09, -0.15)
envelope_freqs = None
q1_envelopes = {}
for label in talkers:
    for q1 in Q1:
        envelope = so.pitch_adaptive_envelope(buzz_vowels[label], exact[label], q1=q1)
        envelope_freqs = envelope.f
        q1_envelopes[label, q1] = 10 * np.log10(envelope.data[0][:, inside].mean(axis=1))
print(f"{'Peak re the true peak [dB]':<38}{'F1':>7}{'F2':>7}{'F3':>7}")
for label in talkers:
    truth = true_envelope_db(envelope_freqs, label)
    in_band = (envelope_freqs > 200) & (envelope_freqs < 4000)
    for q1 in Q1:
        envelope_db = q1_envelopes[label, q1]
        shifted = envelope_db - np.median((envelope_db - truth)[in_band])
        heights = []
        for formant, _ in HOD[label][:3]:
            near = np.abs(envelope_freqs - formant) < 0.15 * formant
            heights.append(shifted[near].max() - truth[near].max())
        print(f"{label + f', q1 = {q1:g}':<38}" + "".join(f"{height:>7.1f}" for height in heights))

stfts = {label: so.STFT(snd, win_dur=0.040, hop_dur=0.005) for label, snd in talkers.items()}
cepstra = {label: so.Cepstrum(stft) for label, stft in stfts.items()}
envelopes = {label: so.pitch_adaptive_envelope(snd, tracks[label]) for label, snd in talkers.items()}
TIMES = {"Male talker": 0.50, "Female talker": 0.60}  # one time window in a vowel of each sentence
window = {}
for label, cep in cepstra.items():
    t, f0, peak = cep.f0()
    i = int(np.argmin(np.abs(t - TIMES[label])))
    window[label] = {"i": i, "t": t[i], "f0": f0[0, i], "cutoff": 0.5 / f0[0, i]}


def above_envelope(label, envelope_db):
    """How far each harmonic peak up to 4 kHz rises above an envelope (dB), in one time window."""
    stft, w = stfts[label], window[label]
    spectrum_db = 20 * np.log10(np.abs(stft.data[0, :, w["i"]]))
    gaps = []
    for harmonic in range(1, int(4000 / w["f0"]) + 1):
        near = np.nonzero(np.abs(stft.f - harmonic * w["f0"]) < w["f0"] / 2)[0]
        top = near[np.argmax(spectrum_db[near])]
        gaps.append(spectrum_db[top] - envelope_db[top])
    return np.array(gaps)


lifted, cheaptrick_db = {}, {}
for label, cep in cepstra.items():
    w = window[label]
    lifted[label] = 20 * np.log10(cep.lifter(w["cutoff"]).envelope()[0, :, w["i"]])
    cheaptrick_db[label] = 10 * np.log10(envelopes[label](np.array([w["t"]]), stfts[label].f)[0, :, 0])
    for name, envelope_db in (("lifter", lifted[label]), ("CheapTrick", cheaptrick_db[label])):
        gaps = above_envelope(label, envelope_db)
        level = f"median {np.median(gaps):5.1f} dB above, " if name == "lifter" else " " * 21
        print(f"{label:14} {name:11}: {level}spread (SD) {np.std(gaps):.1f} dB")

vowel, F0 = vowels["Male talker"], medians["Male talker"]
center = int(0.5 * FS)  # the time window at 0.5 s
half = int(round(2 * FS / F0))  # half of four periods
offsets = np.arange(-half, half + 1)
fit_window = 0.5 + 0.5 * np.cos(np.pi * offsets / (half + 1))
phase = 2 * np.pi * F0 * offsets / FS  # the running phase over the time window
harmonic_phases = np.outer(phase, np.arange(1, int(FS / 2 / F0) + 1))
ramp = (offsets / half)[:, None]
columns = np.column_stack(
    [np.ones(len(offsets)), np.cos(harmonic_phases), np.sin(harmonic_phases)]
    + [ramp * np.cos(harmonic_phases), ramp * np.sin(harmonic_phases)]
)
segment = vowel.data[center + offsets, 0]
root_window = np.sqrt(fit_window)  # weighted least squares: minimize the windowed residual
weights, *_ = np.linalg.lstsq(columns * root_window[:, None], segment * root_window, rcond=None)
residual = segment - columns @ weights
print(f"{columns.shape[1]} columns fitted to {len(offsets)} samples")
residual_power, signal_power = np.sum((fit_window * residual) ** 2), np.sum((fit_window * segment) ** 2)
print(f"windowed residual power over windowed signal power: {residual_power / signal_power:.3f}")

residual_shares = {label: so.harmonic_aperiodicity(vowels[label], exact[label]) for label in vowels}
for label, f0 in medians.items():
    share = residual_shares[label]
    error_db = 10 * np.log10(share.share[0][:, inside].mean(axis=1) / true_aperiodicity(share.f))
    low, high = np.percentile(error_db[(share.f >= 300) & (share.f <= 7000)], [0, 100])
    print(f"{pitch_names[label]}: {int(0.45 * FS / f0)} harmonics below 7.2 kHz, window {4000 / f0:.1f} ms")
    print(f"  measured minus truth, 300 to 7000 Hz: {low:+.1f} to {high:+.1f} dB")

band_edges = [0, 1000, 2000, 4000, 7000]
bands = "".join(f"{f'{low}-{high}':>11}" for low, high in zip(band_edges[:-1], band_edges[1:], strict=True))
print(f"{'Share of noise [dB], band [Hz]':<34}{bands}")
for label, vowel in vowels.items():
    track = exact[label]
    envelope = so.pitch_adaptive_envelope(vowel, track)
    # the truth in each band: noise power over total power, the envelope weighting each frequency
    # as .bands() weights it
    weights = envelope.data[0][:, inside].mean(axis=1)
    truth_db = []
    for low, high in zip(band_edges[:-1], band_edges[1:], strict=True):
        in_band = (envelope.f >= low) & (envelope.f < high)
        share = np.sum(true_aperiodicity(envelope.f[in_band]) * weights[in_band]) / np.sum(weights[in_band])
        truth_db.append(10 * np.log10(share))
    print(f"{label.lower()}'s pitch, {medians[label]:.0f} Hz")
    print(f"{'  truth':<34}" + "".join(f"{value:>11.1f}" for value in truth_db))
    wrong_track = (window_times, track[1] * 1.01)
    for name, measure in [("harmonic residual", so.harmonic_aperiodicity), ("D4C", so.band_aperiodicity)]:
        for f0_name, f0 in [("exact F0", track), ("F0 1% high", wrong_track)]:
            shares = measure(vowel, f0).bands(band_edges, envelope=envelope)[0][:, inside].mean(axis=1)
            row = f"  {name}, {f0_name}"
            print(f"{row:<34}" + "".join(f"{value:>11.1f}" for value in 10 * np.log10(shares)))

harvest_envelopes = {label: so.pitch_adaptive_envelope(talkers[label], harvest[label]) for label in talkers}
d4c_maps = {label: so.band_aperiodicity(talkers[label], harvest[label]) for label in talkers}
residual_maps = {label: so.harmonic_aperiodicity(talkers[label], harvest[label]) for label in talkers}

print(f"{'Median share of noise [dB], band [Hz]':<44}{bands}")
for label in talkers:
    voiced = np.interp(d4c_maps[label].t, *harvest[label]) > 0
    for name, maps in [("D4C", d4c_maps), ("harmonic residual", residual_maps)]:
        shares = maps[label].bands(band_edges, envelope=harvest_envelopes[label])[0][:, voiced]
        row = f"{label}, {name}"
        print(f"{row:<44}" + "".join(f"{value:>11.1f}" for value in 10 * np.log10(np.median(shares, axis=1))))

PHASES = {"cosine": "cosine", "Schroeder": "schroeder+", "random": "random"}
buzzes = {
    name: {
        label: so.harmonic_complex(snd.duration, FS, tracks[label], phases=phases, rng=3)
        for label, snd in talkers.items()
    }
    for name, phases in PHASES.items()
}
for name, pair in buzzes.items():
    print(
        f"{name:9}: crest factor "
        + ", ".join(f"{snd.peak / snd.rms:.1f} ({label})" for label, snd in pair.items())
    )

def liftered_envelope(snd, track):
    """The cepstrum of 40 ms Hann time windows every 5 ms, liftered below half the pitch period
    of each window (the track's voiced F0, interpolated on a log scale across unvoiced gaps)."""
    cepstrum = so.Cepstrum(so.STFT(snd, win_dur=0.040, hop_dur=0.005))
    is_voiced = track.voiced[0]
    periods = 1 / np.exp(np.interp(cepstrum.t, track.t[is_voiced], np.log(track.f0[0][is_voiced])))
    return cepstrum.lifter(0.5 * periods).envelope_view()


def unvoiced_re_voiced(snd, track):
    """The RMS of a sound's unvoiced stretches over that of its voiced ones, in dB."""
    samples = snd.data[:, 0]
    voiced = np.interp(np.arange(len(samples)) / snd.fs, track.t, track.voiced[0].astype(float))
    return 20 * np.log10(np.sqrt(np.mean(((1 - voiced) * samples) ** 2) / np.mean((voiced * samples) ** 2)))


def distances(snd, label):
    """The level-free RMS dB between the log mel spectrograms of snd and the talker's recording,
    the median over the recording's voiced and over its unvoiced time windows."""
    original, rebuilt = so.MFCC(talkers[label], n_mels=40), so.MFCC(snd, n_mels=40)
    count = min(original.mel_db.shape[2], rebuilt.mel_db.shape[2])
    difference = rebuilt.mel_db[0][:, :count] - original.mel_db[0][:, :count]
    per_window = np.sqrt(np.mean((difference - difference.mean(axis=0)) ** 2, axis=0))
    voiced = np.interp(original.t[:count], tracks[label].t, tracks[label].voiced[0].astype(float)) > 0.5
    return np.median(per_window[voiced]), np.median(per_window[~voiced])


rebuilds = {}
for label, snd in talkers.items():
    track = tracks[label]
    rebuilds[label] = {
        "simple (page 2)": so.source_filter_synthesize(
            track, liftered_envelope(snd, track), fs=FS, duration=snd.duration, rng=2
        ),
        "CheapTrick, no aperiodicity": so.source_filter_synthesize(track, envelopes[label], rng=2),
        "CheapTrick and D4C": so.source_filter_synthesize(
            track, envelopes[label], so.band_aperiodicity(snd, track), rng=2
        ),
        "CheapTrick and harmonic residual": so.source_filter_synthesize(
            track, envelopes[label], so.harmonic_aperiodicity(snd, track), rng=2
        ),
    }
comparison = {}
print(f"{'':<50}{'unvoiced re voiced':>19}{'distance, voiced':>18}{'unvoiced':>10}")
for label, snd in talkers.items():
    print(f"{label + ', the recording':<50}{unvoiced_re_voiced(snd, tracks[label]):>16.1f} dB")
    for name, rebuilt in rebuilds[label].items():
        balance = unvoiced_re_voiced(rebuilt, tracks[label])
        comparison[label, name] = (balance, *distances(rebuilt, label))
        print(
            f"{'  ' + name:<50}{balance:>16.1f} dB"
            + "".join(f"{d:>15.2f} dB" for d in comparison[label, name][1:])
        )