{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import json\n",
    "import blahut_arimoto as ba\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>model</th>\n",
       "      <th>in_dataset_name</th>\n",
       "      <th>method</th>\n",
       "      <th>method_kwargs</th>\n",
       "      <th>scores</th>\n",
       "      <th>labels</th>\n",
       "      <th>idx</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>resnet18_cifar10</td>\n",
       "      <td>cifar10</td>\n",
       "      <td>dice</td>\n",
       "      <td>{}</td>\n",
       "      <td>[9.59467601776123, 7.032630443572998, 8.800002...</td>\n",
       "      <td>[0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...</td>\n",
       "      <td>[0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>resnet18_cifar10</td>\n",
       "      <td>cifar10</td>\n",
       "      <td>doctor</td>\n",
       "      <td>{}</td>\n",
       "      <td>[-0.0009341239929199219, -0.15732842683792114,...</td>\n",
       "      <td>[0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...</td>\n",
       "      <td>[0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>resnet18_cifar10</td>\n",
       "      <td>cifar10</td>\n",
       "      <td>energy</td>\n",
       "      <td>{}</td>\n",
       "      <td>[9.047059059143066, 5.1664018630981445, 7.7167...</td>\n",
       "      <td>[0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...</td>\n",
       "      <td>[0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>resnet18_cifar10</td>\n",
       "      <td>cifar10</td>\n",
       "      <td>igeood_logits</td>\n",
       "      <td>{}</td>\n",
       "      <td>[2.813105821609497, 2.739166259765625, 2.79803...</td>\n",
       "      <td>[0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...</td>\n",
       "      <td>[0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>resnet18_cifar10</td>\n",
       "      <td>cifar10</td>\n",
       "      <td>kl_matching</td>\n",
       "      <td>{}</td>\n",
       "      <td>[-0.00017658778233453631, -0.4094327688217163,...</td>\n",
       "      <td>[0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...</td>\n",
       "      <td>[0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               model in_dataset_name         method method_kwargs  \\\n",
       "6   resnet18_cifar10         cifar10           dice            {}   \n",
       "14  resnet18_cifar10         cifar10         doctor            {}   \n",
       "1   resnet18_cifar10         cifar10         energy            {}   \n",
       "4   resnet18_cifar10         cifar10  igeood_logits            {}   \n",
       "2   resnet18_cifar10         cifar10    kl_matching            {}   \n",
       "\n",
       "                                               scores  \\\n",
       "6   [9.59467601776123, 7.032630443572998, 8.800002...   \n",
       "14  [-0.0009341239929199219, -0.15732842683792114,...   \n",
       "1   [9.047059059143066, 5.1664018630981445, 7.7167...   \n",
       "4   [2.813105821609497, 2.739166259765625, 2.79803...   \n",
       "2   [-0.00017658778233453631, -0.4094327688217163,...   \n",
       "\n",
       "                                               labels  \\\n",
       "6   [0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...   \n",
       "14  [0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...   \n",
       "1   [0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...   \n",
       "4   [0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...   \n",
       "2   [0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...   \n",
       "\n",
       "                                                  idx  \n",
       "6   [0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...  \n",
       "14  [0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...  \n",
       "1   [0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...  \n",
       "4   [0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...  \n",
       "2   [0, 6, 2, 0, 1, 2, 8, 7, 4, 9, 9, 0, 9, 4, 7, ...  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scores_path1 = \"results/scod_benchmark_cifar10/scores.csv\"\n",
    "scores_path2 = \"results/scod_benchmark_cifar100/scores.csv\"\n",
    "dataset_list = list(\n",
    "    {\n",
    "        \"cifar100\": \"test\",\n",
    "        \"svhn\": \"test\",\n",
    "        \"isun\": None,\n",
    "        \"lsun_c\": None,\n",
    "        \"lsun_r\": None,\n",
    "        \"tiny_imagenet_c\": None,\n",
    "        \"tiny_imagenet_r\": None,\n",
    "        \"textures\": None,\n",
    "        \"places365\": None,\n",
    "        \"uniform\": None,\n",
    "        \"gaussian\": None,\n",
    "    }.keys()\n",
    ")\n",
    "scores1 = pd.read_csv(scores_path1)\n",
    "scores2 = pd.read_csv(scores_path2)\n",
    "scores = pd.concat([scores1, scores2])\n",
    "# drop duplicates\n",
    "scores = scores.drop_duplicates(subset=[\"model\", \"method\"], keep=\"last\")\n",
    "# drop gradnorm and vim\n",
    "scores = scores[scores.method != \"gradnorm\"]\n",
    "scores = scores[scores.method != \"vim\"]\n",
    "# drop na\n",
    "scores = scores.dropna()\n",
    "# sort by model and method\n",
    "scores = scores.sort_values(by=[\"model\", \"method\"])\n",
    "scores.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 14/14 [00:15<00:00,  1.08s/it]\n",
      "100%|██████████| 14/14 [00:15<00:00,  1.09s/it]\n",
      "100%|██████████| 14/14 [00:18<00:00,  1.29s/it]\n",
      "100%|██████████| 14/14 [00:16<00:00,  1.15s/it]\n",
      "100%|██████████| 4/4 [01:19<00:00, 19.95s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('resnet18_cifar10', 'doctor'), ('resnet18_cifar10', 'kl_matching'), ('resnet18_cifar10', 'knn_euclides'), ('resnet18_cifar10', 'mahalanobis'), ('resnet18_cifar10', 'relative_mahalanobis'), ('resnet18_cifar100', 'doctor'), ('resnet18_cifar100', 'kl_matching'), ('resnet18_cifar100', 'knn_euclides'), ('resnet18_cifar100', 'mahalanobis'), ('resnet18_cifar100', 'relative_mahalanobis'), ('resnet34_cifar10', 'doctor'), ('resnet34_cifar10', 'kl_matching'), ('resnet34_cifar10', 'knn_euclides'), ('resnet34_cifar10', 'mahalanobis'), ('resnet34_cifar10', 'relative_mahalanobis'), ('resnet34_cifar100', 'doctor'), ('resnet34_cifar100', 'kl_matching'), ('resnet34_cifar100', 'knn_euclides'), ('resnet34_cifar100', 'mahalanobis'), ('resnet34_cifar100', 'relative_mahalanobis')]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 66684 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 45292 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 72376 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 50056 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 42516 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 45432 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53164 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 64596 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 50460 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 45856 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 72888 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 61708 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53056 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 67348 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 68844 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48944 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 73836 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48624 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 44208 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48692 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 80764 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 69876 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55068 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 45252 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 83328 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 59368 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 81884 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 54172 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 60000 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 47044 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55968 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 56072 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 46008 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 49256 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 44940 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53440 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 47760 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 43876 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 60436 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 46296 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53320 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 70804 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 62672 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 44136 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 60064 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 52048 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 44320 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 44828 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 57876 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 59240 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 49344 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 43548 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 76636 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 58228 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 66192 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 59428 characters]",
      "text/plain": [
       "<Figure size 4400x300 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wrong_methods = {}\n",
    "for model in tqdm(scores[\"model\"].unique()):\n",
    "    df = scores.query(\"model == @model\")\n",
    "    methods = list(df[\"method\"])\n",
    "    labels_array = np.concatenate(df[\"labels\"].map(lambda x: np.array(json.loads(x)).reshape(-1, 1)).values, 1)\n",
    "    scores_array = np.concatenate(df[\"scores\"].map(lambda x: np.array(json.loads(x)).reshape(-1, 1)).values, 1)\n",
    "    datasets_idx = np.concatenate(df[\"idx\"].map(lambda x: np.array(json.loads(x)).reshape(-1, 1)).values, 1)\n",
    "\n",
    "    # compute blahut arimoto weights with p_value\n",
    "    ## get probs from scores\n",
    "    X = scores_array[datasets_idx[:, 0] == 0][:1000]\n",
    "    probs_array = ba.p_value_fn(scores_array, X=X)\n",
    "    ## compute optimal weights\n",
    "    device = \"mps\"\n",
    "    w = ba.blahut_arimoto(\n",
    "        torch.from_numpy(probs_array).float(),\n",
    "        device=device,\n",
    "        threshold=1e-3,\n",
    "        verbose=False,\n",
    "    ).numpy()\n",
    "    ba_scores = np.sum(w * probs_array, 1)\n",
    "    scores_array = np.concatenate([scores_array, ba_scores.reshape(-1, 1)], 1)\n",
    "    labels_array = np.concatenate([labels_array, labels_array[:, 1].reshape(-1, 1)], 1)\n",
    "    datasets_idx = np.concatenate([datasets_idx, datasets_idx[:, 1].reshape(-1, 1)], 1)\n",
    "    methods.append(\"blahut_arimoto\")\n",
    "\n",
    "    assert labels_array.shape == scores_array.shape\n",
    "    # for each method, plot the histogram\n",
    "    for col in tqdm(range(labels_array.shape[1])):\n",
    "        fig, ax = plt.subplots(1, len(dataset_list), figsize=(4 * len(dataset_list), 3))\n",
    "        for i, ds in enumerate(dataset_list):\n",
    "            filt_in = np.where(datasets_idx[:, col] == 0)[0]\n",
    "            filt_out = np.where(datasets_idx[:, col] == i + 1)[0]\n",
    "\n",
    "            # control over pi (pi=0.5)\n",
    "            min_size = min(filt_in.shape[0], filt_out.shape[0])\n",
    "            filt_in = filt_in[:min_size]\n",
    "            filt_out = filt_out[:min_size]\n",
    "            # TODO: try probs_array\n",
    "            y_in = labels_array[filt_in, col].reshape(-1)\n",
    "            x_in = scores_array[filt_in, col].reshape(-1)\n",
    "            y_out = labels_array[filt_out, col].reshape(-1)\n",
    "            x_out = scores_array[filt_out, col].reshape(-1)\n",
    "\n",
    "            in_scores = x_in[y_in == 0]\n",
    "            out_scores = np.concatenate([x_in[y_in == 1], x_out])\n",
    "\n",
    "            sns.distplot(\n",
    "                in_scores,\n",
    "                ax=ax[i],\n",
    "                label=\"in\",\n",
    "                norm_hist=True,\n",
    "                kde=False,\n",
    "                color=\"green\",\n",
    "                hist_kws={\"alpha\": 0.5},\n",
    "            )\n",
    "            sns.distplot(\n",
    "                out_scores,\n",
    "                ax=ax[i],\n",
    "                label=\"out\",\n",
    "                norm_hist=True,\n",
    "                kde=False,\n",
    "                color=\"red\",\n",
    "                hist_kws={\"alpha\": 0.5},\n",
    "            )\n",
    "            ax[i].set_title(ds)\n",
    "            ax[i].legend()\n",
    "\n",
    "            try:\n",
    "                assert in_scores.mean() > out_scores.mean()\n",
    "                assert in_scores.mean() > 0\n",
    "                assert out_scores.mean() > 0\n",
    "            except:\n",
    "                wrong_methods[(model, methods[col])] = True\n",
    "\n",
    "        plt.suptitle(f\"{model} - {methods[col]}\")\n",
    "        plt.tight_layout()\n",
    "        os.makedirs(\"figures/histograms\", exist_ok=True)\n",
    "        plt.savefig(f\"figures/histograms/{model}_{methods[col]}.png\")\n",
    "print(list(wrong_methods.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "These methods do not satisfy the assumptions:\n",
      "{'relative_mahalanobis', 'doctor', 'mahalanobis', 'kl_matching', 'knn_euclides'}\n"
     ]
    }
   ],
   "source": [
    "wm = list(wrong_methods.keys())\n",
    "print(\"These methods do not satisfy the assumptions:\")\n",
    "print(set(w[1] for w in wm))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
