{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bdb22654-f761-45a2-a033-7c03b76be0d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_env import *\n",
    "from sklearn.neighbors import KNeighborsClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5115bd8f-a519-430a-abdb-ab35bf53e95e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_group(alpha, beta, is_positive):\n",
    "    pos = scipy.stats.uniform(loc=1)\n",
    "    neg = scipy.stats.uniform(loc=-2)\n",
    "    if is_positive:\n",
    "        return evoml.RandomVariablePopulationGroup(\n",
    "            rv_vec=[\n",
    "                evoml.MixtureRandomVariable(\n",
    "                    rv_vec=[\n",
    "                        neg,\n",
    "                        pos,\n",
    "                    ],\n",
    "                    p=[\n",
    "                        (1-beta)/(1-beta+beta*alpha),\n",
    "                        (beta*alpha)/(1-beta+beta*alpha),\n",
    "                    ],\n",
    "                ),\n",
    "                pos,\n",
    "            ],\n",
    "            p = [\n",
    "                1-beta*(1-alpha),\n",
    "                beta*(1-alpha),\n",
    "            ]\n",
    "        )\n",
    "    else:\n",
    "        return evoml.RandomVariablePopulationGroup(\n",
    "            rv_vec=[\n",
    "                neg,\n",
    "                evoml.MixtureRandomVariable(\n",
    "                    rv_vec=[\n",
    "                        pos,\n",
    "                        neg,\n",
    "                    ],\n",
    "                    p=[\n",
    "                        (1-beta)/(1-beta+beta*alpha),\n",
    "                        (beta*alpha)/(1-beta+beta*alpha),\n",
    "                    ],\n",
    "                ),\n",
    "            ],\n",
    "            p = [\n",
    "                beta*(1-alpha),\n",
    "                1-beta*(1-alpha),\n",
    "            ]\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "61b6ff2b-c538-4dd2-b3ee-592c32aef053",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "alpha=0.2\n",
      "beta=0.8\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x16b978fe0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 34560 characters]",
      "text/plain": [
       "<Figure size 1000x300 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 311,
       "width": 989
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "alpha = param(0.2,'alpha')\n",
    "beta = param(0.8,'beta')\n",
    "pop = evoml.Population(groups=[\n",
    "    build_group(alpha=alpha, beta=beta, is_positive=is_pos) for is_pos in [True, False]\n",
    "])\n",
    "\n",
    "fig,axs=evoml.create_figure(1,3)\n",
    "for p,ax in zip([0,0.5,1],axs):\n",
    "    X,y,_ = pop.sample_from_mixture(5000, p=[1-p,p])\n",
    "    ax.hist(\n",
    "        [X[y==0].ravel(), X[y==1].ravel()],\n",
    "        bins=np.linspace(-3,3,40),\n",
    "    )\n",
    "    ax.set(\n",
    "        title=f'$p=({1-p},{p})$',\n",
    "        xlabel='$x$',\n",
    "        ylabel='Count',\n",
    "    )\n",
    "\n",
    "axs[0].legend([\n",
    "    '$y=0$',\n",
    "    '$y=1$',\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4144ec17-544b-408a-9a94-b0b8dc06b684",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f55324ef8e744cf9bb0c6e82862ab0b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "k_vec = [1,3,5]\n",
    "n_reps = 20\n",
    "\n",
    "n_train = 10000\n",
    "n_test = 10000\n",
    "\n",
    "results = []\n",
    "p_vec = np.linspace(0,1,7)\n",
    "for rep in tqdm(range(n_reps)):\n",
    "    for k in k_vec:\n",
    "        clf = KNeighborsClassifier(n_neighbors=k)\n",
    "        for p in p_vec:\n",
    "            clf.fit(*pop.sample_from_mixture(n=n_train,p=[1-p,p])[:2])\n",
    "            results.append({\n",
    "                'k': k,\n",
    "                'p': p,\n",
    "                'rep': rep,\n",
    "                'acc_A': clf.score(*pop.sample_from_group(n=n_test,g=0)),\n",
    "                'acc_B': clf.score(*pop.sample_from_group(n=n_test,g=1)),\n",
    "            })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "553c7bf8-019d-4606-8496-bfe7f75b5b9a",
   "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>k</th>\n",
       "      <th>p</th>\n",
       "      <th>rep</th>\n",
       "      <th>acc_A</th>\n",
       "      <th>acc_B</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0</td>\n",
       "      <td>0.7455</td>\n",
       "      <td>0.8031</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0.166667</td>\n",
       "      <td>0</td>\n",
       "      <td>0.7350</td>\n",
       "      <td>0.7569</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>0.333333</td>\n",
       "      <td>0</td>\n",
       "      <td>0.7238</td>\n",
       "      <td>0.7398</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>0.500000</td>\n",
       "      <td>0</td>\n",
       "      <td>0.7341</td>\n",
       "      <td>0.7338</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>0.666667</td>\n",
       "      <td>0</td>\n",
       "      <td>0.7477</td>\n",
       "      <td>0.7326</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   k         p  rep   acc_A   acc_B\n",
       "0  1  0.000000    0  0.7455  0.8031\n",
       "1  1  0.166667    0  0.7350  0.7569\n",
       "2  1  0.333333    0  0.7238  0.7398\n",
       "3  1  0.500000    0  0.7341  0.7338\n",
       "4  1  0.666667    0  0.7477  0.7326"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results_df = pd.DataFrame(results)\n",
    "results_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "aba32e30-07d8-4df7-b3f2-e5e9ad8af276",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Figure saved as ./output/knn_theoretical_curves.pdf\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"data:text/plain;base64,[encoded data, 351496 characters]\" download=\"knn_theoretical_curves.pdf\">Download knn_theoretical_curves.pdf</a>"
      ],
      "text/plain": [
       "<evoml.analysis.DownloadableIO at 0x177bc8c70>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 142876 characters]",
      "text/plain": [
       "<Figure size 950x250 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 261,
       "width": 940
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,axs=evoml.create_figure(\n",
    "    1,3,\n",
    "    figsize=(9.5,2.5)\n",
    ")\n",
    "\n",
    "for k,ax in zip(k_vec,axs):\n",
    "    (\n",
    "        results_df\n",
    "        .query('k==@k')\n",
    "        .assign(acc_p=lambda df: (1-df['p'])*df['acc_A'] + df['p']*df['acc_B'])\n",
    "        .groupby('p')\n",
    "        .mean()\n",
    "        .pipe(lambda df: df[[c for c in df.columns if c.startswith('acc')]])\n",
    "        .plot.line(\n",
    "            ax=ax,\n",
    "            marker='.',\n",
    "            linestyle='',\n",
    "            title=f'k-NN with $k={k}$',\n",
    "            legend=False,\n",
    "            xlabel='$p_B$',\n",
    "            ylabel='Accuracy',\n",
    "            color=[\n",
    "                evoml.group_colors[0],\n",
    "                evoml.group_colors[1],\n",
    "                'tab:cyan',\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    ax.set(\n",
    "        xlabel=evoml.analysis.composition_label,\n",
    "        xticks=evoml.analysis.composition_tickpos[::-1],\n",
    "        xticklabels=evoml.analysis.composition_ticklabels,\n",
    "    )\n",
    "    ax.yaxis.set_major_formatter(matplotlib.ticker.PercentFormatter(xmax=1,decimals=1))\n",
    "    ## Theory\n",
    "    p = np.linspace(0,1,100)\n",
    "    binom_cdf = lambda p: scipy.stats.binom.cdf(n=k,k=k//2,p=p)\n",
    "    gamma = p - 2*p*beta + beta\n",
    "    noise_coeff_pos = beta*(1-p)/gamma\n",
    "    noise_coeff_neg = beta*p/(1-gamma)\n",
    "    # exp_yhat_pos = 1-2*binom_cdf(1-noise_coeff_pos*alpha)\n",
    "    # exp_yhat_neg = -1+2*binom_cdf(1-noise_coeff_neg*alpha)\n",
    "    exp_yhat_pos = -1+2*binom_cdf(noise_coeff_pos*alpha)\n",
    "    exp_yhat_neg = 1-2*binom_cdf(noise_coeff_neg*alpha)\n",
    "    acc_pos = (1+(1-2*alpha)*exp_yhat_pos)/2\n",
    "    acc_neg = (1+(-1)*exp_yhat_neg)/2\n",
    "    acc_A = beta*acc_pos + (1-beta)*acc_neg\n",
    "    acc_B = acc_A[::-1]\n",
    "    line_params = {\n",
    "        'linestyle': '--',\n",
    "        'zorder': -1,\n",
    "    }\n",
    "    ax.plot(p, acc_A, label='acc_A',color=evoml.group_colors[0], **line_params)\n",
    "    ax.plot(p, acc_B, label='acc_B',color=evoml.group_colors[1], **line_params)\n",
    "    ax.plot(p, (1-p)*acc_A+p*acc_B, color='tab:cyan', **line_params)\n",
    "\n",
    "axs[0].legend(\n",
    "    [\n",
    "        '$\\\\mathrm{acc}_A$',\n",
    "        '$\\\\mathrm{acc}_B$',\n",
    "        '$\\\\mathrm{acc}_p$',\n",
    "    ],\n",
    "    loc=(0.18,0.60),\n",
    ")\n",
    "\n",
    "fig.save_and_download('./output/knn_theoretical_curves.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a4742443-b2c6-4781-9537-65093de8a680",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameters saved as: ./output/knn.json\n"
     ]
    }
   ],
   "source": [
    "param_tracker.save('./output/knn.json')"
   ]
  }
 ],
 "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.23"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
