{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c3ff15b1-c922-47fa-be98-349747bf3765",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_env import *\n",
    "import glob\n",
    "import os\n",
    "import ast\n",
    "import evoml.analysis\n",
    "import evoml.dynamics\n",
    "import zipfile\n",
    "import json\n",
    "\n",
    "import scipy.interpolate\n",
    "\n",
    "def ci(s):\n",
    "    mu = s.mean()\n",
    "    std = s.std()\n",
    "    n = len(s)\n",
    "    return scipy.stats.norm.interval(0.99,loc=mu,scale=std/np.sqrt(n))\n",
    "\n",
    "def ci_pm(s):\n",
    "    l,h = ci(s)\n",
    "    return h-(l+h)/2\n",
    "\n",
    "def report_ci(s):\n",
    "    return '{:.1f} \\\\pm {:.1%}'.format(\n",
    "        s.mean()*100,\n",
    "        ci_pm(s),\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d6c8bee-6656-4ccf-962b-378b5ef4caef",
   "metadata": {},
   "source": [
    "## CIFAR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6e224a5a-92bd-45ac-a8ec-d5a45e16b592",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: xlabel='p'>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 68080 characters]",
      "text/plain": [
       "<Figure size 400x300 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 290,
       "width": 378
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "dcts = ast.literal_eval(open('./data/cifar10_augmentations.repr','r').read())\n",
    "\n",
    "cifar_augmentations_game_df = (\n",
    "    pd.DataFrame(dcts)\n",
    "    .transpose()\n",
    "    .rename({\n",
    "        ('test',False): 'acc_A',\n",
    "        ('test',True): 'acc_B',\n",
    "    },axis=1)\n",
    "    .assign(\n",
    "        p=lambda df: df['config'].map(lambda d: float(d[('training','flip_prob')])),\n",
    "        acc_p=lambda df: (1-df['p'])*df['acc_A'] + df['p']*df['acc_B'],\n",
    "    )\n",
    "    .set_index('p')\n",
    "    [['acc_A','acc_B','acc_p']]\n",
    "    .sort_index()\n",
    ")\n",
    "\n",
    "cifar_augmentations_game_df.groupby(level=0).mean().plot.line(figsize=(4,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b4071236-2651-4289-8add-58cbf998d9df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cifar_augmentation_single_group_accuracy=92.6%\n",
      "cifar_augmentation_single_group_accuracy_confidence=92.6 \\pm 0.1%\n",
      "cifar_augmentation_coexistence_accuracy=93.5%\n",
      "cifar_augmentation_coexistence_accuracy_confidence=93.5 \\pm 0.1%\n",
      "cifar_augmentation_n_repetitions=20\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "np.int64(20)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "param(\n",
    "    cifar_augmentations_game_df.loc[0,'acc_A'].mean(),\n",
    "    'cifar_augmentation_single_group_accuracy',\n",
    "    '{:.1%}',\n",
    ")\n",
    "param(\n",
    "    cifar_augmentations_game_df.loc[0,'acc_A'].pipe(report_ci),\n",
    "    'cifar_augmentation_single_group_accuracy_confidence',\n",
    ")\n",
    "param(\n",
    "    cifar_augmentations_game_df.loc[0.5,'acc_A'].mean(),\n",
    "    'cifar_augmentation_coexistence_accuracy',\n",
    "    '{:.1%}',\n",
    ")\n",
    "param(\n",
    "    cifar_augmentations_game_df.loc[0.5,'acc_A'].pipe(report_ci),\n",
    "    'cifar_augmentation_coexistence_accuracy_confidence',\n",
    ")\n",
    "param(\n",
    "    cifar_augmentations_game_df.loc[0,'acc_A'].count(),\n",
    "    'cifar_augmentation_n_repetitions',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4eac9f55-52a7-4ac9-9450-7b3e3fdc7799",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.9348  , 0.934405])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class InterpolatedFitness:\n",
    "    def __init__(self, game_df, sample=True, interpolation_kind='linear'):\n",
    "        self.game_df = game_df\n",
    "        self.sample = sample\n",
    "        self.interpolation_kind = interpolation_kind\n",
    "        x = self.game_df.groupby(level=0).mean().index.to_numpy()\n",
    "        agg = (lambda dfg: dfg.sample(1)) if self.sample else (lambda dfg: dfg.mean())\n",
    "        a = self.game_df['acc_A'].groupby(level=0).pipe(agg).to_numpy()\n",
    "        b = self.game_df['acc_B'].groupby(level=0).pipe(agg).to_numpy()\n",
    "        interp = lambda y: scipy.interpolate.interp1d(\n",
    "            x=x,\n",
    "            y=y,\n",
    "            kind=self.interpolation_kind,\n",
    "        )\n",
    "        self.interp_a = interp(a)\n",
    "        self.interp_b = interp(b)\n",
    "\n",
    "    def __call__(self, p):\n",
    "        assert len(p)==2, 'Interpolation is implemented in two-group settings'\n",
    "        pb = p[-1]\n",
    "        return np.array([self.interp_a(pb), self.interp_b(pb)])\n",
    "\n",
    "class ReflectedFitness:\n",
    "    def __init__(self, fitness, p_eq):\n",
    "        self.fitness = fitness\n",
    "        self.p_eq = np.array(p_eq)\n",
    "\n",
    "    def __call__(self, p):\n",
    "        return self.fitness(2*self.p_eq - p)\n",
    "        \n",
    "\n",
    "fitness = InterpolatedFitness(\n",
    "    game_df=cifar_augmentations_game_df.groupby(level=0).mean(),\n",
    ")\n",
    "fitness([0.5,0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "55793853-955f-4038-8180-859f1df16e5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 2000\n",
    "p0 = 0.2\n",
    "res1 = evoml.dynamics.discrete_replicator(\n",
    "    p0=[1-p0,p0],\n",
    "    fitness=fitness,\n",
    "    T=T,\n",
    "    return_fitness=True,\n",
    ")\n",
    "welf1 = (res1[0]*res1[1]).sum(axis=1)\n",
    "res2 = evoml.dynamics.discrete_replicator(\n",
    "    p0=[1-p0,p0],\n",
    "    fitness=ReflectedFitness(fitness, p_eq=(0.5,0.5)),\n",
    "    T=T,\n",
    "    return_fitness=True,\n",
    ")\n",
    "welf2 = (res2[0]*res2[1]).sum(axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9c55538c-caa7-4397-ab3b-eaa9ceceaa37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54900 characters]",
      "text/plain": [
       "<Figure size 300x300 with 2 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 311,
       "width": 290
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,axs=evoml.create_figure(2,1,figsize=(3,3))\n",
    "ax = axs[0]\n",
    "ax.plot(\n",
    "    res1[0][:,0],\n",
    "    # label='Unstabilized',\n",
    "    label='$\\\\mathcal{A}(\\\\bm{p})$',\n",
    ")\n",
    "ax.plot(\n",
    "    res2[0][:,0],\n",
    "    # label='Stabilized',\n",
    "    label='$\\\\mathcal{A}\\'(\\\\bm{p})$',\n",
    ")\n",
    "\n",
    "ax.set(\n",
    "    title='Stabilization by Reweighing', #  $F_k(\\\\bm{p})$\n",
    "    ylabel=evoml.analysis.composition_label,\n",
    "    yticks=evoml.analysis.composition_tickpos + [1-p0],\n",
    "    yticklabels=evoml.analysis.composition_ticklabels + ['$\\\\bm{p}^0$'],\n",
    "    # ylim=(-0.03,1.03),\n",
    "    ylim=(0.5-0.03,1+0.03),\n",
    ")\n",
    "ax.set(\n",
    "    xticklabels=[],\n",
    "    # yticklabels=[],\n",
    ")\n",
    "ax.axhline(0.5,**evoml.background_line_style)\n",
    "ax.axhline(1,**evoml.background_line_style)\n",
    "ax.axhline(0.0,**evoml.background_line_style)\n",
    "# ax.legend()\n",
    "\n",
    "ax=axs[1]\n",
    "ax.plot(welf1)\n",
    "ax.plot(welf2)\n",
    "ax.set(\n",
    "    ylabel='Accuracy',\n",
    "    xlabel='Time ($t$)',\n",
    ")\n",
    "ax.yaxis.set_major_formatter(\n",
    "    matplotlib.ticker.PercentFormatter(xmax=1,decimals=1),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8add781-0b5b-4990-b570-83d7dfc29e97",
   "metadata": {},
   "source": [
    "## MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5106c7aa-3e50-48f6-ba5a-c23ccfe6a03b",
   "metadata": {},
   "outputs": [],
   "source": [
    "extract_p = lambda path: (\n",
    "    int(os.path.basename(path).split('.')[0].split('_')[1])\n",
    "    /int(os.path.basename(path).split('.')[0].split('_')[2])\n",
    ")\n",
    "\n",
    "zf = zipfile.ZipFile('./data/mnist_label_noise.zip')\n",
    "results = []\n",
    "for f in zf.filelist:\n",
    "    path = f.filename\n",
    "    for dct in json.load(zf.open(f)):\n",
    "        results.append({\n",
    "            'filename': os.path.basename(path),\n",
    "            'p': extract_p(path=path),\n",
    "            'epoch': dct['epoch'],\n",
    "        } | {\n",
    "            f'{k}_{suffix}': v\n",
    "            for group, suffix in [('group_a','A'), ('group_b','B'), ('training_set','train')]\n",
    "            for k,v in dct[group].items()\n",
    "        })\n",
    "\n",
    "mnist_df = pd.DataFrame(results)\n",
    "selected_epoch = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "edb98b81-3c4f-4d22-841f-a5c5a9bf6e05",
   "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>acc_A</th>\n",
       "      <th>acc_B</th>\n",
       "      <th>acc_p</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>p</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0.0</th>\n",
       "      <td>0.801125</td>\n",
       "      <td>0.811083</td>\n",
       "      <td>0.801125</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0.1</th>\n",
       "      <td>0.801996</td>\n",
       "      <td>0.807925</td>\n",
       "      <td>0.802589</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0.2</th>\n",
       "      <td>0.800528</td>\n",
       "      <td>0.807234</td>\n",
       "      <td>0.801869</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0.3</th>\n",
       "      <td>0.803254</td>\n",
       "      <td>0.805254</td>\n",
       "      <td>0.803854</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0.4</th>\n",
       "      <td>0.802266</td>\n",
       "      <td>0.804365</td>\n",
       "      <td>0.803106</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0.5</th>\n",
       "      <td>0.804024</td>\n",
       "      <td>0.803595</td>\n",
       "      <td>0.803810</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0.6</th>\n",
       "      <td>0.806020</td>\n",
       "      <td>0.801155</td>\n",
       "      <td>0.803101</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        acc_A     acc_B     acc_p\n",
       "p                                \n",
       "0.0  0.801125  0.811083  0.801125\n",
       "0.1  0.801996  0.807925  0.802589\n",
       "0.2  0.800528  0.807234  0.801869\n",
       "0.3  0.803254  0.805254  0.803854\n",
       "0.4  0.802266  0.804365  0.803106\n",
       "0.5  0.804024  0.803595  0.803810\n",
       "0.6  0.806020  0.801155  0.803101"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mnist_game_df = (\n",
    "    mnist_df\n",
    "    .query('epoch==@selected_epoch')\n",
    "    .assign(acc_p=lambda df: df['acc_A']*(1-df['p']) + df['acc_B']*df['p'])\n",
    "    .groupby('p')\n",
    "    [['acc_A','acc_B','acc_p']]\n",
    "    .mean()\n",
    "    # .iloc[1::2]\n",
    ")\n",
    "mnist_game_df.loc[:0.6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c997b8d0-15ae-480c-b708-c4d9a6fc5e20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mnist_label_noise_train_accuracy=98.7%\n",
      "mnist_label_noise_train_accuracy_confidence=98.7 \\pm 0.1%\n",
      "mnist_label_noise_coexistence_accuracy=80.4%\n",
      "mnist_label_noise_coexistence_accuracy_confidence=80.4 \\pm 0.2%\n",
      "mnist_label_noise_n_repetitions=50\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "np.int64(50)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "param(\n",
    "    mnist_df.query('epoch==@selected_epoch')['acc_train'].mean(),\n",
    "    'mnist_label_noise_train_accuracy',\n",
    "    '{:.1%}',\n",
    ")\n",
    "param(\n",
    "    mnist_df.query('epoch==@selected_epoch')['acc_train'].pipe(report_ci),\n",
    "    'mnist_label_noise_train_accuracy_confidence',\n",
    ")\n",
    "param(\n",
    "    mnist_df.query('epoch==@selected_epoch and p==0.5')['acc_A'].mean(),\n",
    "    'mnist_label_noise_coexistence_accuracy',\n",
    "    '{:.1%}',\n",
    ")\n",
    "param(\n",
    "    mnist_df.query('epoch==@selected_epoch and p==0.5')['acc_A'].pipe(report_ci),\n",
    "    'mnist_label_noise_coexistence_accuracy_confidence',\n",
    ")\n",
    "param(\n",
    "    mnist_df.query('epoch==@selected_epoch and p==0.5')['acc_A'].count(),\n",
    "    'mnist_label_noise_n_repetitions',\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19717490-2ee2-4fa0-adae-cfaac8e7437a",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "060b11af-5ccb-48f1-bb91-8de94d8c9354",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Figure saved as ./output/empirical.pdf\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"data:text/plain;base64,[encoded data, 377768 characters]\" download=\"empirical.pdf\">Download empirical.pdf</a>"
      ],
      "text/plain": [
       "<evoml.analysis.DownloadableIO at 0x318db3e50>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 154864 characters]",
      "text/plain": [
       "<Figure size 950x210 with 4 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 230,
       "width": 940
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,axs = plt.subplot_mosaic(\n",
    "    mosaic=[['cifar','stab1','mnist'],['cifar','stab2','mnist']],\n",
    "    figsize=(9.5,2.1),\n",
    "    constrained_layout=dict(\n",
    "        w_pad=0.15,\n",
    "        hspace=0.0,\n",
    "        h_pad=0.0,\n",
    "    ),\n",
    ")\n",
    "\n",
    "ax = axs['cifar']\n",
    "evoml.plot_population_game(\n",
    "    cifar_augmentations_game_df.groupby(level=0).mean(),\n",
    "    ax=ax,\n",
    "    equilibria=[\n",
    "        (0,True),\n",
    "        (0.5,False),\n",
    "        (1,True),\n",
    "    ],\n",
    "    equilibria_kwargs={\n",
    "        'd': 0.1,\n",
    "    },\n",
    "    annotations={\n",
    "        'acc_A': (0.0,0.9327),\n",
    "        'acc_B': (0.88,0.9327),\n",
    "    },\n",
    ")\n",
    "ax.set(\n",
    "    title='CIFAR-10 With Horizontal Flips',\n",
    ")\n",
    "\n",
    "ax = axs['mnist']\n",
    "evoml.plot_population_game(\n",
    "    mnist_game_df,\n",
    "    ax=ax,\n",
    "    equilibria=[\n",
    "        (0.5,True),\n",
    "    ],\n",
    "    equilibria_kwargs={\n",
    "        'd': 0.1,\n",
    "    },\n",
    "    annotations={\n",
    "        'acc_A': (0.82,0.81),\n",
    "        'acc_B': (0.07,0.81),\n",
    "    },\n",
    ")\n",
    "ax.set(\n",
    "    title='MNIST With Label Noise',\n",
    ")\n",
    "\n",
    "\n",
    "ax = axs['stab1']\n",
    "label1='$\\\\mathcal{A}(\\\\bm{p})$'\n",
    "label2='$\\\\mathcal{A}\\'(\\\\bm{p})$'\n",
    "c1=evoml.analysis.group_colors[0]\n",
    "c2='tab:orange'\n",
    "ax.plot(\n",
    "    res1[0][:,0],\n",
    "    label=label1,\n",
    "    color=c1,\n",
    "    linestyle='--',\n",
    "    \n",
    ")\n",
    "ax.plot(\n",
    "    res2[0][:,0],\n",
    "    label=label2,\n",
    "    color=c2,\n",
    ")\n",
    "\n",
    "ax.set(\n",
    "    title='CIFAR-10 Simulated Stabilization', #  $F_k(\\\\bm{p})$\n",
    "    ylabel='Population',\n",
    "    yticks=evoml.analysis.composition_tickpos + [1-p0],\n",
    "    yticklabels=evoml.analysis.composition_ticklabels + ['$\\\\bm{p}^0$'],\n",
    "    ylim=(0.5-0.03,1+0.03),\n",
    ")\n",
    "ax.set(\n",
    "    xticklabels=[],\n",
    ")\n",
    "ax.axhline(0.5,**evoml.background_line_style)\n",
    "ax.axhline(1,**evoml.background_line_style)\n",
    "ax.axhline(0.0,**evoml.background_line_style)\n",
    "ax.legend(loc=(0.6,0.29))\n",
    "\n",
    "\n",
    "with evoml.analysis.autoscale_turned_off(ax=ax):\n",
    "    x = ax.get_xlim()[0]\n",
    "    for y,c,m in [\n",
    "        (1.0, c1, 'o'),\n",
    "        (0.9, c1, '^'),\n",
    "        (0.6, c2, 'v'),\n",
    "        (0.5, c2, 'o'),\n",
    "    ]:\n",
    "        ax.plot(\n",
    "            x,y,\n",
    "            marker=m,\n",
    "            clip_on=False,\n",
    "            zorder=100,\n",
    "            # mec='black',\n",
    "            c=c,\n",
    "            markersize=3.5,\n",
    "        )\n",
    "\n",
    "ax = axs['stab2']\n",
    "ax.plot(welf1, color=c1, linestyle='--', label=label1)\n",
    "ax.plot(welf2, color=c2, label=label2)\n",
    "ax.set(\n",
    "    xlabel='Time ($t$)',\n",
    ")\n",
    "ax.set_ylabel(\n",
    "    ylabel='Welfare',\n",
    "    labelpad=8,\n",
    ")\n",
    "ax.yaxis.set_major_formatter(\n",
    "    matplotlib.ticker.PercentFormatter(xmax=1,decimals=1),\n",
    ")\n",
    "\n",
    "evoml.analysis.save_and_download_fig(fig, './output/empirical.pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a58f56ee-afb8-407c-b05c-b3cc26395e4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameters saved as: ./output/empirical.json\n"
     ]
    }
   ],
   "source": [
    "param_tracker.save('./output/empirical.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "evoml",
   "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
}
