examples/plankton.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f3c181f-6dd5-446c-ac06-382b498f7600",
   "metadata": {},
   "source": [
    "# Counting plankton\n",
    "\n",
    "Assessment of the increases in phytoplankton growth during springtime warming is important for the study of global biogeochemical cycling in response to climate change. The goal is to count the number of plankton observed by the Imaging FlowCytobot (1, 2), an automated, submersible flow cytometry system, at Woods Hole Oceanographic Institution in the year 2014. Data from 2013 is labeled and data from 2014 is imputed with machine-learning predictions from a state-of-the-art ResNet fine-tuned on all data up to and including 2012. The years 2013 and 2014 have a distribution shift, primarily caused by the change in the base frequency of plankton observations. This notebook applies prediction-powered inference for label-shifted problems.\n",
    "\n",
    "1. R. J. Olson, A. Shalapyonok, H. M. Sosik, An automated submersible flow cytometer for analyzing pico-and nanophytoplankton: FlowCytobot. Deep Sea Research Part I: Oceanographic Research Papers 50(2), 301–315 (2003).\n",
    "2. E. C. Orenstein, O. Beijbom, E. E. Peacock, H. M. Sosik, WHOI-Plankton- A large scale fine grained visual recognition benchmark dataset for plankton classification. arXiv preprint arXiv:1510.00745 (2015). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd39cbfb-8736-49f9-adf6-3075f9cfc3cd",
   "metadata": {},
   "source": [
    "### Import necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "644d0704-3c1b-4d6a-97ee-5a16ff07c04e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os, sys\n",
    "\n",
    "sys.path.append(os.path.abspath(os.path.join(os.getcwd(), os.pardir)))\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.stats import binom, norm\n",
    "from ppi_py.datasets import load_dataset\n",
    "from ppi_py import ppi_distribution_label_shift_ci, classical_mean_ci\n",
    "from tqdm import tqdm\n",
    "from scipy.optimize import brentq\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cf90ae6",
   "metadata": {},
   "source": [
    "### Import the plankton data set\n",
    "\n",
    "Load the data. The data set contains hand-annotated labels indicating whether the image contains plankton (```Y```), corresponding labels predicted by a ResNet (```Yhat```), both from year 2013, and the same data from year 2014 (```Y_unlabeled``` and ```Yhat_unlabeled```, respectively). The data from 2014 is used as the unlabeled data set in the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a6da3138",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_folder = \"./data/\"\n",
    "data = load_dataset(dataset_folder, \"plankton\")\n",
    "Y = data[\"Y\"].astype(bool)\n",
    "Yhat = data[\"Yhat\"].astype(bool)\n",
    "Y_unlabeled = data[\"Y_unlabeled\"].astype(bool)\n",
    "Yhat_unlabeled = data[\"Yhat_unlabeled\"].astype(bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8969f9db",
   "metadata": {},
   "source": [
    "### Problem setup\n",
    "\n",
    "Specify the error level (```alpha```), range of values for the labeled data set size (```ns```), number of trials (```num_trials```), and other problem parameters.\n",
    "\n",
    "Compute the ground-truth value of the estimand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5b3c8f29",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 0.05  # Total error budget\n",
    "delta = 0.048  # Rectifier budget\n",
    "K = 2  # Total number of classes\n",
    "nu = np.array([0, 1])  # Linear contrast to estimate\n",
    "n_total = Y.shape[0]  # Total number of labeled examples\n",
    "N = Y_unlabeled.shape[0]  # Total number of unlabeled examples\n",
    "ns = np.linspace(1000, n_total, 20).astype(\n",
    "    int\n",
    ")  # Test for different numbers of labeled examples\n",
    "num_trials = 50\n",
    "\n",
    "# True mean (computed on all labels)\n",
    "true_theta = Y_unlabeled.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83ce18be",
   "metadata": {},
   "source": [
    "### Construct intervals\n",
    "\n",
    "Form confidence intervals for all methods and problem parameters. A dataframe with the following columns is formed:\n",
    "1. ```method``` (one of ```PPI``` and ```Imputation```; due to label shift there is no classical baseline)\n",
    "2. ```n``` (labeled data set size, takes values in ```ns```)\n",
    "3. ```lower``` (lower endpoint of the confidence interval)\n",
    "4. ```upper``` (upper endpoint of the confidence interval)\n",
    "5. ```trial``` (index of trial, goes from ```0``` to ```num_trials-1```)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "812f8fd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:19<00:00,  1.01it/s]\n"
     ]
    }
   ],
   "source": [
    "# Run prediction-powered inference and classical inference for many values of n\n",
    "results = []\n",
    "for i in tqdm(range(ns.shape[0])):\n",
    "    for j in range(num_trials):\n",
    "        # Prediction-Powered Inference\n",
    "        n = ns[i]\n",
    "        rand_idx = np.random.permutation(n_total)\n",
    "        _Yhat = Yhat[rand_idx[:n]]\n",
    "        _Y = Y[rand_idx[:n]]\n",
    "\n",
    "        ppi_count_ci = ppi_distribution_label_shift_ci(\n",
    "            _Y, _Yhat, Yhat_unlabeled, K, nu, alpha, delta, return_counts=True\n",
    "        )\n",
    "        # Append results\n",
    "        results += [\n",
    "            pd.DataFrame(\n",
    "                [\n",
    "                    {\n",
    "                        \"method\": \"PPI\",\n",
    "                        \"n\": n,\n",
    "                        \"lower\": ppi_count_ci[0],\n",
    "                        \"upper\": ppi_count_ci[1],\n",
    "                        \"trial\": j,\n",
    "                    }\n",
    "                ]\n",
    "            )\n",
    "        ]\n",
    "\n",
    "# Imputed CI\n",
    "imputed_ci = classical_mean_ci(Yhat_unlabeled > 0.5, alpha=alpha)\n",
    "\n",
    "results += [\n",
    "    pd.DataFrame(\n",
    "        [\n",
    "            {\n",
    "                \"method\": \"Imputation\",\n",
    "                \"n\": np.nan,\n",
    "                \"lower\": imputed_ci[0] * N,\n",
    "                \"upper\": imputed_ci[1] * N,\n",
    "                \"trial\": 0,\n",
    "            }\n",
    "        ]\n",
    "    )\n",
    "]\n",
    "\n",
    "df = pd.concat(results, axis=0, ignore_index=True)\n",
    "df[\"width\"] = df[\"upper\"] - df[\"lower\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d15ba288",
   "metadata": {},
   "source": [
    "### Plot results\n",
    "\n",
    "Plot:\n",
    "1. Five randomly chosen intervals from the dataframe for PPI, and the imputed interval;\n",
    "2. The average interval width for PPI, together with a scatterplot of the widths from the five random draws.\n",
    "\n",
    "Note that due to label shift there is no classical baseline to plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6077b2c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAnsAAADMCAYAAAALZGVzAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAAkA0lEQVR4nO3de5zcdX3v8dd7Zu8h92wISYncYgsKCgSIJYFIICEIWvRofdiL1nNAOFixttVWQhUvvVjrqdQeKJyqeOkV5FSBkFhIMAkkISACgthwVSCQQG5krzPz6R/z28nsZje7Ibvzm5m8nzzmsb/5zndmPr/ZTPLm+/19fz9FBGZmZmZWnzJpF2BmZmZmY8dhz8zMzKyOOeyZmZmZ1TGHPTMzM7M65rBnZmZmVscc9szMzMzqWMMwj/u8LFYtlHYBVcDfR6sW/j4C559/ftx5551pl2E27PdxuLBnZmZVrveR5aXtxhOXpljJoWXbtm1pl2A2Ig57ZmY1Lr/53tK2w56ZDeRj9szMzMzqmEf2zMxqXMObz0+7BDOrYg57ZmY1rmHOmWmXUFGSfgzsTO4+DXwR+CbFRUyPAldEREHSJcBHgBzwhYi4TVIr8B1gOrAb+GBEbJU0D/hq0ndlRFzzemrrzHey8uUVNGeaWTx9CRl5As3S57BnZmY1Q1ILQEQsLGv7PrAsIlZLuh54l6T7gI8Bc4EWYK2kHwKXA49ExGclvR9YBlwJXA+8B3gKuF3SKRHx4IHWd9fL/8k921YDsKVrC4unL2Zm66zXvb9mo8H/y2FmZrXkLUCbpJWS7k5G5E4F7kkeXw6cC5wOrIuI7ojYCWwGTgLmA3eW95U0AWiOiCcjIoAVwKIDLeyFzud5eOdPSvef6XiGO7bczgudz7+uHTUbLQ57ZmZWSzqALwNLgMuA7wJKQhoUp2YnAhPYO9U7VHt5265B+u5D0qWSNknatHXr1n6PrX1lLV2FrtL9QuTZ1rONta+sfR27aTZ6PI1rZlbjen/8H6XtxpPflWIlFfFzYHMS7n4u6RWKI3t9xgM7KIa38cO0D9d3HxFxA3ADwNy5c/ud6Py5zmdpUlO//s2ZZp7rfHaEu2Y2NjyyZ2ZW4/LPbCrdDgEfBv4GQNJMiqNyKyUtTB5fCqwBNgILJLVImggcT3HxxjrggvK+EbEL6JF0rCRRHDVc83qKU9mCjEIUfNkbqwoe2TMzs1ryj8A3Ja2luPr2w8A24EZJTcDjwM0RkZd0LcXQlgGuioguSdcBNyXP7wE+kLxu35RwluJq3A0HWtgbWo/isd0/Ld3PkaO70MVx4+a83n01GxUOe2ZmNa7hre9Mu4SKiYjygFbu7EH63gjcOKCtA3jvIH3XA/MOprYzp57JS91b2JkrHhJYiALtze2cOfXQOjWOVR9P45qZ1biGo08r3Sw9M1tncd70xaX7GWW44PB3+NQrljqHPTMzs1Fy9LhjStv5yHNEy8wUqzErctgzMzMbJY1qpEHFI6Tykac3elOuyMxhz8zMbNRIojXbVrrfke9IsRqzIi/QMDOrcT33/1tpu+m096VYiQG0ZVvZnSueo7kj38GkxknpFmSHPIc9M7MaV/jlI3vvOOylrq1sZK8z35liJWZFnsY1MzMbRa39wp6ncS19HtkzqxMbXz3gc8BavThr6d7tA/hzcPqUM8agGGvLtpa2O3IOe5Y+hz2zOjGlaQrtzdPTLsNqSFe+i5ZsS9pl1B0v0LBq42lcMzOzUeRj9qzaOOyZmZmNkhc6n+eZjmdK97f2vJxeMWYJhz0zM7NR8ELn86zetoogSm3PdTzHC53Pp1iVmcOeWd2Qv85mqXp418O0ZFsZlx1XaguCh3c9nGJVZl6gYVY3jj3s2LRLsJR03Xp1abvl4s+nWMmhbXvvdiY1TKIn011qy0ee7b3bU6zKzCN7ZmZmo2Jy42Q6C500qrHU1lPoZXLj5BSrMvPInplZzWuc91tpl2DASRNOYvW2VWSULbXlopeTJpyUYlVmDntmdenBHQ+wO7c77TKsUhrLtrdted0vc/a0hQddyqFsZussFk57Ow/tfKjUlo88R7TMTK8oMxz2zOpSS6aVoycck3YZVmM6c520NrQO39GGNLN1FjNbZ/GjbffQEz0EQXeh2yevtlT5mD0zM7NR1tbgq2hY9XDYMzMzG2Wt5dfHddizlDnsmZmZjbLWfpdMc9izdDnsmdUhHx9klq7y6+N2+Pq4ljIv0DCrQ8eM8+KMQ4lPqlx92jyNa1XEYc/MrMY1zv9w2iXYAJ7GtWrisGdmVuOy7UenXYIN0H8a12HP0uVj9szMzEZZW7+RPR+zZ+nyyJ5ZHbltyw/SLsFq0IUzLkq7hLrjU69YNXHYM6sjR7TMZHbr7LTLsBqzJ7eHcQ3j0i7jgEiaDjwAnAfkgG8CATwKXBERBUmXAB9JHv9CRNwmqRX4DjAd2A18MCK2SpoHfDXpuzIirjmY+tp8zJ5VEU/jmplZTZHUCPwD0Dc/+hVgWUQsAAS8S9IM4GPAmcAS4C8kNQOXA48kfb8FLEte43rgA8B84AxJpxxMja0+9YpVEYc9MzOrNV+mGM5eSO6fCtyTbC8HzgVOB9ZFRHdE7AQ2AydRDHN3lveVNAFojognIyKAFcCigymw3wKNnEf2LF0Oe2ZmVjMkfQjYGhErypuTkAbFqdmJwARgZ1mfwdrL23YN0new979U0iZJm7Zu3TpknZ7GtWrisGdWRzL+Slv9+zBwnqTVwFspTsVOL3t8PLCDYngbP0z7cH33ERE3RMTciJjb3t4+ZJHlCzS6Cl0UorD/vTIbQ16gYVZHTp50ctolmI2piDirbzsJfJcBfy1pYUSsBpYCq4CNwBcltQDNwPEUF2+sAy5IHl8KrImIXZJ6JB0LPEXxGL+DWqCRUYaWTAtdhS6CoDPfWXOLYKx+OOyZmVmt+0PgRklNwOPAzRGRl3QtsIbiLNZVEdEl6TrgJklrgR6KizKgGBq/C2QprsbdcLBFtWXb6Cp0ATjsWaoc9szMrCZFxMKyu2cP8viNwI0D2jqA9w7Sdz0wbzTra822Qe+rgI/bs3Q57JnVsa89+Xdpl2A16KPH/n7aJdSFNp9Y2aqEw55ZHTvusOM4Yfyb0i7DasyOnh1MapqUdhk1z9fHtWrhpXtmZmZjoNXXx7Uq4bBnZmY2Bnx9XKsWDntmZmZjoK3B07hWHRz2zOpYlmzaJZgdsto8jWtVwgs0zOrYeYcvTrsEs0NWqxdoWJXwyJ6ZmdkY8GpcqxYOe2ZmZmOg/Dx7PqmypcnTuGZ17NM//ZO0S7BKyHXv3W5oPuiXm9U6iyuO8YmVD5anca1aOOyZ1bE3TziR0yafnnYZVmPu374x7RLqghdoWLXwNK6ZmdkYaM40IwRAd6GbfORTrsgOVQ57ZmZmYyCjjE+sbFXBYc/MzGyM9J/KddizdDjsmdWxBvmwXDtwU5qmpF1C3ei/SMPH7Vk6/C+BWR1736/8ZtolWA2ac9ictEuoG+WnX+nIeWTP0uGRPTMzszHS6mlcqwIe2TMzq3FdP/h8abvloqtTrMQGavM0rlUBhz0zs1qX60m7AhtC/7C3J8VK7FDmaVwzM7Mx0trvkmke2bN0eGTPzKzGNV+4LO0SXhdJlwAfB1oBARERx6Ra1CjzqVesGjjsmZnVODUe/PVwU3IZcAGwJe1Cxoqvj2vVwGHPzMzSsi0ink27iLHU79Qrnsa1lDjsmZlZRUn682SzSdIK4EEgACLi06kVNgbasuNK257GtbQ47JmZWaU9MeBnn6h0IWPN18a1auCwZ2ZW47pu3XtuvZaLP7+fntUhIm4CkPS1iPhoX7ukbwHfSq2wMeDz7Fk18KlXzMysoiRdIelF4BJJLyS3F4FZadc22poyTWSVBSAXvfQUfE5EqzyHPTMzq6iI+PuIOAK4JiJmJrcjImLRcM+VlJX0dUnrJP1I0rGSjpO0VtIaSddJyiR9L5G0SdJ6SRcmba2Sbkn63iGpPWmfJ2lD8rqfGa19leRz7VnqPI1rZlbjamHqtpyk3002f1m2DUBEDDeNe1HS70xJC4GvUDxH37KIWC3peuBdku4DPgbMBVqAtZJ+CFwOPBIRn5X0fmAZcCVwPfAe4CngdkmnRMSDB7+3xanc13KvAcVFGhMbJ47Gy5qNmMOemZlV2vHJzzOATuBe4DSgkWGO2YuI/y/ptuTuG4CXgHcA9yRty4HFQB5YFxHdQLekzcBJwHzgS2V9r5Y0AWiOiCcBkhXCiyiuEj5oPm7P0uawZ2ZmFRURfwog6c6IeEdfu6SVI3x+TtJNwMXA/wAujIi+lby7gYnABGBn2dMGay9v2zWg76BX8pB0KXApwOzZs0dSrk+sbKnzMXtmZpaW6ZImAUiaCkwd6RMj4oPAG4EbKV5urc94YAfF8DZ+mPbh+g72vjdExNyImNve3j6iWvsfs+ewZ5XnsGdmVuOit7t0qzFfBDZJehDYAPzBcE+Q9DuS/jS52wEUktdYmLQtBdYAG4EFklokTaQ4dfwosI7iJdpKfSNiF9CTLPYQsCR5jVHR5pE9S5mncc3Malz3bV8obdfSYo2IuAW4RdJ04JWIyI/gad8DviHpRxSP8fs48Dhwo6SmZPvmiMhLupZiaMsAV0VEl6TrgJskrQV6gA8kr3sZ8F0gC6yMiA2jtZ8Oe5Y2hz0zM6uovpMpS3oA6C5rJyJ+fX/PjYg9wPsGeejsQfreSHGat7ytA3jvIH3XA/NGtgcHxqdesbQ57JmZ1bqGprQrOFB9w49HAyuAByiujH0ttYrGkEf2LG0Oe2ZmNa7loquH71RFIuKl5OcUSccD7wRuAF6muMK2rvQLezmHPas8hz0zM0uFpLcA5wHnJE2Pp1jOmCk/9YpX41oaHPbMzCwtP6J4xYqrIuKOtIsZK+Uje50FH7NnledTr5iZWVqmUjzdygJJd0n657QLGgttZQs0PI1rafDInplZjYvOvRd/UOuEFCs5YJOAWRQve9YGPJtqNWNk4BU0IoLi6fzMKsNhz8ysxnXf+del7Vo6zx7Flbi3Al+MiJ+mXcxYacw00qhGeqOXAgV6Cj00Z5vTLssOIQ57ZmaWiog4Ne0aKqU120ZvrnhJ3o58h8OeVZSP2TMzq3Ut4/ferCr1O27PK3KtwjyyZ2ZW41qWfjLtEmwYbT79iqXII3tmZmZjbOAiDbNKctgzMzMbY/2ncX2uPasshz0zM7Mx5mlcS5OP2TMzq3GF3dtK25nx01KsxIbiaVxLk8OemVmN6/nPr5a2a+w8e4cMXx/X0uRpXDMzszHmY/YsTR7ZMzOrcRo3Ne0SbBhtDZ7GtfQ47JmZ1bjmxR9PuwQbhqdxLU2exjUzMxtjnsa1NDnsmZmZjbG27LjStqdxrdIc9szMzMZYa9nIXle+k0IUUqzGDjU+Zs/MrMYVtj9f2s5MnpViJTaUrLI0Z5rpLnQTBN2Frn7H8ZmNJYc9M7Ma17P6+tK2z7NXvdqybXQXuoHiVK7DnlWKp3HNzMwqoNWLNCwlHtkzM6txmjQz7RJsBHz6FUuLw56ZWY1rfvvlaZdgI9BWfn3cnMOeVY6ncc3MzCqgX9jzyJ5VkMOemZnVDEmNkr4taY2kjZLeKek4SWuTtuskZZK+l0jaJGm9pAuTtlZJtyR975DUnrTPk7RB0jpJnxmL2suP2ev0MXtWQQ57ZmZWS34beCUiFgBLga8BXwGWJW0C3iVpBvAx4ExgCfAXkpqBy4FHkr7fApYlr3s98AFgPnCGpFNGu3CP7FlafMyemVmNy299urSdbT86xUoq4t+Bm8vu54BTgXuS+8uBxUAeWBcR3UC3pM3ASRTD3JfK+l4taQLQHBFPAkhaASwCHhzNwh32LC0e2TMzq3G9a79eutW7iHgtInZLGk8x9C0DFBGRdNkNTAQmADvLnjpYe3nbrkH67kPSpcnU8KatW7ceUO2exrW0OOyZmVlNkXQksAr4dkT8E1B+7bHxwA6K4W38MO3D9d1HRNwQEXMjYm57e/sB1e3r41paHPbMzGqcph1VutU7SYcDK4FPRUTfUOaPJS1MtpcCa4CNwAJJLZImAscDjwLrgAvK+0bELqBH0rGSRPEYvzWjXXtbv5E9hz2rHB+zZ2ZW45oX/M+0S6ikTwOTKR5rd3XSdiVwraQm4HHg5ojIS7qWYmjLAFdFRJek64CbJK0FeiguygC4DPgukAVWRsSG0S681cfsWUoc9szMrGZExJUUw91AZw/S90bgxgFtHcB7B+m7Hpg3SmUOqq3BYc/S4WlcMzOzCmjJtJS2uwvd3LHldl7ofD7FiuxQ4bBnZmZWAVu6XiRT9s/uY7sfc+CzinDYMzOrcfkXf1a6WfVa+8rafvd39uzgxe4X92k3G20+Zs/MrMb1rv9uaTt78edTrMT2Z/Oezf3u90QPhXzs02422jyyZ2ZmVgFd+S4atHeMpTd6KUSernxXilXZocAje2ZmNS4z41fTLsFGoDXbQi7fS4MayEUOgO7opi3TNswzzQ6Ow56ZWY1rettvp12CjcCx447jiXiC3nwve+I1guIV3rryXXTnu2nONqdcodUrT+OamZlVwPyp8zmieQaTmyczrXEaQgB0Fjr59xf+jb2X9zUbXQ57ZmZmFTCzdRYXzHgHx48/nqMOO5oTxr+p9NiDOx7gvlfvTbE6q2eexjUzM6uQma2zmNk6q3T/X3/5L2zYvh6AW1/8Hke2zebI1iPTKs/qlEf2zMxqXP4XPyndrLZcPPPdzGyZCUA+8tz07Dd8KTUbdR7ZMzOrcb2bbi5tZ498S4qV2IFqyjTxwdm/x//Z/Dd0Fbp4tfdV/uyxZYzPjmfB1LM45/BFaZdodcAje2ZmZilqb27nTePfXLpfoMCe/B5WvHwnd790V4qVWb1w2DMzq3GZXzmxdLPa9OSezTSUTbblkv/u3naXV+naQfM0rplZjWs67X1pl2AHqbPQSVumjY7oKJ1wOQg6Ch1c//R1/MbMizmi5YiUq7Ra5bBnZmaWstZMKz2FHsZlDqMnuukqdJVOuvxfe37Ol//rS8xufQPbe16lO7ppzbQyf8oCH9NnI+JpXDMzs5TNn7KAXOTojR6a1ESLWsiQKZ14OQie7XyGXfldKER3vpuVW1f4mD4bEYc9MzOzlJ1z+CIWty+hKdNEV3TRnG3mgunv4I/mfJI5497Yr29XdNEZneQjz6ptd1OIQkpVW63QMAd++qhQqxZKu4Aq4O+jDSr39P2l7YajT6vEW/r7CMydOzc2bdo05u8TEXzqp39MIQoU2DfYTWiYyCmTTqGQz/Pw7ofpLHR6mvfQMuz30cfsmZnVuNxD3y9tVyjsWQVJ4rDsYXTnuwlBT6G7X+jbldvJ6m2rin0RTWqiM9/JipfvBHDgM4c9M7NU/Pxb0LFl3/a2GfDG3618PVbV5k9ZwMqtK2iggcMyh9Ed3fRGLw1qoCd6Sv2CoDu6S/eXv3wHO/I7eEPbUbzY8QIP7NhEV3R55O8Q47BnZpaGTBPMGuIf2j0vwLiZI36p7FFzR6koq1Z9oWztq2uK07TZVs6bspizpy/kid1P8PVn/9+gU7x58qx9ZQ1rX1lTamuggdcKr3Hny8t5LfcaF868iIwy3P3SXXtf32GwrjjsmZnVuMaT35V2CVYB5xy+aNDwdcKEE5jQMIHufDeSyEWOXOQGDX9QPGFzn9WvrmLd9rW0ZcexO7eLrLI0qpGufFe/aWAHwdrmsGdmZlbjyqd5WzOt9EYvvYVeTp54CpOaJ/HDl1eWzts3UG/0sjO3A4BCFOiN3tJjd7x8Oxt2rOeVnlf2CYIRwaIZ5wLUbRisl/1y2DMzM6tx+0zzZlo5Z9rekcCNr26gO99NQ6aRfOTJR47e6CWS/4ZSoMDWnq3F7QFB8Patt7F+53qIYHvvdjJkaFADHfmO0hTxeTPOoyXTyqqX795vaKrGUHX3S3cVQ23y3+7Cbu54+XZ+2fkLzp6+kOnN07lv231VV/dgfOoVqxU+1YO/j/XlueXQOn3wx9qOOKBj9lLg7yOVO/XKaLj7pbuKI39qoFGN9EYvucixuH0Jp089g7964i/oiR6EyEeBAvn9hsADkSFDgQKZ5L++8DRn3BxOnHgST732FA/v+kmptr5p6CXTzx/RFPJoB8WI4JmOZ7j+6f/bL9zub/+yygLFQDxv8ttYMuN8xmXHDRtyR8mw30eHPasV/sfF30cbQu6/1pW2G+acWYm39PeR2gp7sP9QNFgY7C308utTzmTTjvuTIJihQH7I8/2NNiHasm105DsQIqssEUGBAnPGzeGECW/i6dee4tHdj5ammHORIx95zm0/l3MPX4ykEYfFjnxHcf8zjezM7Tzo+hvUQD7ypTAYBBHBmVPmc/6MpTRnm0crqDrsWd3wPy7+PtoQum69urTdcvHnK/GWqX8fJZ0B/FVELJR0HPBNit+RR4ErIqIg6RLgI0AO+EJE3CapFfgOMB3YDXwwIrZKmgd8Nem7MiKuGa6GWgt7wxkqeAw1KnjO1EWcNPkt/P2Tf0dv9CIyFKNYUIg8wD6nhqmkvpBVrE39RhVnNs9k9rg38FLnFp7pfIYMGfJDjGY2qIEmNRMEvYUeJDGteRovdr14UPU1qYne6N0nDJ4xeR7nHn4e4xvGs/rlVSMJgw57VjdS/8elCvj7aIM61MKepE8CvwPsiYh5kr4PfCUiVku6HlgB3Af8EJgLtABrk+0rgAkR8VlJ7wfeFhFXSnoIeA/wFHA7sCwiHtxfHfUW9vbnQEcF+6aIzzl8Edc89hm6C8XjBSOKYTCXhJy3TjqZB3c8QD4Jh8MdQ1hJR7bM5sWuF2jKNO13v7JqoECBQuTJRY4gaMg00FM4uJArRBBkyNCSaQHo9/79ug7DCzTMzGpc9rhfT7uESnsSeDfw7eT+qcA9yfZyYDGQB9ZFRDfQLWkzcBIwH/hSWd+rJU0AmiPiSQBJK4BFwH7D3qFkqNO+9D0GAxaHTNnbf8HUs1i5dQWFyJdCkxDntS/mnMMX0d7YPugU8tlTF7Jhx3p6Cj00qJGgQEApKM6dfBobt29Irg0s+qLiwUwx912BJAj+YM4n9gm5g+1XhgJNaqQ3qWBx+xLePv0crnm8LwxmkzBYIBfF095IKgXcofSF3gIFOgudTMhOLH3OBzrV67BnZlbjGk9cmnYJFRURt0g6qqxJsXeaajcwEZgAlB94NVh7eduuAX2PGey9JV0KXAowe/bsg9qPenIwYXB/K4lbsi2s3LqCoDBoUJzcMHnIUcWz2s/mCz/7HD2FHhozjQTFxRe90UODGlg64wJu33IbuUKOrLJkS6/RQ1Om6aD366ypZ7Ny6woAmtVML8XFHovbl7Bw+tv53OOfLQuDQSEK5JMw2JJtYU9+T+m9smpAEo000lnoPODfj8OemZnVuvKhnPHADorhbfww7cP13UdE3ADcAMVp3IOs+5Cxv9C0v8dfV1AcJHDl+0YVk8B1zrRFnDl1Pt25blZuXZEs8GigN3qKxyNOGdnI2eiFwaaktthnmrgx00SGDFA8J2JrpnVEtZVz2DMzs1r3Y0kLI2I1sBRYBWwEviipBWgGjqe4eGMdcEHy+FJgTUTsktQj6ViKx+wtAYZdoGGV8XqDYt9j8PrD4ljWPtLp73zkyKiRnsKBBdFyXqBhtcILNPx9tOqR+vcxmcb9l2SBxhuBG4Em4HHgkojIJ6txLwUywJ8n079twE3AEUAP8IGI2JKsxv1bIEtxNe5Vw9VwKC3QsHSM8NQsXo1rdSP1f1yqgL+PNqjex+8ubTcef04l3tLfRxz2rGp4Na6ZWb3L/2xVabtCYc/Makgm7QLMzMzMbOwMN41rZlVC0p3AtLTr2I9pwLa0ixhD3r+9tkXE+WNZTC2QtBV4doiH6/3Py+vlz2VwB/O5DPt9dNgzs1EhaVNEzE27jrHi/bMD4c9zcP5cBjfWn4uncc3MzMzqmMOemZmZWR1z2DOz0XJD2gWMMe+fHQh/noPz5zK4Mf1cfMyemZmZWR3zyJ6ZmZlZHXPYM7NBSTpD0upk+xRJz0tandx+M2m/RNImSeslXZi0tUq6RdIaSXdIak/a50naIGmdpM+kuF+Nkr6d1LdR0jslHSdpbdJ2naRMre7ffvaxbn6H1UpSRtL1ku5LPuPj0q5pNFX6uyPpM8n73Cvp9HT2euQkTZf0C0m/VnWfS0T45ptvvvW7AZ8EHgHWJ/f/F/CHA/rMSPo0AxPLtj8BfDbp837gq8n2Q8CxFC/tcwdwSkr79nvA3ybbU4HngO8DC5O264GLa3X/9rOPdfM7rNYb8G7gm8n2POA/0q6pAn+uxuS7k9zuTtpmA/envf/DfDaNwK3Az4Ffq7bPxSN7ZjaYJyn+w9XnVOAdkn4k6R8ljQdOB9ZFRHdE7AQ2AycB84E7k+ctB86VNAFojogno/i32Apgn6t5V8i/A1eX3c9R3L97kvvLgXOp3f2DofexXn6H1ar0uUXEeqDezidXye/OfGBlFD0HNPSNeFWpL1MMdS8k96vqc3HYM7N9RMQtQG9Z00bgjyPiLOAp4DPABGBnWZ/dFP9vtby9vG3XIH0rLiJei4jdSdi5GVhGcbFa32q1wfZjqPaq2z8Ych/r5ndYxQZ+nnlJdXMN+gp/d4Z6jaoj6UPA1ohYUd5cTZ+Lw56ZjcStEfFA3zZwMsW/iMaX9RkP7BjQPlhbeXsqJB0JrAK+HRH/BBTKHh6u5qrfPxh0H+vqd1ilBn5GmYjIpVXMWKjgd6eW/rx9GDhPxWOc3wp8C5he9njqn4vDnpmNxIqyA4EXAQ9QHClaIKlF0kTgeOBRYB1wQdJ3KbAmInYBPZKOlSRgCbCmonuQkHQ4sBL4VER8PWn+saSFyfZSirXV5P7BkPtYN7/DKlb63CTNo3hMVt2o8HdnHbAkWfQym2Jwrspr6kbEWRFxdkQspHic3e8Cy6vpc6mb4WUzG1OXA1+T1ANsAS6NiF2SrqX4F1AGuCoiuiRdB9wkaS3QA3wgeY3LgO8CWYrHnGyo+F4UfRqYDFwtqe/4oyuBayU1AY8DN0dEvkb3Dwbfx08Af1snv8NqdSvFEZ57KR5A/3sp1zPaKvrdkbQGuC95jSsqsoej5w+BG6vlc/FJlc3MzMzqmKdxzczMzOqYw56ZmZlZHXPYMzMzM6tjDntmZmZmdcxhz8zMzKyOOeyZmZmZ1TGHPTMzM7M65rBnZmZmVscc9szMzMzqmMOemZmZWR1z2DMzMzOrYw57ZnZQJC2UtEPSkWVtfynpQwfxmkdJWj8qBe772llJKyStlTR5mL4LJf3L63iP1ZJ+bUDbiZLOOtDXMjM7WA57ZjYaeoBvSFLahYzAEcC0iJgfEdsr+L7vAU6o4PuZmQEOe2Y2Ou4GXgWuKG8cOEInaX3S9llJ35F0p6SNkj4k6QeSnpA0L+neLun7yXOuTp5/pKTlklYlP49MXu+RZDTtkwPe/7ck3Z+M4n1DUiNwAzBH0j8M6Lta0vXJz3skzRjw+Ecl3S1pjaTbJDUldf9bcv/xgaOZki5Kap0FfAj4hKTTJZ0naUPyPt+TNCkZRVwu6T8kPSzpqoP4fZhVteG+Oza6HPbMbLRcDvyBpDkj7N8ZEecD3wMuiIiLgL8E3p88fhjwO8CZwFJJbwG+DFwbEW9Ptv8y6TsDWBwRX+p7cUlTgWuAcyJiPrAD+Ajwv4HHIuIjg9R0b0QsBP4V+HTZa2WAqcC5EbEAaAROSx6eGBEXAu8E/qTstd4NfBS4MCKeB74JfAW4n2LgfHdEnA3cAyxLnvMGiiOAbwP6BVezOjTUd8dGmcOemY2KiHgF+DjFUDPU3y3l07wPJj93AI8l29uBlmT7JxGxMyLywEbgjcCJwKclrQb+DJie9H06InoGvNcxwE8jYndy/0fAm4bZjbuTn/cCv1q2bwWKU9X/LOkfgV+hGPgAHkp+/qKsdoBFwBSgd8B7TAN2JQFwYF2PREQuIvYAncPUalbrHkp+Dvzu2Chz2DOzURMRPwCeoDhlCdAFTE8WRUwCji7vPszLHS/pMEkNwBnAT4GfAZ9KRt8+Atyc9C0M8vyngRMkjUvunw38fJj3PDX5eWbyfgBIOgn4jYj4TeD3Kf7d2Rdch9qPK4AVwOfKaswA24AJko4YpK7hPhOzeuI/7xXSkHYBZlZ3Pk5xVIuI2CLphxSnLjcnt5F6leJ0ajvwrxHxmKQ/Aq6T1AK0AlcO9eSI2CbpM8AqSYXkvf+E4pTvUD4k6RPAHopTyCcm7ZuBPZI2Ad3Ai8DMEezD54CNkm4HHgD+GngcuAT4XlLXdorh+M0jeD0zswOmCAdrM7NkaviyiPhZ2rWYmY0mT+OamZmZ1TGP7JmZmZnVMY/smZmZmdUxhz0zMzOzOuawZ2ZmZlbHHPbMzMzM6pjDnpmZmVkdc9gzMzMzq2P/DT+8yf/Au59yAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 648x216 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "make_plots(\n",
    "    df,\n",
    "    \"./plots/plankton.pdf\",\n",
    "    n_idx=4,\n",
    "    intervals_xlabel=\"Number of plankton\",\n",
    "    plot_classical=False,\n",
    "    true_theta=true_theta,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "850df0cb-c991-487e-be3c-9971c9668596",
   "metadata": {},
   "source": [
    "### Power experiment\n",
    "\n",
    "Find the smallest value of ```n``` such that PPI has power 80% against the null $H_0: \\text{number of plankton} > 10K$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "150cc24f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find n such that we reject H0: Number of plankton > 10K with probability 80% using a test at level alpha\n",
    "num_experiments = 100\n",
    "list_rand_idx = [\n",
    "    np.random.permutation(n_total) for i in range(num_experiments)\n",
    "]\n",
    "\n",
    "\n",
    "def _to_invert_ppi(n):\n",
    "    n = int(n)\n",
    "    nulls_rejected = 0\n",
    "    # Data setup\n",
    "    for i in range(num_experiments):\n",
    "        rand_idx = list_rand_idx[i]\n",
    "        _Yhat = Yhat[rand_idx[:n]]\n",
    "        _Y = Y[rand_idx[:n]]\n",
    "\n",
    "        ppi_count_ci = ppi_distribution_label_shift_ci(\n",
    "            _Y, _Yhat, Yhat_unlabeled, K, nu, alpha, delta, return_counts=True\n",
    "        )\n",
    "        if ppi_count_ci[0] > 10000:\n",
    "            nulls_rejected += 1\n",
    "    return nulls_rejected / num_experiments - 0.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "57d481cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The PPI test requires n=25513 labeled data points to reject the null.\n"
     ]
    }
   ],
   "source": [
    "n_ppi = int(brentq(_to_invert_ppi, 1000, 100000, xtol=1))\n",
    "print(\n",
    "    f\"The PPI test requires n={n_ppi} labeled data points to reject the null.\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}