{
 "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",
    "\n",
    "import scipy.interpolate\n",
    "import sklearn.gaussian_process\n",
    "import sklearn.gaussian_process.kernels\n",
    "\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51f4f4a2",
   "metadata": {},
   "source": [
    "## Load CIFAR data"
   ]
  },
  {
   "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",
    "cifar_augmentations_game_df.groupby(level=0).mean().plot.line(figsize=(4,3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6106f6e9",
   "metadata": {},
   "source": [
    "## Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bf02f2fa-226e-4224-96b8-a1cb46012bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class InterpolatedFitness:\n",
    "    def __init__(self, *, game_df, interpolation_kind='linear'):\n",
    "        self.game_df = game_df\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",
    "        agg = 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 only implemented in two-group settings'\n",
    "        pb = p[-1]\n",
    "        return np.array([self.interp_a(pb), self.interp_b(pb)])\n",
    "\n",
    "\n",
    "class GaussianProcessFitness:\n",
    "    def __init__(self, *, game_df, random_state, gp_length_scale=0.1, gp_alpha=4):\n",
    "        cols = ['acc_A','acc_B']\n",
    "        self.game_df = game_df\n",
    "        self.rng = np.random.default_rng(random_state)\n",
    "        df = game_df[cols]\n",
    "        self.P = df.index.to_numpy().reshape(-1, 1)\n",
    "        self.y = df.to_numpy()\n",
    "        self.gp = sklearn.gaussian_process.GaussianProcessRegressor(\n",
    "            kernel=sklearn.gaussian_process.kernels.RBF(gp_length_scale),\n",
    "            normalize_y=True,\n",
    "            alpha=gp_alpha,\n",
    "            random_state=self.rng.integers(1e6),\n",
    "        )\n",
    "        self.gp.fit(self.P, self.y)\n",
    "\n",
    "    def __call__(self, p):\n",
    "        assert len(p)==2, 'Interpolation is only implemented in two-group settings'\n",
    "        pb = p[-1]\n",
    "        return self.gp.sample_y([[pb]],random_state=self.rng.integers(1e6))[0,:,0]\n",
    "\n",
    "\n",
    "class GaussianProcessMeanFitness:\n",
    "    def __init__(self, *, gp_fitness):\n",
    "        self.fitness=gp_fitness\n",
    "\n",
    "    def __call__(self,p):\n",
    "        assert len(p)==2, 'Interpolation is only implemented in two-group settings'\n",
    "        pb = p[-1]\n",
    "        return self.fitness.gp.predict([[pb]])[0]\n",
    "\n",
    "class GaussianProcessAmplifiedVariance:\n",
    "    def __init__(self, *, gp_fitness, default_variance_multiplier=1):\n",
    "        self.fitness=gp_fitness\n",
    "        self.alpha = default_variance_multiplier\n",
    "\n",
    "    def __call__(self,p,amplification=None):\n",
    "        assert len(p)==2, 'Interpolation is only implemented in two-group settings'\n",
    "        pb = p[-1]\n",
    "        mean = self.fitness.gp.predict([[pb]])[0]\n",
    "        y = self.fitness(p)\n",
    "        noise = y-mean\n",
    "        alpha = amplification or self.alpha\n",
    "        return mean+noise*alpha\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",
    "class ScaledFitness:\n",
    "    def __init__(self, fitness, tau):\n",
    "        self.fitness = fitness\n",
    "        self.tau = tau\n",
    "\n",
    "    def __call__(self, p):\n",
    "        return self.fitness(p)*self.tau\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d61b562a",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fb37d3c3-b27f-48ea-9bae-63391d436771",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_linspace = np.linspace(0,1,100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c6c90c7",
   "metadata": {},
   "source": [
    "## Fit Gaussian process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "63753033-89a6-4519-a219-4729c430d465",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gaussian_process_length_scale=0.1\n",
      "gaussian_process_alpha=4\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x305bcffa0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 166736 characters]",
      "text/plain": [
       "<Figure size 800x200 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 211,
       "width": 811
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "interp_fitness = InterpolatedFitness(\n",
    "    game_df=cifar_augmentations_game_df,\n",
    ")\n",
    "gp_fitness = GaussianProcessFitness(\n",
    "    game_df=cifar_augmentations_game_df,\n",
    "    gp_length_scale=param(0.1,'gaussian_process_length_scale'),\n",
    "    gp_alpha=param(4,'gaussian_process_alpha'),\n",
    "    random_state=1234,\n",
    ")\n",
    "gp_mean_fitness = GaussianProcessMeanFitness(\n",
    "    gp_fitness=gp_fitness,\n",
    ")\n",
    "gp_amplified_fitness = GaussianProcessAmplifiedVariance(\n",
    "    gp_fitness=gp_fitness,\n",
    ")\n",
    "\n",
    "fig,axs=plt.subplots(1,3,figsize=(8,2),layout='constrained')\n",
    "axs_it = iter(axs.ravel())\n",
    "\n",
    "ax = next(axs_it)\n",
    "for i in [0,1]:\n",
    "    ax.plot(\n",
    "        p_linspace,\n",
    "        [interp_fitness([1-p,p])[i] for p in p_linspace],\n",
    "        color=evoml.group_colors[i],\n",
    "        label=['A','B'][i],\n",
    "    )\n",
    "\n",
    "ax.set(\n",
    "    title='CIFAR-10 Interpolated fitness',\n",
    ")\n",
    "ax.legend()\n",
    "\n",
    "ax = next(axs_it)\n",
    "for i in [0,1]:\n",
    "    (\n",
    "        gp_fitness.game_df\n",
    "        .iloc[:,i]\n",
    "        .groupby(level=0)\n",
    "        .agg(['mean','std'])\n",
    "        .plot.line(\n",
    "            y='mean',\n",
    "            yerr='std',\n",
    "            label=['A','B'][i],\n",
    "            color=evoml.group_colors[i],\n",
    "            ax=ax,\n",
    "            capsize=3,\n",
    "        )\n",
    "    )\n",
    "    mean_prediction, std_prediction = gp_fitness.gp.predict(\n",
    "        p_linspace.reshape(-1,1),\n",
    "        return_std=True,\n",
    "    )\n",
    "    ax.fill_between(\n",
    "        p_linspace,\n",
    "        mean_prediction[:,i] - std_prediction[:,i],\n",
    "        mean_prediction[:,i] + std_prediction[:,i],\n",
    "        color=evoml.group_colors[i],\n",
    "        alpha=0.3,\n",
    "    )\n",
    "    ax.plot(\n",
    "        p_linspace,\n",
    "        [gp_fitness([1-p,p])[i] for p in p_linspace],\n",
    "        '.',\n",
    "        color=evoml.group_colors[i],\n",
    "        label=['A','B'][i],\n",
    "        alpha=0.3,\n",
    "    )\n",
    "\n",
    "ax.set(\n",
    "    title='CIFAR-10 Gaussian process fit',\n",
    ")\n",
    "\n",
    "ax = next(axs_it)\n",
    "for i in [0,1]:\n",
    "    ax.plot(\n",
    "        p_linspace,\n",
    "        [gp_mean_fitness([1-p,p])[i] for p in p_linspace],\n",
    "        color=evoml.group_colors[i],\n",
    "        label=['A','B'][i],\n",
    "    )\n",
    "ax.plot(\n",
    "    p_linspace,\n",
    "    [np.array([1-p,p])@gp_mean_fitness([1-p,p]) for p in p_linspace],\n",
    ")\n",
    "\n",
    "ax.set(\n",
    "    title='CIFAR-10 GP mean fitness',\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0077d06f",
   "metadata": {},
   "source": [
    "### Compute equilibrium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6c182bd6-64df-4782-b12b-8b9d021d7900",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.5007507507507508)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p0_eq_vec = np.linspace(0.45,0.55,1000)\n",
    "p_eq = p0_eq_vec[np.array([gp_mean_fitness([1-p0,p0]) for p0 in p0_eq_vec]).min(axis=1).argmax()]\n",
    "p_eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5aca8173",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "initial_state_range_min=0.35\n",
      "initial_state_range_max=0.65\n"
     ]
    }
   ],
   "source": [
    "p0_vec = np.linspace(\n",
    "    param(0.35,'initial_state_range_min'),\n",
    "    param(0.65,'initial_state_range_max'),\n",
    "    100,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0750e405",
   "metadata": {},
   "source": [
    "## Sensitivity to sampling noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c04c7c0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.35, 0.38, 0.41, 0.44, 0.47, 0.5 , 0.53, 0.56, 0.59, 0.62, 0.65])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reduced_p0_vec = np.linspace(p0_vec.min(), p0_vec.max(), 11)\n",
    "reduced_p0_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24030428-b5a9-42c3-9166-019002588f1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "noisy_simulation_termination_threshold=0.1\n",
      "noisy_simulation_num_repetitions=50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [11:23<00:00, 13.66s/it]\n"
     ]
    },
    {
     "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>rep</th>\n",
       "      <th>p0</th>\n",
       "      <th>amplification</th>\n",
       "      <th>tf</th>\n",
       "      <th>dominant_group</th>\n",
       "      <th>res</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0.35</td>\n",
       "      <td>1</td>\n",
       "      <td>1272</td>\n",
       "      <td>0</td>\n",
       "      <td>[[0.65, 0.35], [0.6500822671777533, 0.34991773...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>0.35</td>\n",
       "      <td>5</td>\n",
       "      <td>1270</td>\n",
       "      <td>0</td>\n",
       "      <td>[[0.65, 0.35], [0.6490562930296309, 0.35094370...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>0.38</td>\n",
       "      <td>1</td>\n",
       "      <td>1441</td>\n",
       "      <td>0</td>\n",
       "      <td>[[0.62, 0.38], [0.6199441961438193, 0.38005580...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>0.38</td>\n",
       "      <td>5</td>\n",
       "      <td>1399</td>\n",
       "      <td>0</td>\n",
       "      <td>[[0.62, 0.38], [0.619248015469921, 0.380751984...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>0.41</td>\n",
       "      <td>1</td>\n",
       "      <td>1898</td>\n",
       "      <td>0</td>\n",
       "      <td>[[0.5900000000000001, 0.41], [0.59029592131583...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   rep    p0  amplification    tf  dominant_group  \\\n",
       "0    0  0.35              1  1272               0   \n",
       "1    0  0.35              5  1270               0   \n",
       "2    0  0.38              1  1441               0   \n",
       "3    0  0.38              5  1399               0   \n",
       "4    0  0.41              1  1898               0   \n",
       "\n",
       "                                                 res  \n",
       "0  [[0.65, 0.35], [0.6500822671777533, 0.34991773...  \n",
       "1  [[0.65, 0.35], [0.6490562930296309, 0.35094370...  \n",
       "2  [[0.62, 0.38], [0.6199441961438193, 0.38005580...  \n",
       "3  [[0.62, 0.38], [0.619248015469921, 0.380751984...  \n",
       "4  [[0.5900000000000001, 0.41], [0.59029592131583...  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_noise_results = []\n",
    "amplification_levels = [1,5]\n",
    "noisy_eps = param(0.1, 'noisy_simulation_termination_threshold')\n",
    "for rep in tqdm(range(param(50,'noisy_simulation_num_repetitions'))):\n",
    "    for p0 in reduced_p0_vec:\n",
    "        for amp in amplification_levels:\n",
    "            res = evoml.dynamics.discrete_replicator(\n",
    "                p0=[1-p0,p0],\n",
    "                fitness=GaussianProcessAmplifiedVariance(\n",
    "                    gp_fitness=gp_fitness,\n",
    "                    default_variance_multiplier=amp,\n",
    "                ),\n",
    "                T=100000,\n",
    "                return_fitness=False,\n",
    "                termination_predicate=lambda p: p.min()<noisy_eps,\n",
    "            )\n",
    "            conv_noise_results.append({\n",
    "                'rep': rep,\n",
    "                'p0': p0,\n",
    "                'amplification': amp,\n",
    "                'tf': len(res),\n",
    "                'dominant_group': res[-1].argmax(),\n",
    "                'res': res,\n",
    "            })\n",
    "\n",
    "conv_noise_results_df = pd.DataFrame(conv_noise_results)\n",
    "conv_noise_results_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "940134d9-f2fc-4f46-904c-8c8496682286",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x306d73790>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 37584 characters]",
      "text/plain": [
       "<Figure size 300x200 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 236,
       "width": 306
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,ax=plt.subplots()\n",
    "(\n",
    "    conv_noise_results_df\n",
    "    .assign(\n",
    "        correct=lambda df: df.apply(\n",
    "            lambda row: (row['p0']>=p_eq)==row['dominant_group'],\n",
    "            axis=1,\n",
    "        ),\n",
    "        a_wins=lambda df: 1-df['dominant_group'],\n",
    "        b_wins=lambda df: df['dominant_group'],\n",
    "    )\n",
    "    .groupby(['p0','amplification'])\n",
    "    ['b_wins']\n",
    "    .mean()\n",
    "    .unstack()\n",
    "    .pipe(lambda df: df.rename({c: f'$\\\\alpha={c:g}$' for c in df.columns},axis=1))\n",
    "    .plot.line(\n",
    "        # marker='.',\n",
    "        ax=ax,\n",
    "        figsize=(3,2),\n",
    "    )\n",
    ")\n",
    "ax.axhline(1,**evoml.background_line_style)\n",
    "ax.axhline(0.5,**evoml.background_line_style)\n",
    "ax.axvline(\n",
    "    p_eq,\n",
    "    linestyle='--',\n",
    "    # label='$\\\\vp^*$',\n",
    ")\n",
    "ax.set(\n",
    "    title='Outcomes of stochastic dynamics',\n",
    "    xlabel='$p^0_B$',\n",
    "    ylabel='$\\\\Pr(\\\\vp^*=\\\\text{Only $B$})$',\n",
    "    \n",
    ")\n",
    "ax.legend(loc=(1.1,0.2),title='Variance amplification')\n",
    "ax.legend(loc=4)\n",
    "# ax.axvline(0.56)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9235dca5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "noisy_simulation_trajectories_initial_state=0.47\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [00:23<00:00,  7.80s/it]\n"
     ]
    }
   ],
   "source": [
    "gp_fitness.rng = np.random.default_rng(1234)\n",
    "\n",
    "traj_p0 = param(0.47,'noisy_simulation_trajectories_initial_state')\n",
    "amplification_trajectories = {}\n",
    "for amp in tqdm([0]+amplification_levels):\n",
    "    amplification_trajectories[amp] = []\n",
    "    for rep in range(6):\n",
    "        res = evoml.dynamics.discrete_replicator(\n",
    "            p0=[1-traj_p0,traj_p0],\n",
    "            fitness=GaussianProcessAmplifiedVariance(\n",
    "                gp_fitness=gp_fitness,\n",
    "                default_variance_multiplier=amp,\n",
    "            ),\n",
    "            T=5000,\n",
    "            return_fitness=False,\n",
    "            # termination_predicate=lambda p: p.min()<0.0001,\n",
    "        )\n",
    "        amplification_trajectories[amp].append(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ea57a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Figure saved as ./output/empirical_sensitivity_stochastic.pdf\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"data:text/plain;base64,[encoded data, 468640 characters]\" download=\"empirical_sensitivity_stochastic.pdf\">Download empirical_sensitivity_stochastic.pdf</a>"
      ],
      "text/plain": [
       "<evoml.analysis.DownloadableIO at 0x3142f2c20>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 215376 characters]",
      "text/plain": [
       "<Figure size 950x250 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 261,
       "width": 942
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,axs=plt.subplots(\n",
    "    1,3,\n",
    "    figsize=(9.5,2.5),\n",
    "    constrained_layout={'w_pad':0.15},\n",
    ")\n",
    "axs_it = iter(axs.ravel())\n",
    "\n",
    "\n",
    "ax = next(axs_it)\n",
    "for i in [0,1]:\n",
    "    (\n",
    "        gp_fitness.game_df\n",
    "        .iloc[:,i]\n",
    "        .groupby(level=0)\n",
    "        .agg(['mean','std'])\n",
    "        .plot.line(\n",
    "            y='mean',\n",
    "            yerr='std',\n",
    "            label=['A','B'][i],\n",
    "            color=evoml.group_colors[i],\n",
    "            ax=ax,\n",
    "            capsize=3,\n",
    "        )\n",
    "    )\n",
    "    mean_prediction, std_prediction = gp_fitness.gp.predict(\n",
    "        p_linspace.reshape(-1,1),\n",
    "        return_std=True,\n",
    "    )\n",
    "    ax.fill_between(\n",
    "        p_linspace,\n",
    "        mean_prediction[:,i] - std_prediction[:,i],\n",
    "        mean_prediction[:,i] + std_prediction[:,i],\n",
    "        color=evoml.group_colors[i],\n",
    "        alpha=0.4,\n",
    "        label=['A','B'][i],\n",
    "    )\n",
    "ax.axvline(0.5,**evoml.background_line_style)\n",
    "ax.set(\n",
    "    title='CIFAR-10 Gaussian process fit',\n",
    "    xlabel=evoml.analysis.composition_label,\n",
    "    xticks=evoml.analysis.composition_tickpos,\n",
    "    xticklabels=evoml.analysis.composition_ticklabels[::-1],\n",
    "    ylabel='Accuracy',\n",
    ")\n",
    "ax.yaxis.set_major_formatter(matplotlib.ticker.PercentFormatter(xmax=1,decimals=1))\n",
    "\n",
    "c = 'tab:gray'\n",
    "ax.legend(handles=[\n",
    "    matplotlib.container.ErrorbarContainer(\n",
    "        (\n",
    "            matplotlib.lines.Line2D([],[], color=c),\n",
    "            [matplotlib.lines.Line2D([],[], color=c)],\n",
    "            [matplotlib.collections.LineCollection(np.empty((2,2,2)), color=c)],\n",
    "        ),\n",
    "        has_yerr=True,\n",
    "        label='Observations',\n",
    "    ),\n",
    "    matplotlib.patches.Patch(\n",
    "        color=c,\n",
    "        label='Gaussian process',\n",
    "        alpha=0.3,\n",
    "    ),\n",
    "    matplotlib.lines.Line2D([],[], color=c, ls=':', alpha=0.5, label='$\\\\vp=(0.5,0.5)$'),\n",
    "])\n",
    "\n",
    "ax = next(axs_it)\n",
    "amplification_colors = ['tab:red','tab:blue','tab:orange']\n",
    "for amp,c in zip(amplification_trajectories,amplification_colors):\n",
    "    for traj in amplification_trajectories[amp]: \n",
    "        ax.plot(traj[:,0], c=c, alpha=0.4,ls='-' if amp>0 else '--')\n",
    "\n",
    "ax.axhline(0.5,**evoml.background_line_style)\n",
    "ax.legend(\n",
    "    handles=[\n",
    "        matplotlib.lines.Line2D(\n",
    "            [],[],\n",
    "            color=c,\n",
    "            alpha=0.7,\n",
    "            label=f'$\\\\eta={amp}$',\n",
    "            ls='-' if amp>0 else '--'\n",
    "        )\n",
    "        for amp, c in zip(amplification_trajectories, amplification_colors)\n",
    "    ],\n",
    "    title='Noise level',\n",
    "    loc='lower right',\n",
    ")\n",
    "    \n",
    "ax.set(\n",
    "    title='Stochastic replicator trajectories',\n",
    "    ylabel=evoml.analysis.composition_label,\n",
    "    xlabel='Time ($t$)',\n",
    "    yticks=evoml.analysis.composition_tickpos,\n",
    "    yticklabels=evoml.analysis.composition_ticklabels,\n",
    "\n",
    ")\n",
    "\n",
    "#\n",
    "# Outcome of stochastic dynamics\n",
    "#\n",
    "ax = next(axs_it)\n",
    "ax.plot(p0_vec, p0_vec>p_eq, c='tab:red',ls='--', label='$\\\\eta=0$', alpha=0.4, zorder=-1)\n",
    "(\n",
    "    conv_noise_results_df\n",
    "    .assign(\n",
    "        correct=lambda df: df.apply(\n",
    "            lambda row: (row['p0']>=p_eq)==row['dominant_group'],\n",
    "            axis=1,\n",
    "        ),\n",
    "        a_wins=lambda df: 1-df['dominant_group'],\n",
    "        b_wins=lambda df: df['dominant_group'],\n",
    "    )\n",
    "    .groupby(['p0','amplification'])\n",
    "    ['b_wins']\n",
    "    .mean()\n",
    "    .unstack()\n",
    "    .pipe(lambda df: df.rename({c: f'$\\\\eta={c:g}$' for c in df.columns},axis=1))\n",
    "    .plot.line(\n",
    "        marker='.',\n",
    "        ax=ax,\n",
    "    )\n",
    ")\n",
    "ax.axhline(0.5,**evoml.background_line_style)\n",
    "ax.set(\n",
    "    title='Outcomes of stochastic dynamics',\n",
    "    xlabel='Initial proportion of group $B$ ($p^0_B$)',\n",
    "    ylabel='$\\\\Pr(\\\\vp^*=\\\\text{Only $B$})$',\n",
    "    \n",
    ")\n",
    "ax.legend(\n",
    "    loc='lower right',\n",
    ")\n",
    "\n",
    "# Confidence intervals\n",
    "for amp in amplification_levels:\n",
    "    (\n",
    "        conv_noise_results_df\n",
    "        .assign(\n",
    "            correct=lambda df: df.apply(\n",
    "                lambda row: (row['p0']>=p_eq)==row['dominant_group'],\n",
    "                axis=1,\n",
    "            ),\n",
    "            a_wins=lambda df: 1-df['dominant_group'],\n",
    "            b_wins=lambda df: df['dominant_group'],\n",
    "        )\n",
    "        .groupby(['p0','amplification'])\n",
    "        ['b_wins']\n",
    "        .agg(['sum','count'])\n",
    "        .stack()\n",
    "        .swaplevel(1,2)\n",
    "        .unstack()\n",
    "        .unstack()\n",
    "        [amp]\n",
    "        .agg(lambda row: scipy.stats.binomtest(row['sum'],row['count']).proportion_ci(confidence_level=0.99), axis=1)\n",
    "        .apply(lambda x: pd.Series(index=['l','h'],data=[x.low, x.high]))\n",
    "        .pipe(lambda df: ax.fill_between(\n",
    "            x=df.index,\n",
    "            y1=df['l'],\n",
    "            y2=df['h'],\n",
    "            alpha=0.15,\n",
    "        ))\n",
    "    )\n",
    "\n",
    "evoml.analysis.save_and_download_fig(fig, './output/empirical_sensitivity_stochastic.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e0b027a",
   "metadata": {},
   "source": [
    "## Time to convergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "65eba86d",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_p0_vec = np.linspace(0.1,0.9,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "f1c69c74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time_to_dominance_threshold=0.01\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:10<00:00,  9.71it/s]\n"
     ]
    }
   ],
   "source": [
    "eps = param(0.01,'time_to_dominance_threshold')\n",
    "conv_results = []\n",
    "for p0 in tqdm(full_p0_vec):\n",
    "    T = 100000\n",
    "    tconv = None\n",
    "    while tconv is None:\n",
    "        res = evoml.dynamics.discrete_replicator(\n",
    "            p0=[1-p0,p0],\n",
    "            fitness=gp_mean_fitness,\n",
    "            T=T,\n",
    "            termination_predicate=lambda p: p.min()<eps,\n",
    "        )\n",
    "        tconv = len(res)\n",
    "    conv_results.append(tconv)\n",
    "conv_results=np.array(conv_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "5aefe582",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x3144b18e0>"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 57688 characters]",
      "text/plain": [
       "<Figure size 400x300 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 314,
       "width": 386
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,ax=plt.subplots(figsize=(4,3))\n",
    "d = 1e-4\n",
    "ax.plot(\n",
    "    full_p0_vec[full_p0_vec<p_eq-d],conv_results[full_p0_vec<p_eq-d],\n",
    "    color=evoml.group_colors[0],\n",
    "    label='$\\\\vp^t \\\\to A$',\n",
    ")\n",
    "ax.plot(\n",
    "    full_p0_vec[full_p0_vec>p_eq+d],conv_results[full_p0_vec>p_eq+d],\n",
    "    color=evoml.group_colors[1],\n",
    "    label='$\\\\vp^t \\\\to B$',\n",
    ")\n",
    "ax.axvline(\n",
    "    p_eq,\n",
    "    label='$\\\\vp^*$',\n",
    "    linestyle='--',\n",
    "    # label='Unstable equilibrium $\\\\vp^*$',\n",
    ")\n",
    "ax.set(\n",
    "    xlabel='Initial population state $p^0_B$',\n",
    "    ylabel='Time to dominance [$T/\\\\tau$]',\n",
    "    title=f'CIFAR-10 Replicator Convergence ($\\\\varepsilon={eps}$)',\n",
    "    # xticks=[p_eq,1],\n",
    "    # xticklabels=[\n",
    "    #     f'$\\\\vp^* \\\\approx ({1-p_eq:.2f},{p_eq:.2f})$',\n",
    "    #     'Only $B$',\n",
    "    # ],\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a27d2fd",
   "metadata": {},
   "source": [
    "## Sensitivity to estimation errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ffc1ffd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 11/11 [00:05<00:00,  1.97it/s]\n"
     ]
    },
    {
     "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>p_eq_est</th>\n",
       "      <th>p_B_T</th>\n",
       "      <th>acc_T</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.35</td>\n",
       "      <td>0.199274</td>\n",
       "      <td>0.934517</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.38</td>\n",
       "      <td>0.259044</td>\n",
       "      <td>0.934517</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.41</td>\n",
       "      <td>0.319053</td>\n",
       "      <td>0.934517</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.44</td>\n",
       "      <td>0.379094</td>\n",
       "      <td>0.934517</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.47</td>\n",
       "      <td>0.439122</td>\n",
       "      <td>0.934517</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   p_eq_est     p_B_T     acc_T\n",
       "0      0.35  0.199274  0.934517\n",
       "1      0.38  0.259044  0.934517\n",
       "2      0.41  0.319053  0.934517\n",
       "3      0.44  0.379094  0.934517\n",
       "4      0.47  0.439122  0.934517"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p0 = 0.2\n",
    "eq_results = []\n",
    "for p_eq_hat in tqdm(reduced_p0_vec):\n",
    "    res = evoml.dynamics.discrete_replicator(\n",
    "        p0=[1-p0,p0],\n",
    "        fitness=ReflectedFitness(\n",
    "            fitness=gp_mean_fitness,\n",
    "            p_eq=p_eq_hat,\n",
    "        ),\n",
    "        T=10000,\n",
    "        return_fitness=True,\n",
    "    )\n",
    "    welf = (res[0]*res[1]).sum(axis=1)\n",
    "    eq_results.append({\n",
    "        'p_eq_est': p_eq_hat,\n",
    "        'p_B_T': res[0][-1][-1],\n",
    "        'acc_T': welf[-1],\n",
    "    })\n",
    "eq_results_df = pd.DataFrame(eq_results)\n",
    "eq_results_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3ad620d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x3141e4fa0>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 46448 characters]",
      "text/plain": [
       "<Figure size 300x200 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 233,
       "width": 297
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,ax=plt.subplots(figsize=(3,2))\n",
    "(\n",
    "    eq_results_df\n",
    "    .plot.line(\n",
    "        x='p_eq_est',\n",
    "        y='p_B_T',\n",
    "        marker='.',\n",
    "        ax=ax,\n",
    "        label='$\\\\vp^*_\\\\mathrm{control}$',\n",
    "        color='tab:green',\n",
    "    )\n",
    ")\n",
    "ax.axvline(\n",
    "    p_eq,\n",
    "    linestyle='--',\n",
    "    label='$\\\\vp^*_\\\\mathrm{true}$',\n",
    "    # label='$\\\\vp^*$',\n",
    ")\n",
    "ax.axline(\n",
    "    xy1=(0.5,0.5),\n",
    "    slope=1,\n",
    "    color='tab:gray',\n",
    "    zorder=-1,\n",
    "    linestyle='--',\n",
    "    alpha=0.7,\n",
    "    label='$x=y$',\n",
    ")\n",
    "ax.set(\n",
    "    xlabel='Predicted equilibrium $\\\\hat{\\\\vp}^*$',\n",
    "    ylabel='Control equilibrium $\\\\vp^*$',\n",
    "    title='Control with $\\\\vp^*$ estimation errors',\n",
    "    # aspect='equal',\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "d2606fac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Figure saved as ./output/empirical_sensitivity_deterministic.pdf\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"data:text/plain;base64,[encoded data, 380852 characters]\" download=\"empirical_sensitivity_deterministic.pdf\">Download empirical_sensitivity_deterministic.pdf</a>"
      ],
      "text/plain": [
       "<evoml.analysis.DownloadableIO at 0x3199b3cc0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 159780 characters]",
      "text/plain": [
       "<Figure size 950x250 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 261,
       "width": 944
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,axs=plt.subplots(\n",
    "    1,3,\n",
    "    figsize=(9.5,2.5),\n",
    "    constrained_layout={'w_pad':0.15},\n",
    ")\n",
    "axs_it = iter(axs.ravel())\n",
    "\n",
    "ax = next(axs_it)\n",
    "for i in [0,1]:\n",
    "    ax.plot(\n",
    "        p_linspace,\n",
    "        [gp_mean_fitness([1-p,p])[i] for p in p_linspace],\n",
    "        color=evoml.group_colors[i],\n",
    "        label=f'$\\\\mathrm{{acc}}_{[\"A\",\"B\"][i]}$',\n",
    "    )\n",
    "ax.plot(\n",
    "    p_linspace,\n",
    "    [np.array([1-p,p])@gp_mean_fitness([1-p,p]) for p in p_linspace],\n",
    "    color='tab:cyan',\n",
    "    zorder=-1,\n",
    "    linestyle=':',\n",
    ")\n",
    "ax.legend()\n",
    "\n",
    "ax.set(\n",
    "    title='CIFAR-10 GP mean fitness',\n",
    "    xlabel=evoml.analysis.composition_label,\n",
    "    xticks=evoml.analysis.composition_tickpos,\n",
    "    xticklabels=evoml.analysis.composition_ticklabels[::-1],\n",
    "    ylabel='Accuracy',\n",
    ")\n",
    "ax.yaxis.set_major_formatter(matplotlib.ticker.PercentFormatter(xmax=1,decimals=1))\n",
    "\n",
    "ax = next(axs_it)\n",
    "d = 1e-4\n",
    "ax.plot(\n",
    "    full_p0_vec[full_p0_vec<p_eq-d],conv_results[full_p0_vec<p_eq-d],\n",
    "    color=evoml.group_colors[0],\n",
    "    label='$\\\\vp^t \\\\to A$',\n",
    ")\n",
    "ax.plot(\n",
    "    full_p0_vec[full_p0_vec>p_eq+d],conv_results[full_p0_vec>p_eq+d],\n",
    "    color=evoml.group_colors[1],\n",
    "    label='$\\\\vp^t \\\\to B$',\n",
    ")\n",
    "ax.axvline(\n",
    "    p_eq,\n",
    "    label='$\\\\vp^*$',\n",
    "    linestyle='--',\n",
    ")\n",
    "ax.set(\n",
    "    xlabel='Initial proportion of group $B$ ($p^0_B$)',\n",
    "    ylabel='Steps to convergence',\n",
    "    title=f'Time to replicator convergence',\n",
    ")\n",
    "ax.legend()\n",
    "\n",
    "ax = next(axs_it)\n",
    "(\n",
    "    eq_results_df\n",
    "    .plot.line(\n",
    "        x='p_eq_est',\n",
    "        y='p_B_T',\n",
    "        marker='.',\n",
    "        ax=ax,\n",
    "        label='$\\\\vp^*_\\\\mathrm{control}$',\n",
    "        color='tab:green',\n",
    "    )\n",
    ")\n",
    "# ax.axvline(\n",
    "#     p_eq,\n",
    "#     linestyle='--',\n",
    "#     label='$\\\\vp^*_\\\\mathrm{true}$',\n",
    "# )\n",
    "ax.plot(\n",
    "    p_eq,\n",
    "    p_eq,\n",
    "    'o',\n",
    "    mfc='white',\n",
    "    label='$\\\\vp^*_\\\\mathrm{true}$',\n",
    ")\n",
    "ax.axline(\n",
    "    xy1=(0.5,0.5),\n",
    "    slope=1,\n",
    "    color='tab:gray',\n",
    "    zorder=-1,\n",
    "    linestyle='--',\n",
    "    alpha=0.7,\n",
    "    label='$x=y$',\n",
    ")\n",
    "ax.set(\n",
    "    xlabel='Predicted equilibrium $\\\\hat{\\\\vp}^*$',\n",
    "    ylabel='Outcome of dynamics $\\\\vp^*$',\n",
    "    title='Stabilization with estimation errors',\n",
    ")\n",
    "ax.legend(loc=4)\n",
    "\n",
    "evoml.analysis.save_and_download_fig(fig, './output/empirical_sensitivity_deterministic.pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "26a88ab2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameters saved as: ./output/sensitivity.json\n"
     ]
    }
   ],
   "source": [
    "param_tracker.save('./output/sensitivity.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
}
