speaker-confusion-model/code/plots/confusion_probs.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

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="plot_pred")
parser.add_argument("data")
parser.add_argument("fit")
parser.add_argument("output")
args = parser.parse_args()

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

fit = pd.read_parquet(args.fit)

fig = plt.figure(figsize=set_size(450, 1, 1))
axes = [fig.add_subplot(4, 4, i + 1) for i in range(4 * 4)]

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

n_groups = data["n_groups"]

for i in range(4 * 4):
    ax = axes[i]
    row = i // 4 + 1
    col = i % 4 + 1
    label = f"{col}.{row}"

    # if args.group is None:
    #    data = np.hstack([fit[f'alphas.{k}.{label}']/(fit[f'alphas.{k}.{label}']+fit[f'betas.{k}.{label}']).values for k in range(1,n_groups+1)])
    # else:
    #    data = fit[f'alphas.{args.group}.{label}']/(fit[f'alphas.{args.group}.{label}']+fit[f'betas.{args.group}.{label}']).values
    # data = np.hstack([(fit[f'group_mus.{k}.{label}']).values for k in range(1,59)])
    # data = fit[f'mus.{label}'].values
    if "fixed_bias.1.1" in fit.columns:
        data = expit(
            np.hstack(
                [
                    logit(fit[f"probs.{k+1}.{label}"].values)
                    + fit[f"fixed_bias.{label}"].values
                    for k in range(n_groups)
                ]
            )
        )
    else:
        data = np.hstack([fit[f"probs.{k+1}.{label}"].values for k in range(n_groups)])

    ax.set_xticks([])
    ax.set_xticklabels([])
    ax.set_yticks([])
    ax.set_yticklabels([])
    ax.set_ylim(0, 5)
    ax.set_xlim(0, 1)

    low = np.quantile(data, 0.0275)
    high = np.quantile(data, 0.975)

    if row == 1:
        ax.xaxis.tick_top()
        ax.set_xticks([0.5])
        ax.set_xticklabels([speakers[col - 1]])

    if row == 4:
        ax.set_xticks(np.linspace(0.25, 1, 3, endpoint=False))
        ax.set_xticklabels(np.linspace(0.25, 1, 3, endpoint=False))

    if col == 1:
        ax.set_yticks([2.5])
        ax.set_yticklabels([speakers[row - 1]])

    ax.hist(data, bins=np.linspace(0, 1, 40), density=True, histtype="step")
    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")

fig.suptitle("$p_{ij}$ distribution")
fig.subplots_adjust(wspace=0, hspace=0)

plt.savefig(args.output)

plt.show()