Pitch tracking

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

# WORLD's Harvest track of each recording: times and F0, every 5 ms, 0 where unvoiced.
harvest = {
    label: np.loadtxt(fetch(f"docs/speech/{speaker}_arctic_a0131_f0.csv"), delimiter=",", skiprows=2).T
    for label, speaker in SPEAKERS.items()
}

# The cepstrum of 40 ms Hann windows every 5 ms, and the cepstral F0 of each: (times, F0, peak height).
cepstra = {label: so.Cepstrum(so.STFT(snd, win_dur=0.040, hop_dur=0.005)) for label, snd in talkers.items()}
cepstral = {label: cep.f0(f_lo=75, f_hi=400) for label, cep in cepstra.items()}


def three_tracks(label):
    """The three tracks of one talker at so.f0_track's time windows, 0 where unvoiced. All three
    are on the same 5 ms grid from time 0, so this only picks out the shared time windows."""
    track = tracks[label]
    cepstral_t, cepstral_f0, _ = cepstral[label]
    harvest_t, harvest_f0 = harvest[label]
    return {
        "cepstral": cepstral_f0[0][np.searchsorted(np.round(cepstral_t, 6), np.round(track.t, 6))],
        "so.f0_track": np.where(track.voiced[0], track.f0[0], 0.0),
        "Harvest": harvest_f0[np.searchsorted(np.round(harvest_t, 6), np.round(track.t, 6))],
    }


pitch = {label: three_tracks(label) for label in talkers}
F0_TICKS = [60, 80, 100, 150, 200, 300, 400]
COLORS = {"cepstral": "tab:blue", "so.f0_track": "tab:orange", "Harvest": "0.55"}
medians = {label: np.median(tracks[label].f0[0][tracks[label].voiced[0]]) for label in talkers}
for label, f0 in medians.items():
    print(f"{label}: median F0 {f0:.0f} Hz, period {1e3 / f0:.2f} ms, {0.040 * f0:.1f} periods in 40 ms")

for label, f0 in pitch.items():
    voiced_f0 = f0["so.f0_track"][f0["so.f0_track"] > 0]
    print(f"{label}: two periods inside 13.3 ms on {np.mean(voiced_f0 > 150):.0%} of voiced windows")


def octave_below(f0, reference):
    """True where both are voiced and f0 is half the reference, within 5%."""
    both = (f0 > 0) & (reference > 0)
    return both & (np.abs(f0 / np.where(both, reference, 1.0) - 0.5) < 0.025)

for label, f0 in pitch.items():
    both = (f0["cepstral"] > 0) & (f0["Harvest"] > 0)
    close = np.mean(np.abs(f0["cepstral"][both] / f0["Harvest"][both] - 1) < 0.05)
    low = octave_below(f0["cepstral"], f0["Harvest"])
    times = ", ".join(f"{s:.3f} s" for s in tracks[label].t[low]) or "no time window"
    print(f"{label}: both voiced on {both.sum()} windows, cepstral F0 within 5% of Harvest on {close:.1%};")
    print(f"  an octave below Harvest on {low.sum()} ({low.sum() / both.sum():.1%}): {times}")

for label, f0 in pitch.items():
    track = tracks[label]
    harvest_only = (f0["Harvest"] > 0) & (f0["so.f0_track"] == 0)
    best_score = np.nanmax(np.nan_to_num(track.candidate_scores[0], nan=0.0), axis=1)
    without_rule = so.f0_track(talkers[label], subharmonic_margin=None).f0[0]
    with_rule = f0["so.f0_track"]
    both = (without_rule > 0) & (with_rule > 0)
    moved = np.sum(np.abs(without_rule[both] / with_rule[both] - 1) > 0.05)
    voicing_changed = np.sum((without_rule > 0) != (with_rule > 0))
    print(f"{label}: voiced on " + ", ".join(f"{np.mean(f > 0):.0%} ({name})" for name, f in f0.items()))
    print(
        f"  Harvest voices {harvest_only.sum()} windows the tracker does not; best score below 0.5 on "
        f"{np.mean(best_score[harvest_only] < 0.5):.0%} of them"
    )
    print(
        f"  without the subharmonic rule: {moved} windows move by over 5%, {voicing_changed} change voicing"
    )

PAIRS = [("cepstral", "Harvest"), ("so.f0_track", "Harvest"), ("cepstral", "so.f0_track")]
print(f"{'':24}{'both voiced':>12}{'within 5%':>10}{'octave low':>11}{'octave high':>12}{'one voiced':>11}")
for label, f0 in pitch.items():
    for first, second in PAIRS:
        a, b = f0[first], f0[second]
        both = (a > 0) & (b > 0)
        within = np.mean(np.abs(a[both] / b[both] - 1) < 0.05)
        low, high = octave_below(a, b).sum(), octave_below(b, a).sum()
        one = np.sum((a > 0) != (b > 0))
        print(
            f"{label.split()[0]:7}{first + ' / ' + second:24}{both.sum():5}"
            f"{within:10.1%}{low:11}{high:12}{one:11}"
        )

def show_pair(sounds, title, contours=None, fmax=4000):
    """One column per talker: the waveform, and a narrowband spectrogram (Hann 33 ms) with a 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")
    playhead = {}
    for column, (label, snd) in zip(fig.subfigures(1, 2), 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=60, colorbar=False, fmax=fmax)
        bottom.set_title("Spectrogram (Hann 33 ms)")
        if contours is not None:
            t, f0 = contours[label]
            bottom.plot(t, np.where(f0 > 0, f0, 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


def sawtooths(name):
    """A sawtooth on one of the three tracks, for each talker, and the track itself."""
    contours = {label: (tracks[label].t, pitch[label][name]) for label in talkers}
    buzz = {
        label: finish(so.sawtooth_wave(snd.duration, snd.fs, contours[label]))
        for label, snd in talkers.items()
    }
    return buzz, contours