colab_ua_accuracy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lIsrKf6hTp4B"
   },
   "source": [
    "# Uncertainty-adjusted accuracy evaluation\n",
    "\n",
    "See `README.md` for installation and usage instructions.\n",
    "\n",
    "This notebook re-creates some figures of [2] on the toy dataset\n",
    "introduced in [1] as well as the dermatology dataset.\n",
    "\n",
    "```\n",
    "[1] Stutz, D., Roy, A.G., Matejovicova, T., Strachan, P., Cemgil, A.T.,\n",
    "    & Doucet, A. (2023).\n",
    "    Conformal prediction under ambiguous ground truth. ArXiv, abs/2307.09302.\n",
    "[2] Stutz, D., Cemgil, A.T., Roy, A.G., Matejovicova, T., Barsbey, M.,\n",
    "    Strachan, P., Schaekermann, M., Freyberg, J.V., Rikhye, R.V., Freeman, B.,\n",
    "    Matos, J.P., Telang, U., Webster, D.R., Liu, Y., Corrado, G.S., Matias, Y.,\n",
    "    Kohli, P., Liu, Y., Doucet, A., & Karthikesalingam, A. (2023).\n",
    "    Evaluating AI systems under uncertain ground truth: a case study in\n",
    "    dermatology. ArXiv, abs/2307.02191.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nCL022QZTp4I"
   },
   "source": [
    "# Imports and setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CQCCCpXkTp4J"
   },
   "outputs": [],
   "source": [
    "import functools\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d7I7PYIlTp4L"
   },
   "outputs": [],
   "source": [
    "import agreement\n",
    "import classification_metrics\n",
    "import eval_utils\n",
    "import irn as aggregation\n",
    "import colab_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JSt2-5VpTp4M"
   },
   "outputs": [],
   "source": [
    "compute_rank1_certainties = jax.jit(eval_utils.rankk_certainties)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A-lDDL4a5nA8"
   },
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X3C-mDJ6Tp4M"
   },
   "outputs": [],
   "source": [
    "dataset = 'derm'  #@param ['toy', 'derm']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FPX6rSZHTp4N"
   },
   "outputs": [],
   "source": [
    "model_names = ['A', 'B', 'C', 'D']\n",
    "if dataset == 'toy':\n",
    "  with open('data/toy_data.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "  model_predictions = []\n",
    "  for i in range(4):\n",
    "    with open(f'data/toy_predictions{i}.pkl', 'rb') as f:\n",
    "      model_predictions.append(pickle.load(f))\n",
    "  model_predictions = jnp.array(model_predictions)\n",
    "  num_readers = 3\n",
    "  irn_plausibilities = aggregation.aggregate_irn(\n",
    "      data['test_rankings'][:, :num_readers],\n",
    "      data['test_groups'][:, :num_readers])\n",
    "  num_classes = 3\n",
    "elif dataset == 'derm':\n",
    "  with open('data/dermatology_data.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "  model_predictions = []\n",
    "  for i in range(4):\n",
    "    with open(f'data/dermatology_predictions{i}.txt', 'r') as f:\n",
    "      model_predictions.append(np.loadtxt(f))\n",
    "  model_predictions = jnp.array(model_predictions)\n",
    "  num_readers = 10\n",
    "  num_classes = 419\n",
    "  irn_plausibilities = data['test_irn']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2Nia5SyRTp4P"
   },
   "outputs": [],
   "source": [
    "num_samples = 100  # We used 1000 in the paper.\n",
    "prirn_plausibilities = []\n",
    "temperatures = [1, 3, 5, 10, 20, 30, 50]\n",
    "for temperature in temperatures:\n",
    "  plausibilities_t = aggregation.sample_prirn(\n",
    "    jax.random.PRNGKey(0), irn_plausibilities, num_samples=num_samples,\n",
    "    temperature=temperature, alpha=0.01)\n",
    "  prirn_plausibilities.append(plausibilities_t)\n",
    "  print(f'Computed PrIRN for temperature {temperature}.')\n",
    "prirn_plausibilities = jnp.array(prirn_plausibilities)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mihiWTSKTp4Q"
   },
   "source": [
    "# Certainty analysis and reader agreement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1F-YLiZHTp4R"
   },
   "outputs": [],
   "source": [
    "def compare_rank1_certainties(\n",
    "    plausibilities, names, **kwargs):\n",
    "  \"\"\"Plot rank-1 certainties for plausibilities across reliabilities.\"\"\"\n",
    "  num_models, num_examples, _, num_classes = plausibilities.shape\n",
    "  for m in range(num_models):\n",
    "    certainties = compute_rank1_certainties(\n",
    "        plausibilities[m], jnp.arange(num_classes))\n",
    "    certainties = jnp.max(certainties, axis=-1)\n",
    "    indices = jnp.argsort(certainties)\n",
    "    plt.plot(\n",
    "        jnp.arange(num_examples),\n",
    "        certainties[indices],\n",
    "        label=names[m])\n",
    "  plt.gcf().set_size_inches(\n",
    "      kwargs.get('width', 12), kwargs.get('height', 2.5))\n",
    "  plt.title(kwargs.get('title', f'Certainties across trust'))\n",
    "  plt.ylabel('Certainty')\n",
    "  plt.xlabel('Sorted examples')\n",
    "  plt.xlim(0, num_examples)\n",
    "  plt.ylim(0, 1)\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 2))\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fYer4FkiTp4R"
   },
   "outputs": [],
   "source": [
    "compare_rank1_certainties(\n",
    "    prirn_plausibilities, names=[f'temperature {m}' for m in temperatures],\n",
    "    width=7, title='PrIRN top-1 annotation certainties across temperatures')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NWSoiFgpTp4S"
   },
   "outputs": [],
   "source": [
    "def compute_coverage_agreement(rankings, groups):\n",
    "  \"\"\"Compute mean agreement using coverage against top-1 conditions.\"\"\"\n",
    "  agreements = agreement.leave_one_reader_out_coverage_agreement(\n",
    "      rankings, groups, jnp.array([10] * rankings.shape[0]))\n",
    "  return jnp.sum(agreements, axis=1) / 10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qXbM0UNjTp4S"
   },
   "outputs": [],
   "source": [
    "def plot_rank1_certainties_with_agreement(\n",
    "    agreements, plausibilities, **kwargs):\n",
    "  \"\"\"Plot rank-1 certaninties with mean reader agreement.\"\"\"\n",
    "  num_examples, _, num_classes = plausibilities.shape\n",
    "  certainties = compute_rank1_certainties(\n",
    "      plausibilities, jnp.arange(num_classes))\n",
    "  certainties = jnp.max(certainties, axis=1)\n",
    "  indices = np.argsort(certainties)\n",
    "  correlation = np.corrcoef(certainties, agreements)[0, 1]\n",
    "  plt.plot(\n",
    "      np.arange(num_examples),\n",
    "      certainties[indices],\n",
    "      label='Top-1 certainties')\n",
    "  plt.scatter(\n",
    "      np.arange(num_examples),\n",
    "      agreements[indices],\n",
    "      label='Agreements', s=4, c=colab_utils.COLOR_RED)\n",
    "  m, b = np.polyfit(np.arange(num_examples), agreements[indices], 1)\n",
    "  plt.plot(np.arange(num_examples), m*np.arange(num_examples)+b, color='gray',\n",
    "           label='Regression line')\n",
    "  plt.title(kwargs.get(\n",
    "      'title',\n",
    "      'Top-1 certainty and reader agreement '\n",
    "      f'(corr. {correlation:.2f})'))\n",
    "  plt.ylabel(kwargs.get('ylabel', 'Certainty / agreement'))\n",
    "  plt.xlabel('Sorted examples')\n",
    "  plt.legend()\n",
    "  plt.xlim(0, num_examples)\n",
    "  plt.ylim(0, 1)\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 2))\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ItbscGvvTp4T"
   },
   "outputs": [],
   "source": [
    "for i in range(len(temperatures)):\n",
    "  plot_rank1_certainties_with_agreement(\n",
    "      compute_coverage_agreement(data['test_rankings'], data['test_groups']),\n",
    "      prirn_plausibilities[i], width=7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BbHW83FbTp4T"
   },
   "source": [
    "# Model comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "j8X8HIfTTp4T"
   },
   "outputs": [],
   "source": [
    "def compute_ua_topk_accuracies(\n",
    "    predictions, plausibilities, k, break_ties=False):\n",
    "  \"\"\"Compute uncertainty-adjusted accuracies.\"\"\"\n",
    "  num_examples, _, num_classes = plausibilities.shape\n",
    "  if break_ties:\n",
    "    plausibilities += (jax.random.uniform(\n",
    "        jax.random.PRNGKey(0), plausibilities.shape) - 0.5) * 1e-4\n",
    "  labels = classification_metrics.topk_sets(\n",
    "      plausibilities.reshape(-1, num_classes),\n",
    "      k=1).reshape(num_examples, -1, num_classes)\n",
    "  return eval_utils.map_across_plausibilities(\n",
    "      predictions, labels,\n",
    "      functools.partial(classification_metrics.aggregated_topk_accuracy, k=k))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9jV8Ux8PTp4T"
   },
   "outputs": [],
   "source": [
    "def compare_rank1_certainties_with_ua_accuracies(\n",
    "    predictions, plausibilities, model_names, k=3, **kwargs):\n",
    "  \"\"\"Plot rank-1 certainty with uncertainty-adjusted accuracy.\"\"\"\n",
    "  num_models, _, _ = predictions.shape\n",
    "  num_examples, _, num_classes = plausibilities.shape\n",
    "  certainties = jnp.max(compute_rank1_certainties(\n",
    "      plausibilities, jnp.arange(num_classes)), axis=1)\n",
    "  for m in range(num_models):\n",
    "    accuracies = compute_ua_topk_accuracies(\n",
    "        predictions[m], plausibilities, k)\n",
    "    accuracies = jnp.mean(accuracies, axis=1)\n",
    "    indices = jnp.argsort(accuracies)\n",
    "    plt.plot(\n",
    "        jnp.arange(num_examples), accuracies[indices],\n",
    "        label=model_names[m])\n",
    "  plt.plot(\n",
    "      jnp.arange(num_examples), jnp.sort(certainties),\n",
    "      label='Top-1 certainties',color='gray', linestyle='dashed')\n",
    "  plt.title(f'UA top-{k} accuracy and certainty')\n",
    "  plt.ylabel('Certainty / correct')\n",
    "  plt.xlabel('Sorted examples')\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 2))\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pnaco_o8Tp4U"
   },
   "outputs": [],
   "source": [
    "models_to_compare = jnp.array([0, 2])\n",
    "compare_rank1_certainties_with_ua_accuracies(\n",
    "    model_predictions[models_to_compare], prirn_plausibilities[1],\n",
    "    [model_names[m] for m in models_to_compare], k=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ALv_qjEUTp4U"
   },
   "outputs": [],
   "source": [
    "def plot_ua_accuracies(\n",
    "    predictions,\n",
    "    irn_plausibilities, prirn_plausibilities,\n",
    "    k=3, **kwargs):\n",
    "  \"\"\"Plot uncertainty-adjusted top-k accuracy for different plausibilities.\"\"\"\n",
    "  irn_labels = jnp.argmax(irn_plausibilities, 1)\n",
    "  irn_accuracies = classification_metrics.aggregated_topk_accuracy(\n",
    "      predictions, jax.nn.one_hot(irn_labels, irn_plausibilities.shape[1]), k)\n",
    "  prirn_accuracies = compute_ua_topk_accuracies(\n",
    "      predictions, prirn_plausibilities, k)\n",
    "  prirn_hist, _ = colab_utils.plot_hist(\n",
    "      jnp.mean(prirn_accuracies, axis=0),\n",
    "      alpha=0.5, label='PrIRN accuracies',\n",
    "      color=colab_utils.COLORS[0])\n",
    "  hist_max = jnp.max(prirn_hist)\n",
    "  plt.vlines(\n",
    "      jnp.mean(prirn_accuracies), 0, hist_max,\n",
    "      label='PrIRN UA accuracy', color=colab_utils.COLORS[0])\n",
    "  plt.vlines(\n",
    "      jnp.mean(irn_accuracies), 0, hist_max,\n",
    "      label='IRN accuracy', color=colab_utils.COLORS[0], linestyle='dotted')\n",
    "  plt.legend(loc='upper left')\n",
    "  plt.ylabel('Counts')\n",
    "  plt.xlabel('Accuracy' if k == 1 else f'Top-{k} accuracy')\n",
    "  plt.title(kwargs.get('title', 'UA accuracy and certainty'))\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 2))\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vsBqV_psTp4U"
   },
   "outputs": [],
   "source": [
    "models_to_compare = jnp.array([0, 2])\n",
    "for m in models_to_compare:\n",
    "  plot_ua_accuracies(\n",
    "      model_predictions[m], irn_plausibilities,\n",
    "      prirn_plausibilities[3], k=1, title=model_names[m])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cxWXh4B9Tp4U"
   },
   "outputs": [],
   "source": [
    "def plot_model_comparison_with_certainty(\n",
    "    accuracy_fn, predictions, plausibilities,\n",
    "    temperatures, model_names, k=None,\n",
    "    num_samples=1000, **kwargs):\n",
    "  \"\"\"Compare models across reliabilities.\"\"\"\n",
    "  num_temperatures, _, _, _ = plausibilities.shape\n",
    "  assert len(temperatures) == num_temperatures\n",
    "  num_models = predictions.shape[0]\n",
    "  assert len(model_names) == num_models\n",
    "\n",
    "  vmax = 0\n",
    "  vmin = 1\n",
    "  ax = plt.gca()\n",
    "  for m in range(num_models):\n",
    "    accuracies_m = []\n",
    "    for i, _ in enumerate(temperatures):\n",
    "      accuracies_m_i = accuracy_fn(\n",
    "          predictions[m], plausibilities[i, :, :num_samples])\n",
    "      accuracies_m.append(accuracies_m_i)\n",
    "    # Before: num_temperatures x num_examples x num_samples\n",
    "    accuracies_m = jnp.array(accuracies_m)\n",
    "    mean_accuracies_m = jnp.mean(jnp.mean(accuracies_m, axis=1), axis=1)\n",
    "    std_accuracies_m = jnp.std(jnp.mean(accuracies_m, axis=1), axis=1)\n",
    "    max_accuracies_m = mean_accuracies_m + std_accuracies_m\n",
    "    min_accuracies_m = mean_accuracies_m - std_accuracies_m\n",
    "    ax.plot(\n",
    "        temperatures[:-1], mean_accuracies_m[:-1],\n",
    "        label=model_names[m], color=colab_utils.COLORS[m])\n",
    "    ax.fill_between(\n",
    "        temperatures[:-1], min_accuracies_m[:-1], max_accuracies_m[:-1],\n",
    "        alpha=0.1, color=colab_utils.COLORS[m])\n",
    "    ax.scatter(\n",
    "        temperatures[-1], mean_accuracies_m[-1],\n",
    "        s=25, marker='x', color=colab_utils.COLORS[m])\n",
    "    vmax = max(vmax, jnp.max(max_accuracies_m))\n",
    "    vmin = min(vmin, jnp.min(min_accuracies_m))\n",
    "\n",
    "  ax.vlines(\n",
    "      temperatures[-2], kwargs.get('ymin', vmin), kwargs.get('ymax', vmax + 0.005),\n",
    "      color='gray', linestyle='dotted')\n",
    "\n",
    "  ax.legend(loc='lower right', bbox_to_anchor=(0.85, 0.025))\n",
    "  plt.title(kwargs.get('title', f'Certainty and top-{k} accuracy'))\n",
    "  ax.set_xlabel(kwargs.get('xlabel', 'Repeated readers'))\n",
    "  ax.set_ylabel(kwargs.get('ylabel', 'Accuracy'))\n",
    "  ax.set_ylim(kwargs.get('ymin', vmin), kwargs.get('ymax', vmax + 0.005))\n",
    "  ax.set_xlim(\n",
    "      kwargs.get('xmin', min(temperatures)),\n",
    "      kwargs.get('xmax', max(temperatures)))\n",
    "  ax.set_xticks(\n",
    "      kwargs.get('xticks', []), kwargs.get('xticklabels', []))\n",
    "  ax.set_yticks(\n",
    "      kwargs.get('yticks', []), kwargs.get('yticklabels', None))\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 3.75), kwargs.get('height', 4))\n",
    "  plt.grid()\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4DxWDYofTp4V"
   },
   "outputs": [],
   "source": [
    "kwargs = dict(\n",
    "    temperatures=temperatures + [55], num_samples=10,\n",
    "    xlabel=None, ylabel=None, xticks=temperatures + [55],\n",
    "    xticklabels=['Low', '', '', 'Med', '', '', 'High', 'ML'],\n",
    "    model_names=model_names, yticks=[0.5, 0.6, 0.7, 0.8, 0.9, 1],\n",
    "    ymin=0.5, ymax=1,\n",
    ")\n",
    "prirn_plausibilities_with_irn = jnp.concatenate((\n",
    "    prirn_plausibilities,\n",
    "    jnp.repeat(irn_plausibilities.reshape(1, -1, 1, num_classes), num_samples, axis=2)\n",
    "), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S-nSMKjWTp4V"
   },
   "outputs": [],
   "source": [
    "top_k = 1 if num_classes == 3 else 3\n",
    "plot_model_comparison_with_certainty(\n",
    "    functools.partial(compute_ua_topk_accuracies, k=top_k),\n",
    "    model_predictions, prirn_plausibilities_with_irn,\n",
    "    title=f'Top-{top_k} UA accuracy', **kwargs)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}