colab_mccp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4y3zf-xKT_xG"
   },
   "source": [
    "# Monte Carlo conformal prediction\n",
    "\n",
    "See `README.md` for installation and usage instructions.\n",
    "\n",
    "This notebook re-creates some of the examples and figures from [1] on the toy dataset and results of [1] on the dermatology data.\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",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "spNS_87xT_xW"
   },
   "source": [
    "## Imports and setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Uct1UkQYT_xY"
   },
   "outputs": [],
   "source": [
    "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": "dPPlUEudT_xc"
   },
   "outputs": [],
   "source": [
    "import conformal_prediction\n",
    "import monte_carlo\n",
    "import p_value_combination\n",
    "import plausibility_regions\n",
    "import classification_metrics\n",
    "import colab_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6RuPqQ0FT_xd"
   },
   "outputs": [],
   "source": [
    "colab_utils.set_style()\n",
    "plot_hist = colab_utils.plot_hist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-9F-EYei508g"
   },
   "source": [
    "## Toy dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "msdEDh4BT_xe"
   },
   "outputs": [],
   "source": [
    "with open('data/toy_data.pkl', 'rb') as f:\n",
    "  data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U86LZsIVT_xg"
   },
   "outputs": [],
   "source": [
    "with open('data/toy_predictions0.pkl', 'rb') as f:\n",
    "  predictions = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4fWLDNwyT_xi"
   },
   "source": [
    "### CP with true and top-1 labels\n",
    "\n",
    "Compare standard conformal prediction calibrated against majority voted (top-1) labels and conformal prediction calibrated against the true labels -- which we have access to on the toy dataset but not in practice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bmuS4FaqT_xj"
   },
   "outputs": [],
   "source": [
    "def run_trial(\n",
    "    predictions, data, method='',\n",
    "    alpha=0.1, trials=100, split=0.5):\n",
    "  \"\"\"Run a conformal prediction experiment.\"\"\"\n",
    "  results = {}\n",
    "  tags = ['inefficiencies', 'true_coverages', 'aggregated_coverages', 'top1_coverages']\n",
    "  for tag in tags:\n",
    "    results[tag] = []\n",
    "  keys = jax.random.split(jax.random.PRNGKey(0), 2 * trials)\n",
    "  for t in range(trials):\n",
    "    permutation = jax.random.permutation(keys[2 * t], predictions.shape[0])\n",
    "    val_examples = int(predictions.shape[0]*split)\n",
    "    val_predictions = predictions[permutation[:val_examples]]\n",
    "    test_predictions = predictions[permutation[val_examples:]]\n",
    "    val_human_ground_truth = data['test_smooth_labels'][permutation[:val_examples]]\n",
    "    test_human_ground_truth = data['test_smooth_labels'][permutation[val_examples:]]\n",
    "    val_ground_truth = data['test_labels'][permutation[:val_examples]]\n",
    "    test_ground_truth = data['test_labels'][permutation[val_examples:]]\n",
    "\n",
    "    if method == 'top1':\n",
    "      val_labels = jnp.argmax(val_human_ground_truth, axis=1)\n",
    "      threshold = conformal_prediction.calibrate_threshold(\n",
    "          val_predictions, val_labels, alpha)\n",
    "    elif method == 'mccp':\n",
    "      threshold = monte_carlo.calibrate_mc_threshold(\n",
    "          keys[2 * t + 1], val_predictions, val_human_ground_truth, num_samples=10, alpha=alpha)\n",
    "    else:\n",
    "      val_labels = val_ground_truth\n",
    "      threshold = conformal_prediction.calibrate_threshold(\n",
    "          val_predictions, val_labels, alpha)\n",
    "\n",
    "    confidence_sets = conformal_prediction.predict_threshold(\n",
    "        test_predictions, threshold)\n",
    "\n",
    "    results['inefficiencies'].append(classification_metrics.size(\n",
    "        confidence_sets))\n",
    "    test_one_hot_ground_truth = jax.nn.one_hot(\n",
    "        test_ground_truth, confidence_sets.shape[1])\n",
    "    results['true_coverages'].append(classification_metrics.aggregated_coverage(\n",
    "        confidence_sets, test_one_hot_ground_truth))\n",
    "    test_top1_ground_truth = jax.nn.one_hot(\n",
    "        jnp.argmax(test_human_ground_truth, axis=1),\n",
    "        test_human_ground_truth.shape[1])\n",
    "    results['top1_coverages'].append(\n",
    "        classification_metrics.aggregated_coverage(\n",
    "            confidence_sets, test_top1_ground_truth))\n",
    "    results['aggregated_coverages'].append(\n",
    "        classification_metrics.aggregated_coverage(\n",
    "            confidence_sets, test_human_ground_truth))\n",
    "  for tag in tags:\n",
    "    results[tag] = jnp.array(results[tag])\n",
    "  return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ebWU-rxcT_xp"
   },
   "outputs": [],
   "source": [
    "def plot_coverage_top1_calibration(predictions, data, **kwargs):\n",
    "  \"\"\"Plot coverage when calibrating against true or top-1 labels.\"\"\"\n",
    "  top1_results = run_trial(\n",
    "      predictions, data, method='top1',\n",
    "      alpha=0.05, trials=100)\n",
    "  plot_hist(jnp.mean(top1_results['true_coverages'], axis=-1), normalize=True, label='Coverage of true labels')\n",
    "  plot_hist(jnp.mean(top1_results['top1_coverages'], axis=-1), normalize=True, label='Coverage of voted labels')\n",
    "  plt.vlines(0.95, 0, 0.15, color='black', label='Target')\n",
    "  plt.title('Calibration with voted labels')\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  if kwargs.get('name', False):\n",
    "      plt.savefig(kwargs.get('name') + '.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SsWg5WiCT_xq"
   },
   "outputs": [],
   "source": [
    "plot_coverage_top1_calibration(\n",
    "    predictions, data, name='coverage_top1_calibration')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0cBMsRcCT_xr"
   },
   "outputs": [],
   "source": [
    "def plot_coverage_mccp_calibration(predictions, data, **kwargs):\n",
    "  \"\"\"Plot coverage when calibrating against true or top-1 labels.\"\"\"\n",
    "  top1_results = run_trial(\n",
    "      predictions, data, method='mccp',\n",
    "      alpha=0.05, trials=100)\n",
    "  plot_hist(jnp.mean(top1_results['true_coverages'], axis=-1), normalize=True, label='Coverage of true labels')\n",
    "  plot_hist(jnp.mean(top1_results['top1_coverages'], axis=-1), normalize=True, label='Coverage of voted labels')\n",
    "  plt.vlines(0.95, 0, 0.15, color='black', label='Target')\n",
    "  plt.title('Monte Carlo conformal calibration')\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  if kwargs.get('name', False):\n",
    "      plt.savefig(kwargs.get('name') + '.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3aaKeilfT_xs"
   },
   "outputs": [],
   "source": [
    "plot_coverage_mccp_calibration(\n",
    "    predictions, data, name='coverage_mccp_calibration')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SKvzLyX1T_xs"
   },
   "source": [
    "### Aggregated coverage\n",
    "\n",
    "This experiment illustrates aggregated coverage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jqh5VCeMT_xs"
   },
   "outputs": [],
   "source": [
    "def plot_aggregated_coverage(predictions, data, **kwargs):\n",
    "  true_results = run_trial(\n",
    "      predictions, data, method='',\n",
    "      alpha=0.05, trials=500)\n",
    "\n",
    "  plot_hist(\n",
    "      jnp.mean(true_results['true_coverages'], axis=1), bins=40,\n",
    "      normalize=True, label='True coverage', alpha=0.8)\n",
    "  plot_hist(\n",
    "      jnp.mean(true_results['aggregated_coverages'], axis=1), bins=40,\n",
    "      normalize=True, label='Aggregated coverage', alpha=0.8)\n",
    "  plt.title('Average coverage across trials')\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.xticks([0.94, 0.945, 0.95, 0.955, 0.96])\n",
    "  plt.xlim(0.94, 0.96)\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('aggregated_coverage_histogram1.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()\n",
    "\n",
    "  plt.plot(\n",
    "      jnp.arange(true_results['true_coverages'][0].shape[0]),\n",
    "      jnp.sort(true_results['true_coverages'][0]),\n",
    "      label='True label covered')\n",
    "  plt.fill_between(\n",
    "      jnp.arange(true_results['true_coverages'][0].shape[0]),\n",
    "      jnp.zeros(true_results['true_coverages'][0].shape[0]),\n",
    "      jnp.sort(true_results['true_coverages'][0]),\n",
    "      alpha=0.2, color=plt.rcParams['axes.prop_cycle'].by_key()['color'][0],\n",
    "      label=(\n",
    "          'Realized coverage'\n",
    "          f'({jnp.mean(true_results[\"true_coverages\"][0]):.2f})'))\n",
    "  plt.plot(\n",
    "      jnp.arange(true_results['aggregated_coverages'][0].shape[0]),\n",
    "      jnp.sort(true_results['aggregated_coverages'][0]),\n",
    "      label='Plausibility mass covered')\n",
    "  plt.fill_between(\n",
    "      jnp.arange(true_results['aggregated_coverages'][0].shape[0]),\n",
    "      jnp.zeros(true_results['aggregated_coverages'][0].shape[0]),\n",
    "      jnp.sort(true_results['aggregated_coverages'][0]),\n",
    "      alpha=0.2, color=plt.rcParams['axes.prop_cycle'].by_key()['color'][1],\n",
    "      label=(\n",
    "          'Realized aggregated coverage'\n",
    "          f'({jnp.mean(true_results[\"aggregated_coverages\"][0]):.2f})'))\n",
    "  plt.title('Correctness across examples for single trial')\n",
    "  plt.xlabel('Sorted examples (separate sorting)')\n",
    "  plt.ylabel('Realized coverage')\n",
    "  plt.xlim(0, true_results['aggregated_coverages'][0].shape[0])\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('aggregated_coverage_sorted.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D5NGXsBJT_xt"
   },
   "outputs": [],
   "source": [
    "plot_aggregated_coverage(predictions, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zIj6snZFT_xt"
   },
   "source": [
    "### Monte Carlo conformal prediction\n",
    "\n",
    "Running and evaluation Monte Carlo conformal prediction in comparison to calibrating against the majority-voted (top-1) labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BswmRUq7T_xu"
   },
   "outputs": [],
   "source": [
    "def run_trial(\n",
    "    predictions, data,\n",
    "    alpha=0.05, num_trials=1000, split=0.5, seed=0):\n",
    "  \"\"\"Run a conformal prediction experiment.\"\"\"\n",
    "  results = {}\n",
    "  metrics = ['true_coverages', 'aggregated_coverages']\n",
    "  methods = ['top1_', 'mc_']\n",
    "  for method in methods:\n",
    "    for metric in metrics:\n",
    "      results[method + metric] = []\n",
    "\n",
    "  permutation = jax.random.permutation(\n",
    "      jax.random.PRNGKey(seed), predictions.shape[0])\n",
    "  val_examples = int(predictions.shape[0]*split)\n",
    "  val_predictions = predictions[permutation[:val_examples]]\n",
    "  test_predictions = predictions[permutation[val_examples:]]\n",
    "  val_human_ground_truth = data['test_smooth_labels'][permutation[:val_examples]]\n",
    "  test_human_ground_truth = data['test_smooth_labels'][permutation[val_examples:]]\n",
    "  test_ground_truth = data['test_labels'][permutation[val_examples:]]\n",
    "\n",
    "  keys = jax.random.split(jax.random.PRNGKey(seed + 1), num_trials)\n",
    "  for t in range(num_trials):\n",
    "\n",
    "    def evaluate_method(confidence_sets, key):\n",
    "      test_one_hot_ground_truth = jax.nn.one_hot(\n",
    "          test_ground_truth, confidence_sets.shape[1])\n",
    "      results[f'{key}true_coverages'].append(classification_metrics.aggregated_coverage(\n",
    "          confidence_sets, test_one_hot_ground_truth))\n",
    "      results[f'{key}aggregated_coverages'].append(\n",
    "          classification_metrics.aggregated_coverage(\n",
    "              confidence_sets, test_human_ground_truth))\n",
    "\n",
    "    val_top1_labels = jnp.argmax(val_human_ground_truth, axis=1)\n",
    "    top1_threshold = conformal_prediction.calibrate_threshold(\n",
    "        val_predictions, val_top1_labels, alpha)\n",
    "    top1_confidence_sets = conformal_prediction.predict_threshold(\n",
    "          test_predictions, top1_threshold)\n",
    "    evaluate_method(top1_confidence_sets, 'top1_')\n",
    "\n",
    "    val_mc_predictions, mc_labels = monte_carlo.sample_mc_labels(\n",
    "      keys[t], val_predictions, val_human_ground_truth, num_samples=1)\n",
    "    val_mc_predictions = val_mc_predictions.reshape(\n",
    "        -1, val_mc_predictions.shape[-1])\n",
    "    val_mc_labels = mc_labels.reshape(-1)\n",
    "    mc_threshold = conformal_prediction.calibrate_threshold(\n",
    "        val_mc_predictions, val_mc_labels, alpha)\n",
    "    mc_confidence_sets = conformal_prediction.predict_threshold(\n",
    "        test_predictions, mc_threshold)\n",
    "    evaluate_method(mc_confidence_sets, f'mc_')\n",
    "\n",
    "  for key in results.keys():\n",
    "    results[key] = jnp.array(results[key])\n",
    "  return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HE2QyVG6T_xu"
   },
   "outputs": [],
   "source": [
    "def plot_label_randomness(predictions, data, **kwargs):\n",
    "  alpha = 0.05\n",
    "  results = run_trial(predictions, data, alpha, seed=kwargs.get('seed', 0))\n",
    "  vmax = 0\n",
    "  hist, _ = plot_hist(\n",
    "      jnp.mean(results['mc_aggregated_coverages'], axis=-1), normalize=True,\n",
    "      alpha=0.65, label='Aggregated coverage')\n",
    "  vmax = max(vmax, jnp.max(hist))\n",
    "  plt.vlines(1 - alpha, 0, vmax, color='black', label='Target')\n",
    "  plt.title('Variation in aggregated coverage from sampled labels')\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  if kwargs.get('name', False):\n",
    "      plt.savefig(kwargs.get('name') + '.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mc5wV7m3T_xv"
   },
   "outputs": [],
   "source": [
    "plot_label_randomness(predictions, data, seed=3, name=f'label_randomness3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u5IpwajOT_xv"
   },
   "outputs": [],
   "source": [
    "def run_trial(\n",
    "    predictions, data, alpha=0.05, num_trials=100, split=0.5):\n",
    "  \"\"\"Run a conformal prediction experiment.\"\"\"\n",
    "  results = {}\n",
    "  ms = [1, 5, 10]\n",
    "  metrics = [\n",
    "      'inefficiencies',\n",
    "      'true_coverages',\n",
    "      'aggregated_coverages',\n",
    "      'top1_coverages',\n",
    "  ]\n",
    "  methods = ['true_', 'top1_'] + [f'mc{m}_' for m in ms]\n",
    "  for method in methods:\n",
    "    for metric in metrics:\n",
    "      results[method + metric] = []\n",
    "  keys = jax.random.split(jax.random.PRNGKey(0), 2 * num_trials)\n",
    "  for t in range(num_trials):\n",
    "    permutation = jax.random.permutation(keys[2 * t], predictions.shape[0])\n",
    "    val_examples = int(predictions.shape[0]*split)\n",
    "    val_predictions = predictions[permutation[:val_examples]]\n",
    "    test_predictions = predictions[permutation[val_examples:]]\n",
    "    val_human_ground_truth = data['test_smooth_labels'][permutation[:val_examples]]\n",
    "    test_human_ground_truth = data['test_smooth_labels'][permutation[val_examples:]]\n",
    "    val_ground_truth = data['test_labels'][permutation[:val_examples]]\n",
    "    test_ground_truth = data['test_labels'][permutation[val_examples:]]\n",
    "\n",
    "    def evaluate_method(confidence_sets, key):\n",
    "      results[f'{key}inefficiencies'].append(classification_metrics.size(\n",
    "          confidence_sets))\n",
    "      test_one_hot_ground_truth = jax.nn.one_hot(\n",
    "          test_ground_truth, confidence_sets.shape[1])\n",
    "      results[f'{key}true_coverages'].append(classification_metrics.aggregated_coverage(\n",
    "          confidence_sets, test_one_hot_ground_truth))\n",
    "      test_top1_ground_truth = jax.nn.one_hot(\n",
    "          jnp.argmax(test_human_ground_truth, axis=1),\n",
    "          test_human_ground_truth.shape[1])\n",
    "      results[f'{key}top1_coverages'].append(\n",
    "          classification_metrics.aggregated_coverage(\n",
    "              confidence_sets, test_top1_ground_truth))\n",
    "      results[f'{key}aggregated_coverages'].append(\n",
    "          classification_metrics.aggregated_coverage(\n",
    "              confidence_sets, test_human_ground_truth))\n",
    "\n",
    "    val_true_labels = val_ground_truth\n",
    "    val_top1_labels = jnp.argmax(val_human_ground_truth, axis=1)\n",
    "    true_threshold = conformal_prediction.calibrate_threshold(\n",
    "        val_predictions, val_true_labels, alpha)\n",
    "    top1_threshold = conformal_prediction.calibrate_threshold(\n",
    "        val_predictions, val_top1_labels, alpha)\n",
    "    true_confidence_sets = conformal_prediction.predict_threshold(\n",
    "          test_predictions, true_threshold)\n",
    "    top1_confidence_sets = conformal_prediction.predict_threshold(\n",
    "          test_predictions, top1_threshold)\n",
    "    evaluate_method(true_confidence_sets, 'true_')\n",
    "    evaluate_method(top1_confidence_sets, 'top1_')\n",
    "\n",
    "    for m in ms:\n",
    "      val_mc_predictions, mc_labels = monte_carlo.sample_mc_labels(\n",
    "        keys[2 * t + 1], val_predictions, val_human_ground_truth, num_samples=m)\n",
    "      val_mc_predictions = val_mc_predictions.reshape(-1, val_mc_predictions.shape[-1])\n",
    "      val_mc_labels = mc_labels.reshape(-1)\n",
    "      mc_threshold = conformal_prediction.calibrate_threshold(\n",
    "          val_mc_predictions, val_mc_labels, alpha)\n",
    "      mc_confidence_sets = conformal_prediction.predict_threshold(\n",
    "          test_predictions, mc_threshold)\n",
    "      evaluate_method(mc_confidence_sets, f'mc{m}_')\n",
    "\n",
    "  for key in results.keys():\n",
    "    results[key] = jnp.array(results[key])\n",
    "  return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xaV71KfTT_xw"
   },
   "outputs": [],
   "source": [
    "def plot_std(alpha=0.05, num_trials=100, ms=[1, 5, 10], **kwargs):\n",
    "  \"\"\"Plot standard deviation for MC conformal prediction.\"\"\"\n",
    "  results = []\n",
    "  splits = [0.05, 0.1, 0.2, 0.3, 0.4, 0.5]\n",
    "  for split in splits:\n",
    "    results.append(run_trial(\n",
    "        predictions, data, alpha=0.05, num_trials=num_trials, split=split))\n",
    "\n",
    "  for m in ms:\n",
    "    values = [result[f'mc{m}_aggregated_coverages'] for result in results]\n",
    "    plt.plot(\n",
    "        splits,\n",
    "         [jnp.std(jnp.mean(value, axis=-1), axis=-1) for value in values],\n",
    "        label=r\"$m =$\" + f'{m} sampled labels')\n",
    "  plt.title('Standard deviation in aggregated coverage')\n",
    "  plt.xlabel('Fraction of calibration data')\n",
    "  plt.ylabel('Std in empirical coverage')\n",
    "  plt.legend(bbox_to_anchor=[1, 1], loc='upper right')\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('mc_std_aggregated_lines_wo.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()\n",
    "\n",
    "  vmax = 0\n",
    "  for m in ms:\n",
    "    hist, _ = plot_hist(\n",
    "        jnp.mean(results[0][f'mc{m}_aggregated_coverages'], axis=-1), normalize=True,\n",
    "        alpha=0.65, label=r\"$m =$\" + f'{m} sampled labels', bins=40, range=(0.92, 0.98))\n",
    "    vmax = max(vmax, jnp.max(hist))\n",
    "  plt.vlines(1 - alpha, 0, vmax, color='black', label='Target')\n",
    "  plt.title('Aggregated coverage for 10% calibration data')\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.xlim(0.9, 0.98)\n",
    "  plt.legend(bbox_to_anchor=[0, 1], loc='upper left')\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('mc_std_aggregated_histogram_wo.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zBhfN7QaT_xw"
   },
   "outputs": [],
   "source": [
    "plot_std(alpha=0.05, num_trials=500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_dE_EnXeT_xw"
   },
   "source": [
    "### ECDF-corrected Monte Carlo conformal prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Os_6Bo2DT_xx"
   },
   "outputs": [],
   "source": [
    "def plot_ecdf(**kwargs):\n",
    "  num_tests = 10\n",
    "  num_examples = 10000\n",
    "  val_examples = 10000//2\n",
    "  all_p_values = jax.random.uniform(jax.random.PRNGKey(0), (num_tests, num_examples))\n",
    "  dependent_p_values = all_p_values\n",
    "  all_p_values = jnp.concatenate((all_p_values, dependent_p_values), axis=0)\n",
    "  combined_p_values = jnp.mean(all_p_values, axis=0)\n",
    "  val_combined_p_values = combined_p_values[:val_examples]\n",
    "  test_combined_p_values = combined_p_values[val_examples:]\n",
    "  test_corrected_p_values = p_value_combination.combine_ecdf_p_values(\n",
    "      val_combined_p_values, test_combined_p_values)\n",
    "\n",
    "  plot_hist(\n",
    "      test_combined_p_values, normalize=True,\n",
    "       alpha=0.65, label='Averaged p-values')\n",
    "  plot_hist(\n",
    "      test_corrected_p_values, normalize=True,\n",
    "      alpha=0.65, label=f'ECDF corrected p-values')\n",
    "  plt.legend()\n",
    "  plt.title(f'Distribution of combined p-values')\n",
    "  plt.xlabel('Value')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('ecdf_p_values.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()\n",
    "\n",
    "  baseline_coverages = []\n",
    "  method_coverages = []\n",
    "  max_method_coverages = []\n",
    "  min_method_coverages = []\n",
    "  delta = 0.0001\n",
    "  alphas = jnp.linspace(0, 1, 26)\n",
    "  epsilon = np.sqrt(np.log(2. / delta) / (2 * val_examples))\n",
    "  for alpha in alphas:\n",
    "    baseline_coverages.append(jnp.mean(test_combined_p_values >= alpha))\n",
    "    method_coverage = jnp.mean(test_corrected_p_values >= alpha)\n",
    "    method_coverages.append(method_coverage)\n",
    "    max_method_coverages.append(min(1, method_coverage + epsilon))\n",
    "    min_method_coverages.append(max(0, method_coverage - epsilon))\n",
    "  max_method_coverages = np.array(max_method_coverages)\n",
    "  min_method_coverages = np.array(min_method_coverages)\n",
    "\n",
    "  plt.plot(1 - alphas, baseline_coverages, label='Baseline')\n",
    "  plt.plot(1 - alphas, method_coverages, label=f'ECDF', color='green')\n",
    "  plt.plot(1 - alphas, max_method_coverages, color='green', alpha=0.2)\n",
    "  plt.plot(1 - alphas, min_method_coverages, color='green', alpha=0.2)\n",
    "  plt.fill_between(\n",
    "      alphas, 1 - max_method_coverages, 1 - min_method_coverages, alpha=0.1,\n",
    "      color='green', label='ECDF $1 - \\delta$ band')\n",
    "  plt.title(f'ECDF correction of p-values')\n",
    "  plt.xlabel(r'Target coverage $1 - \\alpha$')\n",
    "  plt.ylabel('Empirical coverage')\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('ecdf_lines.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XMlnlOFaT_xx"
   },
   "outputs": [],
   "source": [
    "plot_ecdf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GxSQENifT_xy"
   },
   "outputs": [],
   "source": [
    "def plot_coverage_ecdf_mccp(\n",
    "    predictions, data,\n",
    "    num_samples=10, split=0.5, alpha=0.05, num_trials=10, **kwargs):\n",
    "  \"\"\"Plot coverage for ECDF based MC conformal prediction.\"\"\"\n",
    "  coverages = []\n",
    "  mc_coverages = []\n",
    "  corrected_mc_coverages = []\n",
    "  rng = jax.random.PRNGKey(0)\n",
    "  for _ in range(num_trials):\n",
    "    permutation_rng, mc_rng, rng = jax.random.split(rng, 3)\n",
    "    permutation = jax.random.permutation(permutation_rng, predictions.shape[0])\n",
    "    _, num_classes = predictions.shape\n",
    "    val_examples = int(predictions.shape[0]*split)\n",
    "    val_predictions = predictions[permutation[:val_examples]]\n",
    "    test_predictions = predictions[permutation[val_examples:]]\n",
    "    val_human_ground_truth = data['test_smooth_labels'][permutation[:val_examples]]\n",
    "    val_labels = data['test_labels'][permutation[:val_examples]]\n",
    "    test_labels = data['test_labels'][permutation[val_examples:]]\n",
    "\n",
    "    p_values = conformal_prediction.compute_p_values(\n",
    "        val_predictions, val_labels, test_predictions)\n",
    "\n",
    "    mc_p_values = monte_carlo.compute_mc_p_values(\n",
    "        mc_rng, val_predictions,\n",
    "        val_human_ground_truth, test_predictions, num_samples)\n",
    "    mc_p_values = jnp.mean(mc_p_values, axis=0)\n",
    "\n",
    "    corrected_mc_p_values = monte_carlo.compute_mc_ecdf_p_values(\n",
    "        mc_rng, val_predictions,\n",
    "        val_human_ground_truth, test_predictions, num_samples)\n",
    "\n",
    "    confidence_sets = conformal_prediction.predict_p_values(p_values, alpha)\n",
    "    mc_confidence_sets = conformal_prediction.predict_p_values(mc_p_values, alpha)\n",
    "    corrected_mc_confidence_sets = conformal_prediction.predict_p_values(\n",
    "        corrected_mc_p_values, alpha)\n",
    "\n",
    "    coverages.append(classification_metrics.aggregated_coverage(\n",
    "        confidence_sets, jax.nn.one_hot(test_labels, num_classes)))\n",
    "    mc_coverages.append(classification_metrics.aggregated_coverage(\n",
    "        mc_confidence_sets, jax.nn.one_hot(test_labels, num_classes)))\n",
    "    corrected_mc_coverages.append(classification_metrics.aggregated_coverage(\n",
    "        corrected_mc_confidence_sets,\n",
    "        jax.nn.one_hot(test_labels, num_classes)))\n",
    "\n",
    "  mc_coverages = jnp.array(mc_coverages)\n",
    "  corrected_mc_coverages = jnp.array(corrected_mc_coverages)\n",
    "\n",
    "  vmax = 0\n",
    "  hist, _ = colab_utils.plot_hist(\n",
    "      jnp.mean(mc_coverages, axis=1),\n",
    "      normalize=True, alpha=0.65, label='Monte Carlo CP')\n",
    "  vmax = max(np.max(hist), vmax)\n",
    "  hist, _ = colab_utils.plot_hist(\n",
    "      jnp.mean(corrected_mc_coverages, axis=1),\n",
    "      normalize=True, alpha=0.65, label='ECDF Monte Carlo CP')\n",
    "  vmax = max(np.max(hist), vmax)\n",
    "  plt.vlines(1 - alpha, 0, vmax, label='Target', color='black')\n",
    "  plt.title('Aggregated coverage of ECDF-based approach')\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.xticks([0.94, 0.95, 0.96])\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  if kwargs.get('name', False):\n",
    "      plt.savefig(kwargs.get('name') + '.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mS_0TTZQT_xz"
   },
   "outputs": [],
   "source": [
    "plot_coverage_ecdf_mccp(\n",
    "    predictions, data, alpha=0.05, num_samples=10, num_trials=100,\n",
    "    name='ecdf_coverage')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zbE94I6IT_xz"
   },
   "source": [
    "### Aggregated conformity scores\n",
    "\n",
    "This corresponds to an experiment calibrating with so-called aggregated conformity scores (in the first version our paper called *expected* conformity scores) which leads to plausibility regions.\n",
    "\n",
    "Please see the first version of our paper on ArXiv: [arxiv.org/abs/2307.09302v1](https://arxiv.org/abs/2307.09302v1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Wo8Q-hJ6T_xz"
   },
   "outputs": [],
   "source": [
    "def plot_conformity_scores(predictions, data, **kwargs):\n",
    "  num_samples = 10\n",
    "  labels = data['test_labels']\n",
    "  smooth_labels = data['test_smooth_labels']\n",
    "  num_examples = predictions.shape[0]\n",
    "  true_scores = predictions[jnp.arange(num_examples), labels]\n",
    "  top1_scores = predictions[\n",
    "      jnp.arange(num_examples), jnp.argmax(smooth_labels, axis=1)]\n",
    "  aggregated_scores = jnp.sum(predictions * smooth_labels, axis=1)\n",
    "\n",
    "  plot_hist(\n",
    "      true_scores, label='True scores $E(x,y)$', range=(0, 1), bins=50,\n",
    "      normalize=True, alpha=0.6)\n",
    "  plot_hist(\n",
    "      top1_scores, label='Voted scores $E(x,argmax_k\\lambda_k)$', range=(0, 1), bins=50,\n",
    "      normalize=True, alpha=0.6)\n",
    "  plot_hist(\n",
    "      aggregated_scores, label=r\"Aggregated scores $e(x,\\lambda)$\",\n",
    "      range=(0, 1), bins=50, normalize=True, alpha=0.6)\n",
    "  plt.legend()\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.xlabel('Conformity score')\n",
    "  plt.title('Conformity score histograms')\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('conformity_scores.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()\n",
    "\n",
    "  plt.plot(\n",
    "      jnp.arange(num_examples) / num_examples,\n",
    "      jnp.sort(true_scores),\n",
    "      label='True scores $E(x,y)$')\n",
    "  plt.plot(\n",
    "      jnp.arange(num_examples) / num_examples, jnp.sort(top1_scores),\n",
    "      label='Voted scores $E(x,argmax_k\\lambda_k)$')\n",
    "  plt.plot(\n",
    "      jnp.arange(num_examples) / num_examples, jnp.sort(aggregated_scores),\n",
    "      label='Aggregated scores $e(x,\\lambda)$')\n",
    "  plt.legend()\n",
    "  plt.xlabel('Frequency')\n",
    "  plt.ylabel('Conformity score')\n",
    "  plt.title('Conformity score CDFs')\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('conformity_scores_cdf.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V9A6OEakT_x0"
   },
   "outputs": [],
   "source": [
    "plot_conformity_scores(predictions, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UevthQB_T_x0"
   },
   "source": [
    "#### Reduced plausibility regions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SZaOyHJ_T_x0"
   },
   "outputs": [],
   "source": [
    "def visualize_confidence_regions(\n",
    "    rng, predictions, data, indices, alpha=0.05, split=0.5, **kwargs):\n",
    "  \"\"\"Visualize confidence regions.\"\"\"\n",
    "  permutation = jax.random.permutation(rng, predictions.shape[0])\n",
    "  val_examples = int(predictions.shape[0]*split)\n",
    "  num_classes = predictions.shape[1]\n",
    "  val_predictions = predictions[permutation[:val_examples]]\n",
    "  test_predictions = predictions[permutation[val_examples:]]\n",
    "  val_human_ground_truth = data['test_smooth_labels'][permutation[:val_examples]]\n",
    "  test_human_ground_truth = data['test_smooth_labels'][permutation[val_examples:]]\n",
    "  val_ground_truth = data['test_labels'][permutation[:val_examples]]\n",
    "  test_ground_truth = data['test_labels'][permutation[val_examples:]]\n",
    "  test_examples = data['test_examples'][permutation[val_examples:]]\n",
    "\n",
    "  baseline_threshold = conformal_prediction.calibrate_threshold(\n",
    "      val_predictions, jnp.argmax(val_human_ground_truth, axis=1), alpha)\n",
    "  baseline_confidence_sets = conformal_prediction.predict_threshold(\n",
    "      test_predictions, baseline_threshold)\n",
    "\n",
    "  threshold = plausibility_regions.calibrate_plausibility_regions(\n",
    "      val_predictions, val_human_ground_truth, alpha)\n",
    "  distributions, coverages = plausibility_regions.predict_plausibility_regions(\n",
    "      test_predictions, threshold, num_grid_points=50)\n",
    "  k = 1\n",
    "  confidence_sets = plausibility_regions.reduce_plausibilities_to_topk(\n",
    "      distributions, coverages, k=k)\n",
    "\n",
    "  colors = np.array([\n",
    "      [228,26,28],\n",
    "      [55,126,184],\n",
    "      [77,175,74],\n",
    "  ]) / 255.\n",
    "  colab_utils.plot_smooth_data(\n",
    "      data['train_examples'], data['train_smooth_labels'],\n",
    "      highlight_points=test_examples[indices], boundary=True,\n",
    "      name='data_smooth_marked', colors=colors)\n",
    "  cmap = matplotlib.cm.get_cmap('viridis')\n",
    "  plt.text(\n",
    "      0.475, 0.8, 'MLP class. boundary', color=cmap(0), fontdict={'fontsize': 14})\n",
    "  plt.show()\n",
    "\n",
    "  projected_distributions = colab_utils.project_simplex(distributions)\n",
    "  for i, n in enumerate(indices):\n",
    "    print('Case:', n)\n",
    "    print('Human ground truth:', test_human_ground_truth[n])\n",
    "    print('Conformity scores:', test_predictions[n])\n",
    "    print('Baseline confidence sets:', baseline_confidence_sets[n])\n",
    "    print(f'Top-{k} confidence set:', confidence_sets[n])\n",
    "\n",
    "    plt.bar(\n",
    "        jnp.arange(num_classes),\n",
    "        test_human_ground_truth[n],\n",
    "        label=r\"Plausibility $\\lambda_k = p(y=k|x)$\", alpha=0.65)\n",
    "    plt.bar(\n",
    "        jnp.arange(num_classes),\n",
    "        test_predictions[n],\n",
    "        label=r\"Conf. scores $E(x, k) = \\pi_k(x)$\", alpha=0.65)\n",
    "    plt.xlabel('Class')\n",
    "    plt.ylabel('Probability')\n",
    "    plt.legend(loc='lower left', bbox_to_anchor=(-0.5, 1))\n",
    "    plt.gcf().set_size_inches(2, 1)\n",
    "    plt.xticks([0, 1, 2])\n",
    "    plt.grid(True)\n",
    "    plt.savefig(f'data_smooth_{i + 1}.pdf', bbox_inches=\"tight\")\n",
    "    plt.show()\n",
    "\n",
    "    colab_utils.plot_simplex(projected_distributions, coverages[n])\n",
    "    plt.savefig(f'data_smooth_{i + 1}1.png', bbox_inches=\"tight\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nFDUvyxiT_x1"
   },
   "outputs": [],
   "source": [
    "visualize_confidence_regions(\n",
    "    jax.random.PRNGKey(0), predictions, data,\n",
    "    indices=np.array([10, 6, 12, 20]), alpha=0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cwJdaqH-T_x1"
   },
   "outputs": [],
   "source": [
    "def run_reduction_trials(\n",
    "    rng, predictions, data, alpha=0.05, trials=100, split=0.5):\n",
    "  \"\"\"Visualize confidence regions.\"\"\"\n",
    "  results = {}\n",
    "  tags = ['true_coverages', 'plausibility_coverages']\n",
    "  for tag in tags:\n",
    "    results[tag] = []\n",
    "  keys = jax.random.split(rng, trials)\n",
    "  for key in keys:\n",
    "    permutation = jax.random.permutation(key, predictions.shape[0])\n",
    "    val_examples = int(predictions.shape[0]*split)\n",
    "    val_predictions = predictions[permutation[:val_examples]]\n",
    "    test_predictions = predictions[permutation[val_examples:]]\n",
    "    val_human_ground_truth = data['test_smooth_labels'][\n",
    "        permutation[:val_examples]]\n",
    "    test_human_ground_truth = data['test_smooth_labels'][\n",
    "        permutation[val_examples:]]\n",
    "    test_ground_truth = data['test_labels'][permutation[val_examples:]]\n",
    "\n",
    "    threshold = plausibility_regions.calibrate_plausibility_regions(\n",
    "        val_predictions, val_human_ground_truth, alpha)\n",
    "    distributions, coverages = plausibility_regions.predict_plausibility_regions(\n",
    "        test_predictions, threshold)\n",
    "    confidence_sets = plausibility_regions.reduce_plausibilities_to_topk(\n",
    "        distributions, coverages, k=1)\n",
    "    num_classes = test_predictions.shape[1]\n",
    "    true_coverages = classification_metrics.aggregated_coverage(\n",
    "        confidence_sets, jax.nn.one_hot(test_ground_truth, num_classes))\n",
    "    plausibility_coverages = plausibility_regions.check_plausibility_regions(\n",
    "        test_predictions, test_human_ground_truth, threshold)\n",
    "    results['true_coverages'].append(true_coverages)\n",
    "    results['plausibility_coverages'].append(plausibility_coverages)\n",
    "  for tag in tags:\n",
    "    results[tag] = jnp.array(results[tag])\n",
    "  return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xOeWrUNZT_x2"
   },
   "outputs": [],
   "source": [
    "def plot_reduced_plausibility_regions(predictions, data, alpha=0.05, **kwargs):\n",
    "  plausibility_results = run_reduction_trials(\n",
    "      jax.random.PRNGKey(0), predictions, data, alpha=alpha)\n",
    "\n",
    "  hist, _ = plot_hist(\n",
    "      jnp.mean(plausibility_results['plausibility_coverages'], axis=-1),\n",
    "      normalize=True, label='Plausibility coverage')\n",
    "  vmax = np.max(hist)\n",
    "  hist, _ = plot_hist(\n",
    "      jnp.mean(plausibility_results['true_coverages'], axis=-1),\n",
    "      normalize=True, label='True label coverage')\n",
    "  vmax = max(vmax, np.max(hist))\n",
    "  plt.vlines(0.95, 0, vmax, color='red', label='Target plausibility coverage')\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.legend()\n",
    "  plt.title('Coverage of reduced plausibility regions')\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 5), kwargs.get('height', 3))\n",
    "  plt.savefig('reduced_plausibility_regions.pdf', bbox_inches=\"tight\")\n",
    "  plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QmHUqH1rT_x2"
   },
   "outputs": [],
   "source": [
    "plot_reduced_plausibility_regions(predictions, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LB0m_533508l"
   },
   "source": [
    "## Dermatology dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZecM988W508o"
   },
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TwZmtLxw508o"
   },
   "outputs": [],
   "source": [
    "with open('data/dermatology_data.pkl', 'rb') as f:\n",
    "  data = pickle.load(f)\n",
    "with open('data/dermatology_predictions0.txt', 'r') as f:\n",
    "  predictions =  np.loadtxt(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sc_u-wxy508p"
   },
   "source": [
    "### Main results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qHlPeeMA508p"
   },
   "outputs": [],
   "source": [
    "def plot_conformal_prediction(\n",
    "    predictions, plausibilities, num_trials=100,\n",
    "    num_samples=10, alpha=0.27, method=False,**kwargs):\n",
    "  \"\"\"Plot expected/standard accuracy for CP against IRN top-1 labels.\"\"\"\n",
    "  _, num_classes = plausibilities.shape\n",
    "  top1_coverages = []\n",
    "  aggregated_coverages = []\n",
    "  sizes = []\n",
    "  keys = jax.random.split(jax.random.PRNGKey(0), 2 * num_trials)\n",
    "  for t in range(num_trials):\n",
    "    permutation = jax.random.permutation(keys[2 * t], predictions.shape[0])\n",
    "    val_examples = int(predictions.shape[0] * 0.5)\n",
    "    val_predicitions = predictions[permutation[:val_examples]]\n",
    "    test_predicitions = predictions[permutation[val_examples:]]\n",
    "    val_plausibilities = plausibilities[permutation[:val_examples]]\n",
    "    test_plausibilities = plausibilities[permutation[val_examples:]]\n",
    "\n",
    "    if method == 'mccp':\n",
    "      threshold = monte_carlo.calibrate_mc_threshold(\n",
    "          keys[2 * t + 1], val_predicitions, val_plausibilities,\n",
    "          alpha, num_samples=num_samples)\n",
    "      confidence_sets = conformal_prediction.predict_threshold(\n",
    "          test_predicitions, threshold)\n",
    "    elif method == 'ecdf':\n",
    "      p_values = monte_carlo.compute_mc_ecdf_p_values(\n",
    "          keys[2 * t + 1], val_predicitions, val_plausibilities,\n",
    "          test_predicitions, num_samples=num_samples, split=0.5)\n",
    "      confidence_sets = conformal_prediction.predict_p_values(\n",
    "          p_values, alpha)\n",
    "    else:\n",
    "      val_labels = jnp.argmax(val_plausibilities, axis=1)\n",
    "      threshold = conformal_prediction.calibrate_threshold(\n",
    "          val_predicitions, val_labels, alpha)\n",
    "      confidence_sets = conformal_prediction.predict_threshold(\n",
    "          test_predicitions, threshold)\n",
    "\n",
    "    test_labels = jnp.argmax(test_plausibilities, axis=1)\n",
    "    top1_coverages.append(classification_metrics.aggregated_coverage(\n",
    "        confidence_sets, jax.nn.one_hot(test_labels, num_classes)))\n",
    "    aggregated_coverages.append(classification_metrics.aggregated_coverage(\n",
    "        confidence_sets, test_plausibilities))\n",
    "    sizes.append(classification_metrics.size(confidence_sets))\n",
    "  top1_coverages = jnp.array(top1_coverages)\n",
    "  aggregated_coverages = jnp.array(aggregated_coverages)\n",
    "  sizes = jnp.array(sizes)\n",
    "\n",
    "  hist1, _ = colab_utils.plot_hist(\n",
    "      jnp.mean(top1_coverages, axis=1),\n",
    "      label='Voted coverage', normalize=True)\n",
    "  hist2, _ = colab_utils.plot_hist(\n",
    "      jnp.mean(aggregated_coverages, axis=1),\n",
    "      label='Aggregated coverage', normalize=True)\n",
    "  vmax = max(jnp.max(hist1), jnp.max(hist2))\n",
    "  plt.vlines(\n",
    "      1 - alpha, 0, vmax,\n",
    "      color='black', label=r'Target')\n",
    "  plt.title(kwargs.get(\n",
    "      'title',\n",
    "      'Empirical coverage for standard CP with voted IRN labels'\n",
    "      f'\\n({num_trials} random calibration/test splits)'))\n",
    "  plt.xlabel('Empirical coverage')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 4.5), kwargs.get('height', 2.5))\n",
    "  plt.show()\n",
    "\n",
    "  _, _ = colab_utils.plot_hist(\n",
    "      jnp.mean(sizes, axis=1), normalize=True)\n",
    "  vmax = max(jnp.max(hist1), jnp.max(hist2))\n",
    "  average = jnp.mean(sizes)\n",
    "  plt.vlines(\n",
    "      average, 0, vmax, label=f'Average: {average:.2f}',\n",
    "      color='black', linewidth=2)\n",
    "  plt.xlabel('Inefficiency')\n",
    "  plt.ylabel('Frequency')\n",
    "  plt.legend()\n",
    "  plt.gcf().set_size_inches(kwargs.get('width', 4.5), kwargs.get('height', 1.5))\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pKZJ9Cj_508p"
   },
   "outputs": [],
   "source": [
    "for alpha in [0.27, 0.1]:\n",
    "  plot_conformal_prediction(\n",
    "      predictions, data['test_irn'], method=False,\n",
    "      title='CP with voted labels', alpha=alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fYZ4D8tj508p"
   },
   "outputs": [],
   "source": [
    "for alpha in [0.27, 0.1]:\n",
    "  plot_conformal_prediction(\n",
    "      predictions, data['test_irn'], method='mccp',\n",
    "      title='Monte Carlo CP', alpha=alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FmBYAYJ2508p"
   },
   "outputs": [],
   "source": [
    "# This runs much longer than the above.\n",
    "plot_conformal_prediction(\n",
    "    predictions, data['test_irn'], method='ecdf',\n",
    "    title='ECDF Monte Carlo CP', name='ecdf', alpha=0.27)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G2tc_Oxf508p"
   },
   "outputs": [],
   "source": [
    "plot_conformal_prediction(\n",
    "    predictions, data['test_irn'], method='ecdf',\n",
    "    title='ECDF Monte Carlo CP', name='ecdf', alpha=0.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
}