colab_pl_sampler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h3YNDKfDUM8g"
   },
   "source": [
    "# PL Sampler\n",
    "\n",
    "See `README.md` for installation and usage instructions.\n",
    "\n",
    "This notebook allows to run the Plackett-Luce Gibbs sampler described in [2] on the toy 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": "markdown",
   "metadata": {
    "id": "50KMayy6UM8q"
   },
   "source": [
    "## Imports and setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Xhk5-JxEUM8r"
   },
   "outputs": [],
   "source": [
    "import functools\n",
    "import jax\n",
    "import jax.numpy as jnp\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": "6Y742AcAUM8v"
   },
   "outputs": [],
   "source": [
    "import selectors_utils\n",
    "import pl_samplers\n",
    "import eval_utils\n",
    "import colab_utils\n",
    "import gaussian_toy_dataset as gtd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HrD9rDboUM8w"
   },
   "outputs": [],
   "source": [
    "colab_utils.set_style()\n",
    "plot_hist = colab_utils.plot_hist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MtjFsjJ95tlR"
   },
   "source": [
    "## Utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "whlT8iKH5tlR"
   },
   "outputs": [],
   "source": [
    "def gibbs_sampler(\n",
    "    selector,\n",
    "    sampler,\n",
    "    shape_lam,\n",
    "    rate_lam,\n",
    "    num_classes,\n",
    "    num_iterations,\n",
    "    warmup_iterations,\n",
    "    rng,\n",
    "):\n",
    "  \"\"\"Helper to run standard Gibbs sampler on one case.\"\"\"\n",
    "  shape_lam = jnp.ones((num_classes)) * shape_lam / num_classes\n",
    "  rate_lam = jnp.ones((num_classes)) * rate_lam\n",
    "  result = sampler.sample(\n",
    "      next(rng),\n",
    "      selector,\n",
    "      shape_lam=shape_lam,\n",
    "      rate_lam=rate_lam,\n",
    "      num_iterations=num_iterations)\n",
    "  result = result[warmup_iterations:]\n",
    "  return jnp.expand_dims(result, axis=0)\n",
    "\n",
    "\n",
    "def gibbs_sampler_from_ranked_classes(\n",
    "    selector,\n",
    "    sampler,\n",
    "    shape_lam,\n",
    "    rate_lam,\n",
    "    num_iterations,\n",
    "    warmup_iterations,\n",
    "    represent_unranked_classes,\n",
    "    normalize_unranked_equally,\n",
    "    num_classes,\n",
    "    rng,\n",
    "):\n",
    "  \"\"\"Helper to run Gibbs sampler on one case while ignoring unranked classes.\"\"\"\n",
    "  num_observed_classes = len(list(set(jax.tree_util.tree_leaves(selector))))\n",
    "  shape_lam /= (num_observed_classes + 1)\n",
    "  if not normalize_unranked_equally:\n",
    "    shape_lam = 1\n",
    "  result = sampler.sample_from_ranked_classes(\n",
    "      next(rng),\n",
    "      selector,\n",
    "      shape_lam_i=shape_lam,\n",
    "      rate_lam_i=rate_lam,\n",
    "      num_classes=num_classes,\n",
    "      num_iterations=num_iterations,\n",
    "      represent_unranked_classes=represent_unranked_classes,\n",
    "      normalize_unranked_equally=normalize_unranked_equally)\n",
    "  result = result[warmup_iterations:]\n",
    "  return jnp.expand_dims(result, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cm3X8RpM5tlS"
   },
   "outputs": [],
   "source": [
    "def plot_plausabilities(\n",
    "    plausabilities, irn_reference, label_names,\n",
    "    limit=10, num_samples=100, **kwargs):\n",
    "  \"\"\"Plot PL plausabilities against reference IRN plausabilities.\"\"\"\n",
    "  indices = jnp.argsort(- irn_reference)\n",
    "  indices = indices[:limit]\n",
    "  plt.bar(\n",
    "      jnp.arange(indices.shape[0]),\n",
    "      irn_reference[indices], alpha=0.5,\n",
    "      label=kwargs.get('reference_label', ''), color=colab_utils.COLORS[0])\n",
    "  if plausabilities is not None:\n",
    "    for s in range(num_samples):\n",
    "      label = f'{num_samples} Samples' if s == 0 else ''\n",
    "      plt.scatter(\n",
    "          jnp.arange(indices.shape[0]),\n",
    "          plausabilities[s, indices],\n",
    "          c='g', s=10, alpha=0.75, label=label)\n",
    "    plt.scatter(\n",
    "        jnp.arange(indices.shape[0]),\n",
    "        jnp.mean(plausabilities[:num_samples], axis=0)[indices],\n",
    "        s=25, c='r', label='Mean')\n",
    "\n",
    "  plt.xticks(np.arange(indices.shape[0]), [label_names[i][:6] for i in indices])\n",
    "  plt.title(kwargs.get('title', f'Plausabilities'))\n",
    "  plt.xlabel(kwargs.get('xlabel', 'Classes'))\n",
    "  plt.ylabel(kwargs.get('ylabel', ''))\n",
    "  plt.ylim(ymin=kwargs.get('ymin', None), ymax=kwargs.get('ymax', None))\n",
    "  if kwargs.get('legend', False):\n",
    "    plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', limit), kwargs.get('height', 3))\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dTAxvyF55tlS"
   },
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TCHyDnrZ5tlS"
   },
   "outputs": [],
   "source": [
    "dataset = 'derm'  #@param ['toy', 'derm']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WEjzB_ko5tlS"
   },
   "outputs": [],
   "source": [
    "if dataset == 'toy':\n",
    "  with open('data/toy_data.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "  label_names = ['0', '1', '2']\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_conditions.txt', 'r') as f:\n",
    "    label_names = [condition.strip() for condition in f.readlines()]\n",
    "  indices = [1026, 1057, 357]  # Examples from the paper."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_unOebHP5tlS"
   },
   "source": [
    "## PL Sampling\n",
    "\n",
    "For the Plackett-Luce Gibbs sampler, we provide several options. Most notably, `sample_from_ranked` can be set to `True` if there are many classes but individual annotations will never include all classes. This is true on the Dermatology dataset, but not on the toy dataset. So setting this to `True` will speed up sampling on the Dermatology data.\n",
    "\n",
    "For the 3 dermatology cases, the below should run in a few minutes. Running on the whole dataset might take considerably longer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i-IIdeU35tlT"
   },
   "outputs": [],
   "source": [
    "shape_lam = 1.0\n",
    "rate_lam = 1.0\n",
    "warmup_iterations = 500\n",
    "total_iterations = 1500\n",
    "reader_repetitions = 3\n",
    "# sample_from_ranked can be set to True if there are many classes and\n",
    "# the partial rankings will never include all classes.\n",
    "sample_from_ranked = (dataset == 'derm')\n",
    "represent_unranked_classes = True\n",
    "normalize_unranked_equally = False\n",
    "jit_strategy = 'jit_per_reader'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u0P6FJbi5tlT"
   },
   "outputs": [],
   "source": [
    "rng = gtd.PRNGSequence(0)\n",
    "num_classes = data['test_irn'].shape[1]\n",
    "plausibilities = []\n",
    "for index in indices:\n",
    "  selector = data['test_selectors'][index]\n",
    "\n",
    "  if reader_repetitions > 1:\n",
    "    selector = selectors_utils.repeat_selectors(\n",
    "        [selector], reader_repetitions\n",
    "    )[0]\n",
    "\n",
    "  if sample_from_ranked:\n",
    "    worker_fn = functools.partial(\n",
    "        gibbs_sampler_from_ranked_classes,\n",
    "        represent_unranked_classes=represent_unranked_classes,\n",
    "        normalize_unranked_equally=normalize_unranked_equally,\n",
    "    )\n",
    "  else:\n",
    "    worker_fn = gibbs_sampler\n",
    "  worker = functools.partial(\n",
    "      worker_fn,\n",
    "      sampler=pl_samplers.GibbsSamplerPlackettLuce(jit_strategy),\n",
    "      shape_lam=shape_lam,\n",
    "      rate_lam=rate_lam,\n",
    "      num_iterations=total_iterations,\n",
    "      warmup_iterations=warmup_iterations,\n",
    "      num_classes=num_classes,\n",
    "      rng=rng)\n",
    "\n",
    "  plausibilities.append(worker(selector))\n",
    "plausibilities = jnp.concatenate(plausibilities, axis=0)\n",
    "plausibilities = eval_utils.normalize_plausibilities(plausibilities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5KsdYkTP5tlT"
   },
   "outputs": [],
   "source": [
    "for i, index in enumerate(indices):\n",
    "  plot_plausabilities(plausibilities[i], data['test_irn'][index], label_names)"
   ]
  }
 ],
 "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
}