speaker-confusion-model/code/models/enumeration.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,
    }
)

from collections import defaultdict

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=["causal"], nargs="+")
parser.add_argument("--confusion-clips", choices=["recording", "clip"], default="clip")
parser.add_argument("--confusion-groups", choices=["child", "recording"], default="recording")
parser.add_argument("--time-scale", type=int, 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 children_siblings(corpus):
    siblings = pd.read_csv("input/siblings.csv")
    siblings = siblings[siblings["corpus"]==corpus].set_index("child_id")
    siblings["child_id"] = siblings["child_id"].astype(str)
    siblings = siblings["n_siblings"].to_dict()

    n = defaultdict(lambda: -1, **siblings)

    return n



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

    project = ChildProject(parameters["path"])

    am = AnnotationManager(project)
    am.read()
    project.recordings["age"] = project.compute_ages()

    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", "age"]], how="left"
    )
    intersection["child"] = corpus + "_" + intersection["child_id"].astype(str)
    intersection["recording"] = corpus + "_" + intersection["recording_filename"].astype(str)
    intersection["clip_id"] = (
        corpus
        + "_"
        + intersection["recording_filename"].str.cat(intersection["range_onset"].astype(str))
    )

    if args.time_scale:
        intersection["onset"] = intersection.apply(
            lambda r: np.arange(
                r["range_onset"], r["range_offset"], args.time_scale * 1000
            ),
            axis=1,
        )
        intersection = intersection.explode("onset")
        intersection["range_onset"] = intersection["onset"]
        intersection["range_offset"] = (
            intersection["range_onset"] + args.time_scale * 1000
        ).clip(upper=intersection["range_offset"])

    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"):
    groupby = {
        "recording": ["child_id", "recording"],
        "clip": ["child_id", "recording", "range_onset"],
    }[args.confusion_clips]

    for clip, ann in intersection.groupby(groupby):
        child = clip[0]
        recording = clip[1]
        # 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",
            )
            # .support(collar=200)
            .get_timeline()
            for speaker in speakers
        }

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

        d = {}
        for i, speaker_A in enumerate(speakers):
            d[f"truth_total_{i}"] = len(truth[speaker_A])
            d[f"vtc_total_{i}"] = len(vtc[speaker_A])

        d["child"] = child
        d["recording"] = recording
        d["duration"] = ann["duration"].sum() / 2 / 1000 / 3600
        d["clip_id"] = ann["clip_id"].iloc[0]
        d["clip_age"] = ann["age"].iloc[0]

        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="recording_filename",
    )
    metrics = pipeline.extract()
    metrics = pd.DataFrame(metrics).assign(corpus=corpus, annotator=annotator)
    project.recordings["age"] = project.compute_ages()
    project.recordings["siblings"] = project.recordings.child_id.astype(str).map(
        children_siblings(corpus)
    )
    metrics = metrics.merge(
        project.recordings[["recording_filename", "age", "siblings"]]
    )
    metrics["duration"] = metrics[f"duration_{annotator}"] / 1000 / 3600
    metrics = metrics[metrics["duration"] > 0.01]
    metrics["child"] = corpus + "_" + metrics["child_id"].astype(str)

    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,
        max_treedepth=12,
        # 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)

    return {"samples": samples}


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

    project.recordings["age"] = project.compute_ages()
    project.recordings["siblings"] = project.recordings.child_id.astype(str).map(
        children_siblings(corpus)
    )

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

    if "normative" in annotations.columns:
        print("filtering out non-normative kids")
        annotations = annotations[annotations["normative"] != "N"]

    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()
        siblings = _annotations["siblings"].max()

        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
        rec["siblings"] = siblings

        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")))

    corpora = list(set(data["corpus"].unique()) | set(recs["corpus"].unique()))
    corpora_map = {
        corpus: i
        for i, corpus in enumerate(corpora)
    }
    print(corpora_map)

    duration = data["duration"].sum()
    clip_duration = data["duration"].values
    clip_age = data["clip_age"].values

    truth_total = np.transpose([data[f"truth_total_{i}"].values for i in range(4)])
    vtc_total = np.transpose([data[f"vtc_total_{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_rate_matrix = np.transpose(
        [speech_rates[f"speech_rate_{i}"].values for i in range(4)]
    )
    speech_rate_age = speech_rates["age"].values
    speech_rate_siblings = speech_rates["siblings"].values.astype(int)

    speech_rates.to_csv("rates.csv")

    data["corpus"] = data["corpus"].astype("category")
    
    confusion_data = {
        "n_clips": truth_total.shape[0],
        "n_classes": truth_total.shape[1],
        "n_groups": data[args.confusion_groups].nunique(),
        "n_corpora": len(corpora),
        "n_validation": 0,
        "group": 1 + data[args.confusion_groups].astype("category").cat.codes.values,
        "conf_corpus": 1 + data["corpus"].map(corpora_map).astype(int).values,
        "truth_total": truth_total.astype(int),
        "vtc_total": vtc_total.astype(int),
        "clip_duration": clip_duration,
        "clip_age": clip_age,
        "clip_id": 1 + data["clip_id"].astype("category").cat.codes.values,
        "n_unique_clips": data["clip_id"].nunique(),
        "speech_rates": speech_rate_matrix.astype(int),
        "speech_rate_age": speech_rate_age,
        "speech_rate_siblings": speech_rate_siblings,
        "group_corpus": (
            1 + speech_rates["corpus"].map(corpora_map).astype(int).values
        ),
        "speech_rate_child": 1+speech_rates["child"].astype("category").cat.codes.values,
        "n_speech_rate_children": speech_rates["child"].nunique(),
        "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_map).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"],
        "siblings": recs["siblings"].astype(int),
        "corpus": children_corpus,
        "recs_duration": args.duration,
    }

    data = {**analysis_data, **confusion_data}

    data["threads"] = args.threads_per_chain

    print(data)

    output = {}

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