colab_partial_ao.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yuukhcFxUC6Z"
   },
   "source": [
    "# Partial average overlap\n",
    "\n",
    "See `README.md` for installation and usage instructions.\n",
    "\n",
    "This notebook shows an example of using the partial average overlap\n",
    "implementation of [2] on the dataset of [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": "code",
   "execution_count": null,
   "metadata": {
    "id": "EY62HMyaUC6l"
   },
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import matplotlib\n",
    "import os\n",
    "import pickle\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "skKpu2ySUC6q"
   },
   "outputs": [],
   "source": [
    "import ranking_metrics\n",
    "import colab_utils\n",
    "import formats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WRMW7Eh85vj6"
   },
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y7ptcO3jUC6s"
   },
   "outputs": [],
   "source": [
    "dataset = 'derm'  #@param ['toy', 'derm']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zsd2QjwGUC6t"
   },
   "outputs": [],
   "source": [
    "if dataset == 'toy':\n",
    "  with open('data/toy_data.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "  with open(f'data/toy_predictions0.pkl', 'rb') as f:\n",
    "    model_predictions = pickle.load(f)\n",
    "  indices = [0, 1, 2]  # Random examples.\n",
    "elif dataset == 'derm':\n",
    "  with open('data/dermatology_data.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "  with open('data/dermatology_predictions0.txt', 'rb') as f:\n",
    "    model_predictions = np.loadtxt(f)\n",
    "  indices = [1026, 1057, 357]  # Examples from the paper."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "21CIE9aa5vj7"
   },
   "source": [
    "## Compute (partial) average overlap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xOzsEWqYUC6v"
   },
   "outputs": [],
   "source": [
    "indices = np.array(indices, dtype=int)\n",
    "num_examples = indices.size\n",
    "model_rankings = jnp.argsort(- model_predictions, axis=1)\n",
    "model_groups = jnp.array([jnp.arange(model_rankings.shape[1]) for _ in model_rankings])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KW4yQbFTUC6w"
   },
   "outputs": [],
   "source": [
    "for i, index in enumerate(indices):\n",
    "  for j, selector in enumerate(data['test_selectors'][index]):\n",
    "    print('Example', i, ', annotation', j, ':', data['test_selectors'][index][j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TxT7sUcA5vj8"
   },
   "outputs": [],
   "source": [
    "model_selectors = formats.convert_rankings_to_selectors(\n",
    "    model_rankings[indices].reshape(num_examples, 1, -1),\n",
    "    model_groups[indices].reshape(num_examples, 1, -1))\n",
    "for i, selector in enumerate(model_selectors):\n",
    "  print('Example', i, ', model prediction: ', selector[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2Lra56ZGUC62"
   },
   "outputs": [],
   "source": [
    "ranking_metrics.average_overlap(\n",
    "    model_rankings[indices],\n",
    "    data['test_rankings'][indices, 0],\n",
    "    jnp.sum(data['test_groups'][indices, 0] >= 0, axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yy-GTJHgUC64"
   },
   "outputs": [],
   "source": [
    "ranking_metrics.partial_average_overlap(\n",
    "    model_rankings[indices],\n",
    "    model_groups[indices],\n",
    "    data['test_rankings'][indices, 0],\n",
    "    data['test_groups'][indices, 0],\n",
    "    jnp.sum(data['test_groups'][indices, 0] >= 0, axis=1))"
   ]
  }
 ],
 "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
}