speaker-confusion-model/code/models/analysis.py

#!/usr/bin/env python

import pandas as pd
import numpy as np

from ChildProject.projects import ChildProject
from ChildProject.annotations import AnnotationManager
from ChildProject.metrics import segments_to_grid, segments_to_annotation
from ChildProject.pipelines.metrics import AclewMetrics
from ChildProject.utils import TimeInterval

from cmdstanpy import CmdStanModel

import datetime

import matplotlib
from matplotlib import pyplot as plt

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

import pickle

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

from pyannote.core import Annotation, Segment, Timeline

import argparse


parser = argparse.ArgumentParser()
parser.add_argument("--corpora", default=["input/bergelson"], nargs="+")
parser.add_argument("--duration", type=int, help="duration in hours", default=8)
parser.add_argument("--run")
parser.add_argument("--chains", type=int, default=1)
parser.add_argument("--warmup", type=int, default=250)
parser.add_argument("--samples", type=int, default=1000)
parser.add_argument("--threads-per-chain", type=int, default=4)
parser.add_argument("--models", default=["analysis"], nargs='+')
parser.add_argument("--validation", type=float, default=0)
args = parser.parse_args()

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

def extrude(self, removed, mode: str = "intersection"):
    if isinstance(removed, Segment):
        removed = Timeline([removed])

    truncating_support = removed.gaps(support=self.extent())
    # loose for truncate means strict for crop and vice-versa
    if mode == "loose":
        mode = "strict"
    elif mode == "strict":
        mode = "loose"

    return self.crop(truncating_support, mode=mode)


def compute_counts(parameters):
    corpus = parameters["corpus"]
    annotator = parameters["annotator"]

    project = ChildProject(parameters["path"])
    am = AnnotationManager(project)
    am.read()

    intersection = AnnotationManager.intersection(am.annotations, ["vtc", annotator])

    intersection["path"] = intersection.apply(
        lambda r: opj(
            project.path, "annotations", r["set"], "converted", r["annotation_filename"]
        ),
        axis=1,
    )
    datalad.api.get(list(intersection["path"].unique()))

    intersection = intersection.merge(
        project.recordings[["recording_filename", "child_id"]], how="left"
    )
    intersection["child"] = corpus + "_" + intersection["child_id"].astype(str)
    intersection["duration"] = (
        intersection["range_offset"] - intersection["range_onset"]
    )
    print(corpus, annotator, (intersection["duration"] / 1000 / 2).sum() / 3600)

    data = []
    for child, ann in intersection.groupby("child"):
        # print(corpus, child)

        segments = am.get_collapsed_segments(ann)
        if "speaker_type" not in segments.columns:
            continue

        segments = segments[segments["speaker_type"].isin(speakers)]

        vtc = {
            speaker: segments_to_annotation(
                segments[
                    (segments["set"] == "vtc") & (segments["speaker_type"] == speaker)
                ],
                "speaker_type",
            ).get_timeline()
            for speaker in speakers
        }

        truth = {
            speaker: segments_to_annotation(
                segments[
                    (segments["set"] == annotator)
                    & (segments["speaker_type"] == speaker)
                ],
                "speaker_type",
            ).get_timeline()
            for speaker in speakers
        }

        for speaker_A in speakers:
            vtc[f"{speaker_A}_vocs_explained"] = vtc[speaker_A].crop(
                truth[speaker_A], mode="loose"
            )
            vtc[f"{speaker_A}_vocs_fp"] = extrude(
                vtc[speaker_A], vtc[f"{speaker_A}_vocs_explained"]
            )
            vtc[f"{speaker_A}_vocs_fn"] = extrude(
                truth[speaker_A], truth[speaker_A].crop(vtc[speaker_A], mode="loose")
            )

            for speaker_B in speakers:
                vtc[f"{speaker_A}_vocs_fp_{speaker_B}"] = vtc[
                    f"{speaker_A}_vocs_fp"
                ].crop(truth[speaker_B], mode="loose")

                for speaker_C in speakers:
                    if speaker_C != speaker_B and speaker_C != speaker_A:
                        vtc[f"{speaker_A}_vocs_fp_{speaker_B}"] = extrude(
                            vtc[f"{speaker_A}_vocs_fp_{speaker_B}"],
                            vtc[f"{speaker_A}_vocs_fp_{speaker_B}"].crop(
                                truth[speaker_C], mode="loose"
                            ),
                        )

        d = {}
        keep_child = True
        for i, speaker_A in enumerate(speakers):
            for j, speaker_B in enumerate(speakers):
                if i != j:
                    z = len(vtc[f"{speaker_A}_vocs_fp_{speaker_B}"])
                else:
                    z = min(
                        len(vtc[f"{speaker_A}_vocs_explained"]), len(truth[speaker_A])
                    )

                d[f"vtc_{i}_{j}"] = z

                if z > len(truth[speaker_B]):
                    keep_child = False

            d[f"truth_{i}"] = len(truth[speaker_A])
            d["child"] = child

        d["duration"] = ann["duration"].sum() / 2 / 1000

        if keep_child:
            data.append(d)

    return pd.DataFrame(data).assign(
        corpus=corpus,
    )

def rates(parameters):
    corpus = parameters["corpus"]
    annotator = parameters["annotator"]
    speakers = ["CHI", "OCH", "FEM", "MAL"]

    project = ChildProject(parameters["path"])
    am = AnnotationManager(project)
    am.read()

    pipeline = AclewMetrics(
        project,
        vtc=annotator,
        alice=None,
        vcm=None,
        from_time="10:00:00",
        to_time="18:00:00",
        by="child_id",
    )
    metrics = pipeline.extract()
    metrics = pd.DataFrame(metrics).assign(corpus=corpus,annotator=annotator)
    metrics["duration"] = metrics[f"duration_{annotator}"]/1000/3600
    metrics = metrics[metrics["duration"] > 0.01]
    
    speakers = ['CHI', 'OCH', 'FEM', 'MAL']

    # metrics.dropna(subset={f"voc_{speaker.lower()}_ph" for speaker in speakers}&set(metrics.columns), inplace=True)

    for i, speaker in enumerate(speakers):
        # if f"voc_{speaker.lower()}_ph" not in metrics.columns:
        #     metrics[f"speech_rate_{i}"] = pd.NA
        # else:
        metrics[f"speech_rate_{i}"] = (metrics[f"voc_{speaker.lower()}_ph"]*(metrics["duration"])).fillna(0).astype(int)

    return metrics

def run_model(data, run, model_name):
    model = CmdStanModel(
        stan_file=f"code/models/{model_name}.stan",
        cpp_options={"STAN_THREADS": "TRUE"},
        compile="force",
    )
    fit = model.sample(
        data=data,
        chains=args.chains,
        threads_per_chain=args.threads_per_chain,
        iter_sampling=args.samples,
        iter_warmup=args.warmup,
        step_size=0.1,
        # save_profile=True,
        # show_console=True,
    )

    vars = fit.stan_variables()
    samples = {}
    for (k, v) in vars.items():
        samples[k] = v

    np.savez_compressed(f"output/aggregates_{run}_{model_name}.npz", **samples)
    samples = np.load(f"output/aggregates_{run}_{model_name}.npz")

    with open(f"output/aggregates_{run}_{model_name}.pickle", "wb") as f:
        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)

    print(samples["R2"].mean(axis=0))

    return {
        "evidence": samples["evidence"].mean(),
        "priors": samples["priors"].mean(),
        "total_evidence": (samples["evidence"]+samples["priors"]).mean(),
        "samples": samples
    }

def compile_recordings(corpus):
    project = ChildProject(corpus)
    am = AnnotationManager(project)
    am.read()

    project.recordings["age"] = project.compute_ages()

    annotations = am.annotations[am.annotations["set"] == "vtc"]
    annotations = annotations.merge(
        project.recordings,
        left_on="recording_filename",
        right_on="recording_filename",
        how="inner",
    )
    recs = []

    for recording_filename, _annotations in annotations.groupby("recording_filename"):
        _annotations = am.get_within_time_range(
            _annotations,
            TimeInterval(
                datetime.datetime(1900, 1, 1, 10, 0),
                datetime.datetime(1900, 1, 1, 10 + args.duration, 0),
            ),
        )
        child_id = _annotations["child_id"].max()
        age = _annotations["age"].max()
        duration = (_annotations["range_offset"] - _annotations["range_onset"]).sum()

        if duration < args.duration * 3600 * 1000:
            continue

        duration = args.duration * 3600 * 1000

        _annotations["path"] = _annotations.apply(
            lambda r: opj(
                project.path, "annotations", r["set"], "converted", r["annotation_filename"]
            ),
            axis=1,
        )

        missing_annotations = _annotations[~_annotations["path"].map(exists)]
        if len(missing_annotations):
            datalad.api.get(list(missing_annotations["path"].unique()))

        segments = am.get_segments(_annotations)
        segments["segment_onset"] -= segments["segment_onset"].min()
        segments = segments[segments["segment_onset"] >= 0]
        segments = segments[segments["segment_onset"] < duration]

        if len(segments) == 0:
            continue

        segments = segments[segments["speaker_type"].isin(["CHI", "OCH", "FEM", "MAL"])]

        rec = {
            f"vtc_{i}": len(segments[segments["speaker_type"] == speaker])
            for i, speaker in enumerate(speakers)
        }
        rec["recording"] = recording_filename
        rec["children"] = f"{corpus}_{child_id}"
        rec["corpus"] = basename(corpus)
        rec["age"] = age

        recs.append(rec)

    recs = pd.DataFrame(recs)
    return recs

if __name__ == "__main__":
    recs = pd.concat([compile_recordings(corpus) for corpus in args.corpora])
    recs["children"] = recs["children"].astype("category").cat.codes.astype(int) + 1

    annotators = pd.read_csv("input/annotators.csv")
    annotators["path"] = annotators["corpus"].apply(lambda c: opj("input", c))

    with mp.Pool(processes=args.chains*args.threads_per_chain) as pool:
        data = pd.concat(pool.map(compute_counts, annotators.to_dict(orient="records")))

    data = data.sample(frac=1)
    duration = data["duration"].sum()

    vtc = np.moveaxis(
        [[data[f"vtc_{j}_{i}"].values for i in range(4)] for j in range(4)], -1, 0
    )
    truth = np.transpose([data[f"truth_{i}"].values for i in range(4)])

    # speech rates at the child level
    annotators = annotators[~annotators['annotator'].str.startswith('eaf_2021')]
    with mp.Pool(processes=args.chains*args.threads_per_chain) as pool:
        speech_rates = pd.concat(pool.map(rates, annotators.to_dict(orient="records")))

    speech_rates.reset_index(inplace=True)
    speech_rates = speech_rates.groupby(["corpus", "child_id"]).sample(1)
    speech_rate_matrix = np.transpose([speech_rates[f"speech_rate_{i}"].values for i in range(4)])

    speech_rates.to_csv("rates.csv")

    print(vtc.shape)

    data["corpus"] = data["corpus"].astype("category")
    corpora = data["corpus"].cat.codes.values
    corpora_codes = dict(enumerate(data["corpus"].cat.categories))
    corpora_codes = {v: k for k, v in corpora_codes.items()}

    confusion_data = {
        "n_clips": truth.shape[0],
        "n_classes": truth.shape[1],
        "n_groups": data["child"].nunique(),
        "n_corpora": data["corpus"].nunique(),
        "n_validation": max(1, int(truth.shape[0] * args.validation)),
        "group": 1 + data["child"].astype("category").cat.codes.values,
        "conf_corpus": 1 + corpora,
        "truth": truth.astype(int),
        "vtc": vtc.astype(int),
        "speech_rates": speech_rate_matrix.astype(int),
        "group_corpus": 1+speech_rates["corpus"].map(corpora_codes).astype(int).values,
        "durations": speech_rates["duration"].values,
        "n_rates": len(speech_rates)
    }

    n_recs = len(recs)

    children_corpus = recs.groupby("children").agg(corpus=("corpus", "first")).sort_index()
    children_corpus = 1+children_corpus.corpus.map(corpora_codes).astype(int).values

    analysis_data = {
        "n_recs": n_recs,
        "n_children": len(recs["children"].unique()),
        "children": recs["children"],
        "vocs": np.transpose([recs[f"vtc_{i}"].values for i in range(4)]),
        "age": recs["age"],
        "corpus": children_corpus,
        "duration": args.duration,
    }

    data = {**analysis_data, **confusion_data}

    output = {}

    for model_name in args.models:
        output[model_name] = run_model(data, args.run, model_name)