Two talkers

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()}


def show_pair(sounds, title, contours=None, fmax=5000):
    """One column per talker: the waveform, and a narrowband spectrogram (Hann 33 ms, which
    resolves the harmonics of both voices) with the pitch track over it if one is given.
    Returns the figure and, for each talker, the panels the playhead follows."""
    fig = plt.figure(figsize=(10, 4.4), layout="constrained")
    columns = fig.subfigures(1, 2)
    playhead = {}
    for column, (label, snd) in zip(columns, sounds.items(), strict=True):
        top, bottom = column.subplots(2, 1, sharex=True, height_ratios=[0.4, 1])
        snd.plot(top, color="k", lw=0.4)
        top.set(title=f"{label}: {title}", xlabel="", ylabel="")
        so.STFT(snd, win_dur=0.0333, hop_dur=0.002).plot(bottom, db_range=70, colorbar=False, fmax=fmax)
        bottom.set_title("Spectrogram (Hann 33 ms)")
        if contours is not None:
            track = contours[label]
            bottom.plot(track.t, np.where(track.voiced[0], track.f0[0] / 1000, np.nan), color="c", lw=1.2)
        for ax in (top, bottom):
            ax.set_xlim(0, snd.duration)
        playhead[label] = [top, bottom]
    return fig, playhead

voiced_f0 = {label: track.f0[0][track.voiced[0]] for label, track in tracks.items()}
for label, f0 in voiced_f0.items():
    low, middle, high = np.percentile(f0, [5, 50, 95])
    span = 12 * np.log2(high / low)
    print(f"{label}: median {middle:.0f} Hz, middle 90% {low:.0f} to {high:.0f} Hz ({span:.1f} semitones)")
ratio = np.median(voiced_f0["Female talker"]) / np.median(voiced_f0["Male talker"])
print(f"female / male median: {ratio:.2f}, {12 * np.log2(ratio):.1f} semitones")

envelopes = {label: so.cheaptrick(talkers[label], tracks[label]) for label in talkers}
frequencies = envelopes["Male talker"].f
scored = (frequencies >= 100) & (frequencies <= 5000)
average_db = {
    label: np.mean(10 * np.log10(env(tracks[label].t[tracks[label].voiced[0]], frequencies)[0]), axis=1)
    for label, env in envelopes.items()
}


def misfit(ratio, tilt=False):
    """RMS dB between the female average and the male one read at f / ratio, 100 to 5000 Hz,
    with the level (and, if tilt is True, a tilt) left free."""
    male_read = np.interp(frequencies[scored] / ratio, frequencies, average_db["Male talker"])
    difference = average_db["Female talker"][scored] - male_read
    columns = [np.ones(scored.sum())] + ([np.log2(frequencies[scored])] if tilt else [])
    basis = np.column_stack(columns)
    return np.sqrt(np.mean((difference - basis @ np.linalg.lstsq(basis, difference, rcond=None)[0]) ** 2))


ratios = np.arange(0.80, 1.60, 0.005)
for tilt in (False, True):
    errors = [misfit(r, tilt) for r in ratios]
    best = ratios[np.argmin(errors)]
    print(f"{'level and tilt free' if tilt else 'level free':20}: r = {best:.3f}, {min(errors):.2f} dB RMS")
warp = ratios[np.argmin([misfit(r, True) for r in ratios])]

spectra = {label: so.long_term_spectrum([snd]) for label, snd in talkers.items()}
octave_edges = [63, 125, 250, 500, 1000, 2000, 4000, 8000]
print("octave from [Hz]        " + "".join(f"{low:>6}" for low in octave_edges[:-1]))
for label, spectrum in spectra.items():
    power = 10 ** (spectrum.level / 10)
    shares = [
        10 * np.log10(power[(spectrum.f >= low) & (spectrum.f < high)].sum() / power.sum())
        for low, high in zip(octave_edges[:-1], octave_edges[1:], strict=True)
    ]
    print(f"{label + ', share [dB]':24}" + "".join(f"{share:6.1f}" for share in shares))

FS = 16000
HOD = {"F1": 730, "F2": 1090, "F3": 2440, "F4": 3500, "F5": 4500}  # Peterson and Barney's male "hod"
medians = {label: float(np.median(f0)) for label, f0 in voiced_f0.items()}
vowels = {label: so.klatt_synthesize(0.6, FS, F0=f0, **HOD) for label, f0 in medians.items()}
for label, f0 in medians.items():
    print(f"{label}: F0 {f0:.0f} Hz, {int(1000 // f0)} harmonics below 1 kHz")