{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "63725ff7-a234-492d-9f69-ea4fd018915b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_env import *\n",
    "import scipy.interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cc74ed62-b346-4c2f-a5ef-21a97975bcf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "build_game = lambda f_A, f_B=None, p_vec=np.linspace(0,1,101): (\n",
    "    pd.DataFrame([\n",
    "        {\n",
    "            'p_A': 1-p_B,\n",
    "            'p_B': p_B,\n",
    "            'acc_A': f_A(p_B),\n",
    "            'acc_B': (f_B or (lambda p: f_A(1-p)))(p_B),\n",
    "        }\n",
    "        for p_B in p_vec\n",
    "    ])\n",
    "    .assign(acc_p=lambda df: df['acc_A']*df['p_A'] + df['acc_B']*df['p_B'])\n",
    "    .set_index('p_B')\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c3481fd6-2d54-4bec-8dd1-94dfadb276b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def annotate(ax, coords, alpha=1):\n",
    "    ax.annotate(\n",
    "        text='$\\\\mathrm{acc}_A$',\n",
    "        xy=coords[0],\n",
    "        color=evoml.analysis.group_cmap(1.0),\n",
    "        alpha=alpha,\n",
    "    )\n",
    "    ax.annotate(\n",
    "        text='$\\\\mathrm{acc}_B$',\n",
    "        xy=coords[1],\n",
    "        color=evoml.analysis.group_cmap(0.0),\n",
    "        alpha=alpha,\n",
    "    )\n",
    "    if len(coords)>3:\n",
    "        ax.annotate(\n",
    "            text='$\\\\mathrm{acc}_{\\\\bm{p}}$',\n",
    "            xy=coords[2],\n",
    "            color='tab:blue',\n",
    "            alpha=alpha,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e2b03467-5d31-4de3-ad1b-cb3a015780d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Figure saved as ./output/prediction_game_types.pdf\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"data:text/plain;base64,[encoded data, 265348 characters]\" download=\"prediction_game_types.pdf\">Download prediction_game_types.pdf</a>"
      ],
      "text/plain": [
       "<evoml.analysis.DownloadableIO at 0x162398e00>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 122080 characters]",
      "text/plain": [
       "<Figure size 950x190 with 4 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 201,
       "width": 949
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "fig, axs = evoml.create_figure(\n",
    "    1,4,\n",
    "    figsize=(9.5,1.9),\n",
    "    sharey=True,\n",
    "    constrained_layout={'w_pad':0.1},\n",
    ")\n",
    "axs_it = iter(axs.ravel())\n",
    "\n",
    "def plot_population_game(**kwargs):\n",
    "    kwargs['annotations'].pop('acc_p')\n",
    "    return evoml.plot_population_game(**kwargs)\n",
    "    \n",
    "ax = next(axs_it)\n",
    "plot_population_game(\n",
    "    game_df=build_game(lambda p: 0.95-p**1.8/7),\n",
    "    ax=ax,\n",
    "    equilibria=[\n",
    "        (0,True),\n",
    "        (0.5,False),\n",
    "        (1.0,True),\n",
    "    ],\n",
    "    equilibria_kwargs=dict(\n",
    "        d=0.1,\n",
    "        markersize=3,\n",
    "    ),\n",
    "    ax_kwargs={'ylim':(0.79,0.99)},\n",
    "    annotations={\n",
    "        'acc_A': (0.0,0.965),\n",
    "        'acc_B': (0.86,0.965),\n",
    "        'acc_p': (0.78,0.903),\n",
    "    }\n",
    ")\n",
    "ax.set(\n",
    "    title='Competitive Exclusion',\n",
    "    yticklabels=[],\n",
    ")\n",
    "\n",
    "ax = next(axs_it)\n",
    "plot_population_game(\n",
    "    # game_df=build_game(lambda p: 0.95-(1-p)**2.2/7),\n",
    "    game_df=build_game(\n",
    "        lambda p: scipy.interpolate.interp1d(\n",
    "            (0,0.5,1.0),\n",
    "            (0.84,0.96,0.80),\n",
    "            kind='quadratic',\n",
    "        )(1-p).item()\n",
    "    ),\n",
    "    ax=ax,\n",
    "    equilibria=[\n",
    "        (0.5,True),\n",
    "    ],\n",
    "    equilibria_kwargs=dict(\n",
    "        d=0.1,\n",
    "        markersize=3,\n",
    "    ),\n",
    "    annotations={\n",
    "        'acc_A': (0.81,0.95),\n",
    "        'acc_B': (0.05,0.95),\n",
    "        'acc_p': (0.455,0.932),\n",
    "    }\n",
    ")\n",
    "ax.set(\n",
    "    title='Stable Mutualistic Coexistence',\n",
    "    yticklabels=[],\n",
    ")\n",
    "\n",
    "ax = next(axs_it)\n",
    "plot_population_game(\n",
    "    game_df=build_game(\n",
    "        lambda p: scipy.interpolate.interp1d(\n",
    "            (0,0.5,1.0),\n",
    "            (0.84,0.96,0.80),\n",
    "            kind='quadratic',\n",
    "        )(p).item()\n",
    "    ),\n",
    "    ax=ax,\n",
    "    equilibria=[\n",
    "        (0,True),\n",
    "        (0.5,False),\n",
    "        (1.0,True),\n",
    "    ],\n",
    "    equilibria_kwargs=dict(\n",
    "        d=0.1,\n",
    "        markersize=3,\n",
    "    ),\n",
    "    annotations={\n",
    "        'acc_A': (0.05,0.95),\n",
    "        'acc_B': (0.81,0.95),\n",
    "        'acc_p': (0.455,0.932),\n",
    "    }\n",
    ")\n",
    "ax.set(\n",
    "    title='Unstable Mutualistic Coexistence',\n",
    "    yticklabels=[],\n",
    ")\n",
    "\n",
    "ax = next(axs_it)\n",
    "plot_population_game(\n",
    "    game_df=build_game(\n",
    "        lambda p: scipy.interpolate.interp1d(\n",
    "            (0,0.5,1.0),\n",
    "            (0.95,0.935,0.89),\n",
    "            kind='quadratic',\n",
    "        )(p).item(),\n",
    "        lambda p: scipy.interpolate.interp1d(\n",
    "            (0,0.5,1.0),\n",
    "            (0.815,0.85,0.86),\n",
    "            kind='quadratic',\n",
    "        )(p).item(),\n",
    "    ),\n",
    "    ax=ax,\n",
    "    equilibria=[\n",
    "        (0,True),\n",
    "        # (0.5,False),\n",
    "        # (1.0,True),\n",
    "    ],\n",
    "    equilibria_kwargs=dict(\n",
    "        d=0.1,\n",
    "        markersize=3,\n",
    "    ),\n",
    "    annotations={\n",
    "        'acc_A': (0.0,0.965),\n",
    "        'acc_B': (0.85,0.837),\n",
    "        'acc_p': (0.55,0.898),\n",
    "    },\n",
    ")\n",
    "ax.set(\n",
    "    title='No Coexistence Equilibria',\n",
    "    yticklabels=[],\n",
    ")\n",
    "\n",
    "fig.save_and_download('./output/prediction_game_types.pdf')"
   ]
  }
 ],
 "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
