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

#!/usr/bin/env python3

from ChildProject.projects import ChildProject
from ChildProject.annotations import AnnotationManager
from ChildProject.metrics import segments_to_annotation

import argparse

import datalad.api
from os.path import join as opj
from os.path import basename, exists

import multiprocessing as mp

import numpy as np
from scipy.stats import binom
import pandas as pd
from pyannote.core import Annotation, Segment, Timeline

import matplotlib
matplotlib.use("pgf")
matplotlib.rcParams.update({
    "pgf.texsystem": "pdflatex",
    'font.family': 'serif',
    'text.usetex': True,
    'pgf.rcfonts': False,
})
from matplotlib import pyplot as plt

effects = ["beta_sib_och", "beta_sib_adu", "alpha_dev", "beta_dev"]

variables = {
    "beta_sib_och": "$\\beta_{\\mathrm{OCH}}^{\\mathrm{sib}}$",
    "beta_sib_adu": "$\\beta_{\\mathrm{ADU}}^{\\mathrm{sib}}$",
    "alpha_dev": "$\\alpha_{\\mathrm{dev}}$",
    "beta_dev": "$\\beta_{\\mathrm{dev}}$",
}

prior_distribution = {
    effect: np.random.randn(10000) for effect in effects
}

samples = {
    "prior": prior_distribution,
    "truth": np.load("output/aggregates_truth_truth_only.npz"),
    "lena_raw": np.load("output/aggregates_lena_cougar_sibs_algo_siblings_adu.npz"),
    "vtc_raw": np.load("output/aggregates_vtc_cougar_sibs_algo_siblings_adu.npz"),
    # "vtc_calibrated": np.load("output/aggregates_vtc_dev_siblings_effect.npz"),
    # "lena_calibrated": np.load("output/aggregates_lena_dev_siblings_effect.npz"),
    # "vtc_calibrated": np.load("output/aggregates_vtc_fausey_15_dev_siblings.npz"),
    # "lena_calibrated": np.load("output/aggregates_lena_fausey_30_dev_siblings.npz")
    "vtc_calibrated": np.load("output/aggregates_vtc_sibs_dev_siblings.npz"),
    "lena_calibrated": np.load("output/aggregates_lena_sibs_dev_siblings.npz")
}

labels = {
    "prior": "Prior",
    "truth": "Manual annotations",
    "lena_raw": "LENA (uncalibrated)",
    "vtc_raw": "VTC (uncalibrated)",
    "lena_calibrated": "LENA (calibrated)",
    "vtc_calibrated": "VTC (calibrated)"
}

positions = {
    "prior": -2,
    "truth": 0,
    "lena_raw": 3,
    "vtc_raw": 2,
    "vtc_calibrated": 5,
    "lena_calibrated": 6,
}

def plot_effect(effect):
    fig, ax = plt.subplots(figsize=[6.4*0.8,3.2*0.8])

    for key in samples:
        mean = samples[key][effect].mean()
        up = np.quantile(samples[key][effect], q=1-0.05/2)
        low = np.quantile(samples[key][effect], q=0.05/2)
        ax.text(
            mean,
            positions[key]+0.25,
            f"\\footnotesize{{$\\mu={mean:.2f}$, $\\mathrm{{CI}}_{{95\\%}}=\\left[{low:.2f}, {up:.2f}\\right]$}}", ha="center"
        )
        ax.scatter([mean],[positions[key]], color="black" if key=="prior" else None)
        ax.errorbar([mean], [positions[key]], xerr=([mean-low], [up-mean]), ls="none", color="black" if key=="prior" else None)

    ax.set_xlabel(variables[effect])
    ax.set_yticks(list(positions.values()))
    ax.set_yticklabels([labels[key] for key in positions])
    ax.set_ylim(np.min(list(positions.values()))-0.25, np.max(list(positions.values()))+1)

    ax.axvline(0, 0, 1, color="black", ls="dashed")

    fig.savefig(f"output/effect_comparison_{effect}.eps", bbox_inches="tight")
    fig.savefig(f"output/effect_comparison_{effect}.png", bbox_inches="tight", dpi=720)
    fig.savefig(f"output/effect_comparison_{effect}.pdf", bbox_inches="tight")

for effect in effects:
    plot_effect(effect)