speaker-confusion-model/code/plots/speech_rate.py

#!/usr/bin/env python3

import pandas as pd
import pickle
import numpy as np
from scipy.special import logit, expit

import argparse
import matplotlib
import matplotlib.pyplot as plt

from scipy.stats import gamma

from os.path import basename


matplotlib.use("pgf")
matplotlib.rcParams.update(
    {
        "pgf.texsystem": "pdflatex",
        "font.family": "serif",
        "font.serif": "Times New Roman",
        "text.usetex": True,
        "pgf.rcfonts": False,
    }
)


def set_size(width, fraction=1, ratio=None):
    fig_width_pt = width * fraction
    inches_per_pt = 1 / 72.27
    if ratio is None:
        ratio = (5**0.5 - 1) / 2
    fig_width_in = fig_width_pt * inches_per_pt
    fig_height_in = fig_width_in * ratio
    return fig_width_in, fig_height_in


parser = argparse.ArgumentParser(description="speech_rate")
parser.add_argument("data")
parser.add_argument("fit")
parser.add_argument("output")
parser.add_argument("--selected-corpus", type=int, default=-1)

args = parser.parse_args()

with open(args.data, "rb") as fp:
    data = pickle.load(fp)

fit = pd.read_parquet(args.fit)

corpora = pd.read_csv("output/training_set.csv")
corpora = corpora["corpus"].map(basename).tolist()

speakers = ["CHI", "OCH", "FEM", "MAL"]

n_groups = data["n_groups"]
n_corpora = data["n_corpora"]

colors = ['#377eb8', '#ff7f00', '#4daf4a', '#f781bf']

mu = np.zeros((data['n_corpora'], 4))
mu_low = np.zeros((data['n_corpora'], 4))
mu_high = np.zeros((data['n_corpora'], 4))

inputs = np.zeros((data['n_corpora'], len(fit)))
input = np.zeros(data['n_corpora'])
input_low = np.zeros(data['n_corpora'])
input_high = np.zeros(data['n_corpora'])

output_speech =  np.zeros((data['n_corpora'], len(fit)))
input_speech =  np.zeros((data['n_corpora'], len(fit)))

output_dur = np.zeros(data['n_corpora'])
output_dur_low = np.zeros(data['n_corpora'])
output_dur_high = np.zeros(data['n_corpora'])

input_dur = np.zeros(data['n_corpora'])
input_dur_low = np.zeros(data['n_corpora'])
input_dur_high = np.zeros(data['n_corpora'])

for c in range(data['n_corpora']):
    for i in range(4):
        # alphas = fit[f"speech_rate_alpha.{i+1}.{c+1}"].values
        # betas = fit[f"speech_rate_beta.{i+1}.{c+1}"].values/1000

        # voc_dur_alphas = fit[f"voc_dur_alpha.{i+1}.{c+1}"].values
        # voc_dur_betas = fit[f"voc_dur_beta.{i+1}.{c+1}"].values

        mus = 1000*fit[f"speech_rate_mu.{i+1}.{c+1}"].values
        voc_dur_means = fit[f"voc_dur_mu.{i+1}.{c+1}"].values

        print(corpora[c], voc_dur_means.mean())

        mu[c,i] = np.mean(mus)
        mu_low[c,i] = np.quantile(mus, q=0.05/2)
        mu_high[c,i] = np.quantile(mus, q=1-0.05/2)

        if i > 0:
            inputs[c,:] += mus
            input_speech[c,:] += voc_dur_means*mus
        else:
            output_speech[c,:] = voc_dur_means*mus

    
    input[c] = inputs[c,:].mean()
    input_low[c] = np.quantile(inputs[c,:], q=0.05/2)
    input_high[c] = np.quantile(inputs[c,:], q=1-0.05/2)

    output_dur[c] = output_speech[c,:].mean()/60
    output_dur_low[c] = np.quantile(output_speech[c,:], q=0.05/2)/60
    output_dur_high[c] = np.quantile(output_speech[c,:], q=1-0.05/2)/60

    input_dur[c] = input_speech[c,:].mean()/60
    input_dur_low[c] = np.quantile(input_speech[c,:], q=0.05/2)/60
    input_dur_high[c] = np.quantile(input_speech[c,:], q=1-0.05/2)/60


keep = [c not in ["winnipeg", "tsimane2017"] for c in corpora]

corpora = [corpora[i] for i in range(len(corpora)) if keep[i]]
mu = mu[keep,:]
mu_low = mu_low[keep,:]
mu_high = mu_high[keep,:]
input = input[keep]
input_low = input_low[keep]
input_high = input_high[keep]

input_dur = input_dur[keep]
input_dur_low = input_dur_low[keep]
input_dur_high = input_dur_high[keep]

output_dur = output_dur[keep]
output_dur_low = output_dur_low[keep]
output_dur_high = output_dur_high[keep]

print(corpora)
print(mu.shape)

fig, axes = plt.subplots(2, 2, sharex=True, sharey=True)

for row in range (2):
    for col in range(2):
        i = row+2*col
        axes[row,col].scatter(np.arange(mu.shape[0]), mu[:,i])
        axes[row,col].errorbar(np.arange(mu.shape[0]), mu[:,i], (mu[:,i]-mu_low[:,i], mu_high[:,i]-mu[:,i]) ,ls="none")
        axes[row,col].set_title(speakers[i])
        axes[row,col].set_ylim(0,1200)
        axes[row,col].set_xticks(np.arange(mu.shape[0]))
        axes[row,col].set_xticklabels(corpora)
        axes[row,col].xaxis.set_tick_params(rotation=90)
        
        if col==0:
            axes[row,col].set_ylabel("voc/h")

fig.suptitle("Latent population mean voc/h\n(human annotations)")
fig.savefig("output/quantities.png", bbox_inches="tight")

fig, axes = plt.subplots(1, 2, sharex=True, sharey=True)

for col in range(2):
    if col == 0:
        axes[col].scatter(np.arange(mu.shape[0]), mu[:,0])
        axes[col].errorbar(np.arange(mu.shape[0]), mu[:,0], (mu[:,0]-mu_low[:,0], mu_high[:,0]-mu[:,0]) ,ls="none")
    else:
        axes[col].scatter(np.arange(len(input)), input)
        axes[col].errorbar(np.arange(len(input)), input, (input-input_low, input_high-input) ,ls="none")

    axes[col].set_title("output" if col == 0 else "input")
    axes[col].set_ylim(0,2000)
    axes[col].set_xticks(np.arange(mu.shape[0]))
    axes[col].set_xticklabels(corpora)
    axes[col].xaxis.set_tick_params(rotation=90)
    
    if col==0:
        axes[col].set_ylabel("voc/h")

fig.suptitle("Latent population mean voc/h\n(human annotations)")
fig.savefig("output/input_output.png", bbox_inches="tight")


fig, axes = plt.subplots(1, 2, sharex=True, sharey=True)

for col in range(2):
    if col == 0:
        axes[col].scatter(np.arange(len(output_dur)), output_dur)
        axes[col].errorbar(np.arange(len(output_dur)), output_dur, (output_dur-output_dur_low, output_dur_high-output_dur) ,ls="none")
    else:
        axes[col].scatter(np.arange(len(input_dur)), input_dur)
        axes[col].errorbar(np.arange(len(input_dur)), input_dur, (input_dur-input_dur_low, input_dur_high-input_dur) ,ls="none")

    axes[col].set_title("output" if col == 0 else "input")
    axes[col].set_ylim(0,60)
    axes[col].set_xticks(np.arange(mu.shape[0]))
    axes[col].set_xticklabels(corpora)
    axes[col].xaxis.set_tick_params(rotation=90)
    
    if col==0:
        axes[col].set_ylabel("min/h")

fig.suptitle("Latent population mean min/h\n(human annotations)")
fig.savefig("output/input_output_dur.png", bbox_inches="tight")

fig, ax = plt.subplots(1,1)

for i in range(4):
    alphas = fit[f"speech_rate_alpha.{i+1}.{args.selected_corpus}"].values
    scale = 1000*fit[f"speech_rate_mu.{i+1}.{args.selected_corpus}"].values/alphas

    x = np.linspace(0,500,200,True)
    pdf = np.zeros((len(x), len(alphas)))
    for k in range(len(x)):
        pdf[k,:] = gamma.pdf(x[k], alphas, np.zeros(len(alphas)), scale)

    pdf_low = np.quantile(pdf, q=0.05, axis=1)
    pdf_high = np.quantile(pdf, q=0.95, axis=1)
    pdf_mean = np.mean(pdf, axis=1)

    ax.plot(x, pdf_mean, color=colors[i], label=speakers[i])
    ax.fill_between(x, pdf_low, pdf_high, color=colors[i], alpha=0.2)

    ax.set_xlim(0, 500)

    ax.set_xlabel("voc/h")

    # ax.axvline(np.mean(data), linestyle="--", linewidth=0.5, color="#333", alpha=1)
    # ax.text(0.5, 4.5, f"{low:.2f} - {high:.2f}", ha="center", va="center")

ax.legend()

corpus = pd.read_csv("output/training_set.csv")["corpus"].map(basename).tolist()[args.selected_corpus-1]
fig.suptitle(f"voc/h distribution for each kind of speaker ({corpus})")
fig.savefig(f"output/speech_distribution_{corpus}.png", bbox_inches="tight")

plt.show()