create_figures.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiments on Quantile Recalibration Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first part of this notebook creates various graphics to analyse the performance of Quantile Recalibration Training on the test dataset. The second part focuses on visualizing metrics obtained during training on the training and validation datasets, and examples of predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from tueplots import fonts\n",
    "\n",
    "from uq.analysis.dataframes import (\n",
    "\tload_config, load_df, make_test_df_for_tuning, get_datasets_df, get_duplication_df,\n",
    "\tset_hparams_columns, make_df_abb\n",
    ")\n",
    "from uq.analysis.plot_metrics import make_plot_df, plot_metric_comparison_per_epoch\n",
    "from uq.analysis.plot_cohen_d import build_cohen_d\n",
    "from uq.analysis.plot_reliability_diagrams import plot_reliability_diagrams, make_reliability_df\n",
    "from uq.analysis.plot_predictions import plot_predictions\n",
    "from uq.analysis.queries import disjunction\n",
    "from uq.utils.general import filter_dict, set_notebook_options, savefig, op_without_index\n",
    "from uq.analysis.constants import ext, main_metrics, other_metrics, metric_names, cat_cmap\n",
    "from uq.analysis.helpers import standard_setting, preprocess, plot_all_cohen_d, filter_df, add_name_index, latex_names\n",
    "\n",
    "set_notebook_options()\n",
    "np.seterr(invalid='ignore')\n",
    "\n",
    "path = Path('results')\n",
    "\n",
    "# LaTeX fonts (requires more time to create figures)\n",
    "plt.rcParams.update(fonts.aistats2023_tex())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Metrics on the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = load_config('logs/full/')\n",
    "reload_df = True\n",
    "df = load_df(config, tuning=True) if reload_df else None\n",
    "test_df = make_test_df_for_tuning(df, config)\n",
    "test_df = preprocess(test_df)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of some methods against a baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_some(name, mixture_size=[3], cal_size=[2048], base_model='nn', **kwargs):\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method.isna() and posthoc_dataset == \"train\"'\n",
    "    )\n",
    "    join_by = []\n",
    "\n",
    "    plot_df = filter_df(test_df, mixture_size=mixture_size, cal_size=cal_size, base_model=base_model)\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'base_loss == \"nll_inhoc_mc\" and inhoc_method in [\"QRT\", \"QREG\"] and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "        'inhoc_method.isna() and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "\n",
    "    names = [r'\\textbf{QRTC}', 'QRC', 'QREGC']\n",
    "    cmap_cohen = dict(zip(names, cat_cmap))\n",
    "    cmap_diff = dict(zip(names[:2], cat_cmap))\n",
    "\n",
    "    plot_all_cohen_d(\n",
    "        config, name, plot_df, baseline_query, join_by, path,\n",
    "        cohen_kwargs={'cmap': cmap_cohen},\n",
    "        cd_kwargs=None,\n",
    "        diff_kwargs={'cmap': cmap_diff},\n",
    "        **kwargs\n",
    "    )\n",
    "    return plot_df\n",
    "\n",
    "plot_df_some = plot_some('some')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_some('some_mixture_1', mixture_size=[1])\n",
    "plot_some('some_mixture_10', mixture_size=[10])\n",
    "plot_df_some_mixture_3_resnet = plot_some('some_mixture_3_resnet', base_model='resnet');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of variants against no variant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_variants(name, **kwargs):\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method.isna() and posthoc_dataset == \"calib\"'\n",
    "    )\n",
    "    join_by = []\n",
    "\n",
    "    plot_df = filter_df(test_df, no_inhoc_variant=False)\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'base_loss == \"nll_inhoc_mc\" and inhoc_method == \"QRT\" and inhoc_dataset == \"batch\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "\n",
    "    cmap = {\n",
    "        r'\\textbf{QRTC}': cat_cmap[0],\n",
    "        'QRIC': cat_cmap[3],\n",
    "        'QRGC': cat_cmap[4],\n",
    "        'QRLC': cat_cmap[5],\n",
    "    }\n",
    "\n",
    "    plot_all_cohen_d(config, name, plot_df, baseline_query, join_by, path, cmap=cmap,\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "plot_variants('inhoc_variants')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of base models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_base_models(name, **kwargs):\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method.isna() and posthoc_dataset == \"train\" and mixture_size == 3 and base_model == \"nn\" '\n",
    "    )\n",
    "    join_by = []\n",
    "\n",
    "    plot_df = filter_df(test_df, base_model=None, mixture_size=None)\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'base_loss == \"nll_inhoc_mc\" and inhoc_method == \"QRT\" and inhoc_dataset == \"batch\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "\n",
    "    \n",
    "\n",
    "    plot_all_cohen_d(\n",
    "        config, name, plot_df, baseline_query, join_by, path,\n",
    "        add_base_model=True, add_mixture_size=True,\n",
    "        cohen_kwargs={'wspace': 0.85},\n",
    "        cd_kwargs={'textspace': 0.85},\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "plot_base_models('base_models')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of inhoc_cal_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cal_size(name, **kwargs):\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method.isna() and posthoc_dataset == \"train\"'\n",
    "    )\n",
    "    join_by = []\n",
    "\n",
    "    plot_df = filter_df(test_df, inhoc_cal_size=None, inhoc_dataset=None)\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'cal_size == 2048 and base_loss == \"nll_inhoc_mc\" and inhoc_method == \"QRT\" and inhoc_dataset == \"batch\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "        'not inhoc_cal_size.isna() and base_loss == \"nll_inhoc_mc\" and inhoc_method == \"QRT\" and inhoc_dataset == \"train\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "\n",
    "    plot_all_cohen_d(\n",
    "        config, name, plot_df, baseline_query, join_by, path, \n",
    "        only_inhoc_cal_size=True,\n",
    "        cohen_kwargs={'wspace': 0.58},\n",
    "        diff_kwargs={},\n",
    "        cd_kwargs={'textspace': 0.75},\n",
    "        cmap='inhoc_cal_size',\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "plot_cal_size('inhoc_cal_size')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unbounded vs Truncated vs Reflected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_reflected(name, **kwargs):\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method.isna() and posthoc_dataset == \"train\"'\n",
    "    )\n",
    "    join_by = []\n",
    "\n",
    "    plot_df = filter_df(test_df, only_posthoc_reflected=False, cal_size=[2047, 2048])\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'base_loss == \"nll_inhoc_mc\" and inhoc_method == \"QRT\" and inhoc_dataset == \"batch\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "\n",
    "    cmap = {\n",
    "        'QRTC-REFL': cat_cmap[0],\n",
    "        'QRTC-TRUNC': cat_cmap[3],\n",
    "        'QRTC-KDE': cat_cmap[4],\n",
    "    }\n",
    "\n",
    "    plot_all_cohen_d(\n",
    "        config, name, plot_df, baseline_query, join_by, path, \n",
    "        only_posthoc_reflected=True,\n",
    "        main_metrics=['test_nll', 'test_calib_l1'],\n",
    "        cohen_kwargs={'wspace': 0.65},\n",
    "        cd_kwargs={'textspace': 0.75},\n",
    "        cmap=cmap,\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "plot_reflected('reflected')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of inhoc_alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_inhoc_alpha(name, **kwargs):\n",
    "    # Note: the baseline has the posthoc_method set to smooth_ecdf\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"'\n",
    "    )\n",
    "    join_by = []\n",
    "\n",
    "    plot_df = filter_df(test_df, inhoc_alpha_to_tune=False)\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'base_loss == \"nll_inhoc_mc\" and inhoc_dataset == \"batch\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "    plot_all_cohen_d(config, name, plot_df, baseline_query, join_by, path, \n",
    "        only_inhoc_alpha=True, \n",
    "        cohen_kwargs={'wspace': 0.4},\n",
    "        diff_kwargs=None,\n",
    "        cd_kwargs={'textspace': 0.6},\n",
    "        cmap='inhoc_alpha',\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "def check_tuning():\n",
    "    from uq.analysis.helpers import tune_inhoc_alpha      \n",
    "\n",
    "    plot_df = test_df\n",
    "    plot_df = plot_df.query('dataset == \"CPU\"')\n",
    "    plot_df = plot_df.query('inhoc_b.isna() or inhoc_b == 0.1')        \n",
    "    plot_df = plot_df.query('inhoc_alpha <= 0')\n",
    "    plot_df = plot_df.query('run_id == 0')\n",
    "    plot_df = plot_df.query('posthoc_method == \"smooth_ecdf\"')\n",
    "    plot_df = plot_df.query('mixture_size == 3')\n",
    "    plot_df = plot_df.query('inhoc_dataset == \"batch\"')\n",
    "    plot_df = plot_df.query('cal_size.isna() or cal_size == 2048')\n",
    "    plot_df = plot_df.query('posthoc_dataset == \"calib\"')                  \n",
    "\n",
    "    display(tune_inhoc_alpha(plot_df))\n",
    "\n",
    "plot_inhoc_alpha('inhoc_alpha')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of inhoc_b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_inhoc_b(name, **kwargs):\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method.isna() and posthoc_dataset == \"train\"'\n",
    "    )\n",
    "    join_by = []\n",
    "    plot_df = filter_df(test_df, keep_non_tuned_inhoc_b=True)\n",
    "\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'base_loss == \"nll_inhoc_mc\" and inhoc_method == \"QRT\" and inhoc_dataset == \"batch\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "\n",
    "    plot_all_cohen_d(config, name, plot_df, baseline_query, join_by, path, \n",
    "        only_inhoc_b=True, \n",
    "        cohen_kwargs={'wspace': 0.4},\n",
    "        diff_kwargs={},\n",
    "        cd_kwargs={'textspace': 0.6},\n",
    "        cmap='inhoc_b',\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "plot_inhoc_b('inhoc_b')\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Relationship with other metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_query = ('regul.isna() '\n",
    "    'and inhoc_method.isna() and inhoc_variant.isna() '\n",
    "    'and posthoc_method.isna() and posthoc_dataset == \"calib\"'\n",
    ")\n",
    "join_by = []\n",
    "plot_df = filter_df(test_df)\n",
    "queries = [\n",
    "    'inhoc_method == \"QRT\" and inhoc_dataset == \"batch\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    'inhoc_method == \"QRT\" and inhoc_dataset == \"batch\" and posthoc_method.isna() and posthoc_dataset == \"calib\"',\n",
    "    'inhoc_method.isna() and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    'inhoc_method.isna() and posthoc_method.isna() and posthoc_dataset == \"calib\"',\n",
    "]\n",
    "plot_df = plot_df.query(disjunction(queries))\n",
    "plot_df = standard_setting(config, plot_df)\n",
    "cohen_df = build_cohen_d(plot_df, main_metrics, baseline_query, join_by)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_dup = get_duplication_df(config)\n",
    "cohen_df = cohen_df.merge(df_dup, on=['dataset'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uq.analysis.df_tuning import make_test_df_mean\n",
    "\n",
    "cohen_df = cohen_df.sort_values('Proportion of top 10 classes')\n",
    "\n",
    "for name, df_name in cohen_df.groupby('name', dropna=False):\n",
    "    groups = df_name.groupby('metric', dropna=False)\n",
    "    fig, axes = plt.subplots(1, len(groups), figsize=(10.5, 2.8), sharey=True, sharex=True, dpi=150)\n",
    "    for (metric, df_metric), axis in zip(groups, axes):\n",
    "        from sklearn.isotonic import IsotonicRegression\n",
    "        from sklearn.linear_model import LinearRegression\n",
    "        from sklearn.tree import DecisionTreeRegressor\n",
    "\n",
    "        model = IsotonicRegression(increasing='auto')\n",
    "        y_ = model.fit_transform(df_metric['Proportion of top 10 classes'], df_metric['Cohen\\'s d'])\n",
    "        axis.plot(y_, df_metric['Proportion of top 10 classes'], color='red')\n",
    "        #model = DecisionTreeRegressor(max_depth=2)\n",
    "        # model = LinearRegression()\n",
    "        # model.fit(df_metric['Proportion of top 10 classes'].to_numpy().reshape(-1, 1), df_metric['Cohen\\'s d'])\n",
    "        # x = np.linspace(0, 1, 1000)\n",
    "        # y_ = model.predict(x.reshape(-1, 1))\n",
    "        # axis.plot(y_, x, color='orange')\n",
    "\n",
    "\n",
    "        axis.scatter(df_metric[\"Cohen's d\"], df_metric['Proportion of top 10 classes'], alpha=0.5)\n",
    "        metric_name = metric_names[metric.split('_', 1)[1]]\n",
    "        axis.set_xscale('symlog', linthresh=1)\n",
    "        axis.set_ylim(0 - 0.02, 1 + 0.02)\n",
    "        axis.set(xlabel=f\"Cohen's d of {metric_name}\", ylabel='Discreteness level')\n",
    "        axis.axvline(0, color='black', linestyle='--')\n",
    "    axes[0].set(ylabel='Proportion of top 10 classes')\n",
    "    fig.tight_layout()\n",
    "    savefig(path / 'cohens_d_vs_top10' / f'{name}.{ext}', fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Measure of the time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_time_df(plot_df):\n",
    "    plot_df = plot_df.copy()\n",
    "    cols = [col for col in test_df.columns if 'time' in col]\n",
    "    plot_df = make_test_df_mean(plot_df)\n",
    "    plot_df = add_name_index(plot_df)\n",
    "    plot_df = plot_df.reset_index()\n",
    "    plot_df['time_per_epoch'] = plot_df['train_time'] / plot_df['best_iter']\n",
    "    replace_dict = {\n",
    "        'Rec (ca)': 'Base + Rec',\n",
    "        'RT (ba)': 'RT',\n",
    "        'RT (ba) + Rec (ca)': 'RT + Rec',\n",
    "        'QR (ba) + Rec (ca)': 'QR + Rec',\n",
    "    }\n",
    "    plot_df['name'] = plot_df['name'].replace(replace_dict)\n",
    "    metrics = ['train_time', 'best_iter', 'time_per_epoch']\n",
    "    plot_df = plot_df[['dataset', 'name'] + metrics].set_index(['dataset', 'name']).stack().rename_axis(index={None: 'metric'}).to_frame(name='value').reset_index()\n",
    "    plot_df = plot_df[['dataset', 'name', 'metric', 'value']].pivot_table(index='dataset', columns=['metric', 'name'], values='value', sort=False)\n",
    "    # Sort using the list `metrics`\n",
    "    order_map = {metric: i for i, metric in enumerate(metrics)}\n",
    "    plot_df = plot_df.sort_index(axis='columns', level='metric', key=lambda x: x.map(order_map))\n",
    "    return plot_df\n",
    "\n",
    "time_df = make_time_df(plot_df_some.query('inhoc_dataset != \"calib\" and not (inhoc_method.isna() and posthoc_method == \"smooth_ecdf\")'))\n",
    "print(time_df.style\n",
    "    .format(precision=2)\n",
    "    .format_index(escape='latex', axis=0)\n",
    "    .format_index(escape='latex', axis=1)\n",
    "    .to_latex(hrules=True, convert_css=True, multirow_align='t')\n",
    ")\n",
    "time_df"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reliability diagrams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_rel_diags(name):\n",
    "    plot_df = filter_df(test_df)\n",
    "\n",
    "    baseline_query = ('regul.isna() '\n",
    "        'and inhoc_method.isna() and inhoc_dataset == \"batch\" '\n",
    "        'and posthoc_method.isna() and posthoc_dataset == \"train\"'\n",
    "    )\n",
    "    queries = [\n",
    "        baseline_query,\n",
    "        'base_loss == \"nll_inhoc_mc\" and inhoc_method == \"QRT\" and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "        'inhoc_method.isna() and posthoc_method == \"smooth_ecdf\" and posthoc_dataset == \"calib\"',\n",
    "    ]\n",
    "    plot_df = plot_df.query(disjunction(queries))\n",
    "    plot_df = standard_setting(config, plot_df)\n",
    "\n",
    "    df_abb = make_df_abb(plot_df.reset_index()['dataset'].unique())\n",
    "    index = plot_df.index.names\n",
    "    plot_df = plot_df.reset_index().merge(df_abb, on='dataset')\n",
    "    plot_df['dataset'] = plot_df['abb']\n",
    "    plot_df = plot_df.set_index(index)\n",
    "\n",
    "    rel_df = make_reliability_df(plot_df)\n",
    "    fig = plot_reliability_diagrams(rel_df, ncols=7, agg_run=True)\n",
    "    savefig(path / 'rel_diags' / f'{name}.{ext}')\n",
    "\n",
    "plot_rel_diags('some')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Metrics on the training and validation datasets per epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = load_config('logs/per_epoch/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uq.analysis.dataframes import set_test_metrics_columns, get_hparams\n",
    "\n",
    "\n",
    "plot_df = load_df(config, tuning=True)\n",
    "set_hparams_columns(plot_df, config)\n",
    "set_test_metrics_columns(plot_df, ['val_nll'], add_infos=False)\n",
    "index = ['dataset_group', 'dataset', 'run_id'] + list(get_hparams(plot_df))\n",
    "plot_df = plot_df.set_index(index)\n",
    "plot_df = filter_df(plot_df, all_base_loss=True, inhoc_alpha_to_tune=False)\n",
    "plot_df = plot_df.drop(columns=['val_nll'])\n",
    "plot_df = preprocess(plot_df)\n",
    "\n",
    "plot_df = standard_setting(config, plot_df).reset_index()\n",
    "\n",
    "df_abb = make_df_abb(plot_df.reset_index()['dataset'].unique())\n",
    "plot_df = plot_df.merge(df_abb, on='dataset')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uq.analysis.plot_metrics import add_train_val_metrics, make_df_per_metric_and_epoch, plot_metric_comparison_per_dataset_per_epoch\n",
    "\n",
    "def preprocess_plot_df(plot_df, metrics_to_plot):\n",
    "    plot_df = add_train_val_metrics(plot_df, set(metrics_to_plot) | {'train_inhoc_base_loss', 'val_base_loss', 'val_nll', 'val_es_loss'})\n",
    "    mask = plot_df.eval('inhoc_method.isna()')\n",
    "    plot_df.loc[mask, 'val_inhoc_regul_scaled'] = plot_df.loc[mask, 'val_es_loss'].apply(lambda d: {x: 0 for x in d})\n",
    "    plot_df.loc[mask, 'val_inhoc_base_loss'] = plot_df.loc[mask, 'val_unnormalized_nll']\n",
    "    mask = plot_df.eval('posthoc_method.isna()')\n",
    "    plot_df.loc[~mask, 'val_inhoc_regul_scaled'] = pd.NA\n",
    "    plot_df.loc[~mask, 'val_inhoc_base_loss'] = pd.NA\n",
    "    plot_df = make_df_per_metric_and_epoch(plot_df, metrics_to_plot + ['val_es_loss'])\n",
    "    return plot_df\n",
    "\n",
    "\n",
    "def plot_metric_per_dataset_per_epoch(plot_df, all_datasets=False):\n",
    "    plot_df = plot_df.copy()\n",
    "    if all_datasets:\n",
    "        datasets_to_plot = df_abb['abb'].unique()\n",
    "    else:\n",
    "        datasets_to_plot = ['YAC', 'NYC']\n",
    "\n",
    "    plt.rcParams.update(fonts.aistats2023_tex())\n",
    "    plt.rcParams.update({\n",
    "        'text.latex.preamble': r'\\usepackage{amsfonts}'\n",
    "    })\n",
    "    val_metric_names = {\n",
    "        'val_unnormalized_nll': r\"$\\sum_{i=1}^N -\\log f'_\\theta(Y_i \\mid X_i)$\", \n",
    "        'val_inhoc_base_loss': r\"$\\sum_{i=1}^N -\\log f_\\theta(Y_i \\mid X_i)$\", \n",
    "        'val_inhoc_regul_scaled': r'$\\alpha \\hat{H}(Z)$', \n",
    "        'val_calib_l1': r\"$\\textrm{PCE}(F'_\\theta)$\"\n",
    "    }\n",
    "    metric_names = val_metric_names\n",
    "    metrics_to_plot = list(metric_names.keys())\n",
    "    plot_df = preprocess_plot_df(plot_df, metrics_to_plot)\n",
    "\n",
    "    plot_df['metric'] = plot_df['metric'].replace(metric_names)\n",
    "    plot_df = plot_df.merge(df_abb, on='dataset')\n",
    "    plot_df['dataset'] = plot_df['abb']\n",
    "\n",
    "    cmap = {\n",
    "        'BASE': cat_cmap[1],\n",
    "        'QRC': cat_cmap[1],\n",
    "        'QRT': cat_cmap[0],\n",
    "        r'\\textbf{QRTC}': cat_cmap[0],\n",
    "    }\n",
    "    plot_df['name'] = pd.Categorical(plot_df['name'], cmap.keys())\n",
    "    plot_df = plot_df.sort_values('name')\n",
    "    fig = plot_metric_comparison_per_dataset_per_epoch(plot_df, datasets_to_plot, metric_names.values(), cmap=cmap, ncols_legend=4)\n",
    "    filename = f'metric_per_dataset_per_epoch_all.{ext}' if all_datasets else f'metric_per_dataset_per_epoch.{ext}'\n",
    "    savefig(path / filename, fig)\n",
    "\n",
    "\n",
    "plot_metric_per_dataset_per_epoch(\n",
    "    plot_df.query('mixture_size == 3 and posthoc_dataset == \"calib\"'),\n",
    "    all_datasets=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_metric_per_epoch(plot_df, metric):\n",
    "    plot_df = plot_df.copy()\n",
    "    datasets_to_plot = plot_df['dataset'].unique()\n",
    "    order = get_datasets_df(config).sort_values('Total instances').reset_index()['Dataset']\n",
    "    datasets_to_plot = order.to_frame().query('Dataset in @datasets_to_plot')['Dataset']\n",
    "    plot_df['dataset'] = plot_df['abb']\n",
    "    datasets_to_plot = pd.DataFrame({'dataset': datasets_to_plot}).merge(df_abb, on='dataset')['abb']\n",
    "\n",
    "    plot_df = preprocess_plot_df(plot_df, [metric])\n",
    "\n",
    "    cmap = {\n",
    "        'BASE': cat_cmap[1],\n",
    "        'QRC': cat_cmap[1],\n",
    "        'QRT': cat_cmap[0],\n",
    "        r'\\textbf{QRTC}': cat_cmap[0],\n",
    "    }\n",
    "    plot_df['name'] = pd.Categorical(plot_df['name'], cmap.keys())\n",
    "    plot_df = plot_df.sort_values('name')\n",
    "\n",
    "    fig = plot_metric_comparison_per_epoch(plot_df, 'dataset', datasets_to_plot, cmap=cmap, ncols=6, ncols_legend=4)\n",
    "    savefig(path / 'metric_per_epoch' / f'{metric}.{ext}', fig)\n",
    "\n",
    "for metric in ['val_nll', 'val_calib_l1', 'val_inhoc_regul_scaled', 'train_nll', 'train_calib_l1', 'train_inhoc_regul_scaled']:\n",
    "    plot_metric_per_epoch(plot_df.query('posthoc_dataset == \"calib\"'), metric)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_dataset_per_epoch(plot_df, dataset):\n",
    "    plot_df = plot_df.query('dataset == @dataset')\n",
    "    plot_df = make_plot_df(plot_df)\n",
    "    metrics = [\n",
    "        'train_calib_l1', 'train_nll', \n",
    "        'val_calib_l1', 'val_nll',\n",
    "    ]\n",
    "    plot_df = plot_df.query('metric in @metrics')\n",
    "\n",
    "    fig = plot_metric_comparison_per_epoch(plot_df, 'metric', metrics, ncols=2)\n",
    "    savefig(path / 'dataset_per_epoch' / f'{dataset}.{ext}', fig)\n",
    "\n",
    "for dataset in plot_df.dataset.unique():\n",
    "    plot_dataset_per_epoch(plot_df, dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examples of predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uq.utils.checkpoints import load_rc_checkpoint, load_model_checkpoint\n",
    "from uq.models.dist.base_dist_module import DistModule\n",
    "from uq.train import load_datamodule\n",
    "\n",
    "\n",
    "def get_checkpoint(series):\n",
    "    def get_original_hparams(d):\n",
    "        return {\n",
    "            k: v \n",
    "            for k, v in d.items() \n",
    "            if k.endswith('_grid') or k.endswith('_dataset') or not k.startswith('inhoc_') and not k.startswith('posthoc_')\n",
    "        }\n",
    "\n",
    "    rc = load_rc_checkpoint(config, series.dataset_group, series.dataset, hparams=get_original_hparams(series.hparams), model_cls=DistModule)\n",
    "    datamodule = load_datamodule(rc)\n",
    "    datamodule.load_datasets()\n",
    "    model = load_model_checkpoint(rc, datamodule)\n",
    "    model.eval()\n",
    "    model.module.get_name('inhoc').inhoc_model.build(-1, -1, 'predict')\n",
    "    return model, datamodule\n",
    "\n",
    "methods = {\n",
    "    'QRTC': 'abb == @abb and inhoc_method == \"QRT\" and posthoc_dataset == \"calib\" and posthoc_method == \"smooth_ecdf\" and inhoc_b == 0.1',\n",
    "    'QRC': 'abb == @abb and inhoc_method.isna() and posthoc_dataset == \"calib\" and posthoc_method == \"smooth_ecdf\"',\n",
    "    'BASE': 'abb == @abb and inhoc_method.isna() and posthoc_dataset == \"calib\" and posthoc_method.isna()',\n",
    "}\n",
    "\n",
    "for abb in ['CP2', 'ELE', 'ABA', 'MER', 'MIP']:\n",
    "    print(abb)\n",
    "    for query in methods.values():\n",
    "        selected = plot_df.query(query).query('mixture_size == 3')\n",
    "        assert len(selected) == 1, selected\n",
    "        series = selected.iloc[0]\n",
    "        model, datamodule = get_checkpoint(series)\n",
    "\n",
    "        for posthoc_module in model.posthoc_manager.modules:\n",
    "            if posthoc_module.hparams['method'] == series.hparams['posthoc_method']:\n",
    "                break\n",
    "        else:\n",
    "            raise ValueError(f'Could not find posthoc method {series.hparams[\"posthoc_method\"]}')\n",
    "        posthoc_module.build(0, -1, 'test', batch=None)\n",
    "\n",
    "        plot_predictions(series.name, model, posthoc_module, datamodule, nrows=1)\n",
    "        savefig(path / 'predictions' / abb / f'{series.name}.{ext}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}