{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d988592c-f979-4185-bfc3-8abe6781a10f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from notebook_env import *\n",
    "from evoml.analysis import simplex_projection as PROJ\n",
    "import evoml.dynamics\n",
    "\n",
    "from mpltern.datasets import get_triangular_grid\n",
    "\n",
    "import folktables\n",
    "from folktables import ACSDataSource, ACSIncome\n",
    "\n",
    "import xgboost\n",
    "\n",
    "import sklearn.linear_model\n",
    "import sklearn.svm\n",
    "import sklearn.model_selection\n",
    "\n",
    "from scipy.interpolate import LinearNDInterpolator\n",
    "\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "820c52e8-1094-4c1c-a360-64c24c2a32d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "P_singletons = np.array(get_triangular_grid(2))\n",
    "rng = np.random.default_rng(12345)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "362b930b-9eec-43c9-bbfd-eb4cf11127bd",
   "metadata": {},
   "source": [
    "## Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "cd62ec93-8bfc-430e-bdc6-de57b427aa17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state_texas_num_datapoints=135,924\n",
      "state_new_york_num_datapoints=103,021\n",
      "state_california_num_datapoints=195,665\n"
     ]
    }
   ],
   "source": [
    "# Fetch folktables data\n",
    "data_source = ACSDataSource(\n",
    "    survey_year='2018', horizon='1-Year', survey='person',\n",
    "    root_dir=os.path.expanduser('~/Documents/data/folktables'),\n",
    ")\n",
    "selected_states = ['TX','NY','CA']\n",
    "state_names = {\n",
    "    'CA':'California',\n",
    "    'TX':'Texas',\n",
    "    'NY':'New-York',\n",
    "}\n",
    "data = {\n",
    "    s: data_source.get_data(states=[s], download=True)\n",
    "    for s in selected_states\n",
    "}\n",
    "\n",
    "prob = folktables.ACSIncome\n",
    "\n",
    "for s in selected_states:\n",
    "    param(\n",
    "        prob.df_to_numpy(data[s])[0].shape[0],\n",
    "        f'state_{state_names[s].replace(\"-\",\"_\").lower()}_num_datapoints',\n",
    "        '{:,}'.format,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "119eb8b3-ee42-47af-b24f-b258fa5edc82",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test_set_size=5,000\n",
      "Dataset sizes:\n",
      "TX: n_train_pool=130924  n_test=5000\n",
      "NY: n_train_pool=98021   n_test=5000\n",
      "CA: n_train_pool=190665  n_test=5000\n"
     ]
    }
   ],
   "source": [
    "n_test = param(5000,'test_set_size', '{:,}'.format)\n",
    "Xy_train_pool = {}\n",
    "Xy_test = {}\n",
    "for s in selected_states:\n",
    "    X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(\n",
    "        *prob.df_to_numpy(data[s])[:2],\n",
    "        random_state=rng.integers(1e9),\n",
    "        test_size=n_test,\n",
    "    )\n",
    "    Xy_train_pool[s] = (X_train, y_train)\n",
    "    Xy_test[s] = (X_test, y_test)\n",
    "\n",
    "print('Dataset sizes:')\n",
    "for s in selected_states:\n",
    "    print(f'{s}: n_train_pool={len(Xy_train_pool[s][0]):<7} n_test={len(Xy_test[s][0])}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "34167687",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TX->TX: 0.7622\n",
      "TX->NY: 0.569\n",
      "TX->CA: 0.5448\n",
      "NY->TX: 0.637\n",
      "NY->NY: 0.753\n",
      "NY->CA: 0.5552\n",
      "CA->TX: 0.637\n",
      "CA->NY: 0.5852\n",
      "CA->CA: 0.7598\n"
     ]
    }
   ],
   "source": [
    "for s_src in selected_states:\n",
    "    for s_dst in selected_states:\n",
    "        clf = sklearn.svm.LinearSVC()\n",
    "        clf.fit(Xy_train_pool[s_src][0][:1000],Xy_train_pool[s_dst][1][:1000])\n",
    "        print(f'{s_src}->{s_dst}: {clf.score(*Xy_test[s_dst])}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e1e01424",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5000, 10)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Xy_test[selected_states[0]][0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63f863a7-7bba-46f9-8604-5ee55be03d28",
   "metadata": {},
   "source": [
    "## Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b55d028e-0551-4051-9bbe-c41716b37b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample(n,p):\n",
    "    \" Sample n datapoints from the population mixture p \"\n",
    "    n_vec = rng.multinomial(n,p)\n",
    "    X = []\n",
    "    y = []\n",
    "    for i,s in enumerate(selected_states):\n",
    "        if n_vec[i]>0:\n",
    "            X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(\n",
    "                *Xy_train_pool[s],\n",
    "                train_size=n_vec[i],\n",
    "                random_state=rng.integers(1e9),\n",
    "            )\n",
    "            X.append(X_train)\n",
    "            y.append(y_train)\n",
    "    X = np.vstack(X)\n",
    "    y = np.hstack(y)\n",
    "    assert len(X)==len(y)\n",
    "    assert len(y)==n\n",
    "    perm = rng.permutation(n)\n",
    "    return X[perm],y[perm]\n",
    "\n",
    "def group_score(clf):\n",
    "    return np.array([clf.score(*Xy_test[s]) for s in selected_states])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "86ee7d5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TX->TX: 0.7642, 0.7616 (d=-0.00)\n",
      "TX->NY: 0.7416, 0.7294 (d=-0.01)\n",
      "TX->CA: 0.7542, 0.7516 (d=-0.00)\n",
      "NY->TX: 0.7692, 0.7516 (d=-0.02)\n",
      "NY->NY: 0.7452, 0.7346 (d=-0.01)\n",
      "NY->CA: 0.7604, 0.7804 (d=0.02)\n",
      "CA->TX: 0.7576, 0.7534 (d=-0.00)\n",
      "CA->NY: 0.746, 0.7412 (d=-0.00)\n",
      "CA->CA: 0.7636, 0.7584 (d=-0.01)\n"
     ]
    }
   ],
   "source": [
    "# Sample procedure sanity check\n",
    "\n",
    "for s_src in selected_states:\n",
    "    for s_dst in selected_states:\n",
    "        clf = sklearn.svm.LinearSVC()\n",
    "        inds = np.random.choice(len(Xy_train_pool[s_src][0]),size=1000)\n",
    "        clf.fit(Xy_train_pool[s_src][0][inds],Xy_train_pool[s_src][1][inds])\n",
    "        score1 = clf.score(*Xy_test[s_dst])\n",
    "        clf.fit(*sample(1000,p=(np.arange(3)==selected_states.index(s_src)).astype(float)))\n",
    "        score2 = clf.score(*Xy_test[s_dst])\n",
    "        print(f'{s_src}->{s_dst}: {score1}, {score2} (d={score2-score1:.2f})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "34bd3f50-4180-4fb8-ad6f-262d63178025",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "triangular_grid_size=20\n",
      "grid_num_repetitions=10\n",
      "training_set_size=1000\n",
      "svm\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 210/210 [00:53<00:00,  3.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xgb\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 210/210 [03:58<00:00,  1.14s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logreg\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 210/210 [05:34<00:00,  1.59s/it]\n"
     ]
    }
   ],
   "source": [
    "P = np.array(get_triangular_grid(param(20,'triangular_grid_size')))\n",
    "n_reps = param(10,'grid_num_repetitions')\n",
    "\n",
    "def compute_fp(*, P, clf_f, n_train, n_reps):\n",
    "    from tqdm import tqdm\n",
    "    f_vec = []\n",
    "    for p in tqdm(P.T):\n",
    "        f_p = np.zeros((n_reps,len(selected_states)))\n",
    "        for rep in range(n_reps):\n",
    "            Xy_train = sample(n=n_train,p=p)\n",
    "            clf = clf_f()\n",
    "            clf.fit(*Xy_train)\n",
    "            f = group_score(clf)\n",
    "            f_p[rep] = f\n",
    "        f_vec.append(f_p.mean(axis=0))\n",
    "    return np.vstack(f_vec)\n",
    "\n",
    "algs = {\n",
    "    'svm': sklearn.svm.LinearSVC,\n",
    "    'xgb': xgboost.XGBClassifier,\n",
    "    'logreg': functools.partial(sklearn.linear_model.LogisticRegression, max_iter=10000),\n",
    "}\n",
    "algs_f_vec = {}\n",
    "n_train = param(1000,'training_set_size')\n",
    "for alg in algs:\n",
    "    print(alg)\n",
    "    algs_f_vec[alg] = compute_fp(\n",
    "        P=P,\n",
    "        n_reps=n_reps,\n",
    "        n_train=n_train,\n",
    "        clf_f=algs[alg],\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b2fa0a0-5dbc-4072-9643-6edd2e435592",
   "metadata": {},
   "source": [
    "## Analyze results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03782cf0-6349-4e21-91f6-1ea6148f9de2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.quiver.Quiver at 0x30fea4340>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 260168 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 464,
       "width": 533
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Vector field and interpolation\n",
    "\n",
    "f_vec = algs_f_vec['svm']\n",
    "F = f_vec.T\n",
    "acc_P = (F*P).sum(axis=0)\n",
    "dP = P*(F-acc_P)\n",
    "\n",
    "ax=plt.subplot(projection='ternary')\n",
    "ax.quiver(\n",
    "    *P,\n",
    "    *dP,\n",
    "    **(dict(scale=0.12, clip_on=True)),\n",
    "    alpha=0.1,\n",
    ")\n",
    "ax_order = ['t','l','r']\n",
    "ax.set(**{f'{x}label': s for x,s in zip(ax_order,selected_states)})\n",
    "ax.plot(0,0,1,'o',c='tab:red')\n",
    "\n",
    "\n",
    "interp = LinearNDInterpolator(\n",
    "    points = (PROJ@P).T,\n",
    "    values = F.T,\n",
    ")\n",
    "P2 = np.array(get_triangular_grid(26))\n",
    "F2 = interp((PROJ@P2).T).T\n",
    "acc_P2 = (F2*P2).sum(axis=0)\n",
    "dP2 = P2*(F2-acc_P2)\n",
    "ax.quiver(\n",
    "    *P2,\n",
    "    *dP2,\n",
    "    **(dict(scale=0.09, clip_on=True)),\n",
    "    color='r',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "241e6ff3-c703-49ab-aad8-7fdf506cdf9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimplexInterpolatedFitness:\n",
    "    def __init__(self, *, P, F):\n",
    "        self.P = P\n",
    "        self.F = F\n",
    "        self.interp = LinearNDInterpolator(\n",
    "            points = (PROJ@P).T,\n",
    "            values = F.T,\n",
    "        )\n",
    "\n",
    "    def __call__(self, p):\n",
    "        assert len(p)==3, 'Interpolation is implemented in three-group settings'\n",
    "        return self.interp(PROJ@p)[0]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "dd8cdcfc-ce95-4dde-8326-79aed25c9ac2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 210/210 [00:04<00:00, 43.22it/s]\n"
     ]
    }
   ],
   "source": [
    "convergence = np.zeros_like(P.T)\n",
    "\n",
    "T=2000\n",
    "for i,p0 in enumerate(tqdm(P.T)):\n",
    "    res = evoml.dynamics.discrete_replicator(\n",
    "        p0=p0,\n",
    "        fitness=SimplexInterpolatedFitness(P=P,F=F),\n",
    "        T=T,\n",
    "    )\n",
    "    convergence[i] = res[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5fe11a6-4e70-41d9-a46d-d07975fa0d50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Text(0.5, nan, 'TX'), Text(0.5, nan, 'NY'), Text(0.5, nan, 'CA')]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 316520 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 464,
       "width": 533
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Basins of attraction\n",
    "\n",
    "ax=plt.subplot(projection='ternary')\n",
    "\n",
    "matplotlib.colors.to_rgba('blue')\n",
    "ax.tripcolor(\n",
    "    *P,\n",
    "    convergence.T[0],\n",
    "    shading='gouraud',\n",
    "    cmap=matplotlib.colors.LinearSegmentedColormap.from_list(\n",
    "        name='c1',\n",
    "        colors=[\n",
    "            matplotlib.colors.to_rgba('tab:blue'),\n",
    "            matplotlib.colors.to_rgba('tab:red'),\n",
    "        ]\n",
    "    ),\n",
    "    alpha=0.5,\n",
    "    edgecolor='white',\n",
    ")\n",
    "\n",
    "ax.quiver(\n",
    "    *P,\n",
    "    *dP,\n",
    "    **(dict(scale=0.12, clip_on=True)),\n",
    "    color='white',\n",
    ")\n",
    "ax_order = ['t','l','r']\n",
    "ax.set(**{f'{x}label': s for x,s in zip(ax_order,selected_states)})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "913671d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0., 1.],\n",
       "       [0., 0., 1.],\n",
       "       [0., 0., 1.],\n",
       "       ...,\n",
       "       [0., 0., 1.],\n",
       "       [0., 0., 1.],\n",
       "       [0., 0., 1.]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5c295e89",
   "metadata": {},
   "outputs": [],
   "source": [
    "whiten = lambda c,a=0.15: a*np.array(matplotlib.colors.to_rgb(c)) + (1-a)*np.ones(3)\n",
    "darken = lambda c,a=0.7: a*np.array(matplotlib.colors.to_rgb(c))\n",
    "normalize = lambda *args: np.array(args)/sum(args)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "8cd1ce3c-5d7d-4577-8cb7-3e29ca3a9925",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "simulation_length=600\n",
      "disparity_proportion_threshold=0.01\n",
      "svm_disparity_time=316\n",
      "Figure saved as ./output/folktables.pdf\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"data:text/plain;base64,[encoded data, 313576 characters]\" download=\"folktables.pdf\">Download folktables.pdf</a>"
      ],
      "text/plain": [
       "<evoml.analysis.DownloadableIO at 0x32c73d130>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 201580 characters]",
      "text/plain": [
       "<Figure size 950x270 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 281,
       "width": 937
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "state_colors = {\n",
    "    'CA': 'tab:cyan',\n",
    "    'TX': 'tab:orange',\n",
    "    'NY': 'tab:green',\n",
    "}\n",
    "\n",
    "def analyze(*, F, ternary_ax, composition_ax, fitness_ax, convergence, p0=None):\n",
    "    T=param(600,'simulation_length')\n",
    "    res, fit = evoml.dynamics.discrete_replicator(\n",
    "        p0=p0,\n",
    "        fitness=SimplexInterpolatedFitness(P=P,F=F),\n",
    "        T=T,\n",
    "        return_fitness=True,\n",
    "    )\n",
    "\n",
    "    disparity_population_threshold = param(0.01,'disparity_proportion_threshold')\n",
    "    group_disparity = (\n",
    "        (fit - (res<disparity_population_threshold).astype(float)).max(axis=1)\n",
    "        -\n",
    "        (fit + (res<disparity_population_threshold).astype(float)).min(axis=1)\n",
    "    )\n",
    "    disparity_accuracy_threshold = 0.01\n",
    "    t_disparity = (group_disparity<disparity_accuracy_threshold).argmax()\n",
    "    t0 = t_disparity\n",
    "\n",
    "    #\n",
    "    # Ternary\n",
    "    #\n",
    "    if ternary_ax is not None:\n",
    "        ax = ternary_ax\n",
    "        ax.tripcolor(\n",
    "            *P,\n",
    "            convergence.T[0],\n",
    "            shading='gouraud',\n",
    "            cmap=matplotlib.colors.LinearSegmentedColormap.from_list(\n",
    "                name='c1',\n",
    "                colors=[\n",
    "                    whiten(state_colors['CA']),\n",
    "                    whiten(state_colors['TX']),\n",
    "                ]\n",
    "            ),\n",
    "        )\n",
    "\n",
    "        ax_order = ['t','l','r']\n",
    "        ax.set(**{f'{x}label': s for x,s in zip(ax_order,selected_states)})\n",
    "        ax.plot(*res.T,clip_on=False,zorder=100,color='r')\n",
    "        ax.plot(*res[-1].T,'.',clip_on=False,color='r')\n",
    "\n",
    "        ax.set(\n",
    "            title='ACSIncome Replicator Dynamics',\n",
    "        )\n",
    "        ax.raxis.set_ticks([])\n",
    "        ax.taxis.set_ticks([])\n",
    "        ax.laxis.set_ticks([])\n",
    "\n",
    "        interp = LinearNDInterpolator(\n",
    "            points = (PROJ@P).T,\n",
    "            values = F.T,\n",
    "        )\n",
    "        P2 = np.array(get_triangular_grid(18))\n",
    "        F2 = interp((PROJ@P2).T).T\n",
    "        acc_P2 = (F2*P2).sum(axis=0)\n",
    "        dP2 = P2*(F2-acc_P2)\n",
    "        ax.quiver(\n",
    "            *P2,\n",
    "            *dP2,\n",
    "            **(dict(scale=0.11, clip_on=True)),\n",
    "        )\n",
    "\n",
    "    #\n",
    "    # Composition\n",
    "    #\n",
    "    if composition_ax is not None:\n",
    "        ax = composition_ax\n",
    "        for s in state_colors:\n",
    "            if s=='NY':\n",
    "                continue\n",
    "            ax.plot(\n",
    "                res[:,selected_states.index(s)],\n",
    "                label=f'$p_\\\\mathrm{{({s})}}$',\n",
    "                color=state_colors[s],\n",
    "            )\n",
    "        ax.plot(\n",
    "            np.arange(t_disparity),\n",
    "            res[:t_disparity,selected_states.index('NY')],\n",
    "            color=state_colors['NY'],\n",
    "        )\n",
    "        ax.plot(\n",
    "            np.arange(T-t_disparity)+t_disparity,\n",
    "            res[t_disparity:,selected_states.index('NY')],\n",
    "            color=state_colors['NY'],\n",
    "            linestyle='--'\n",
    "        )\n",
    "        ax.set(\n",
    "            title='Population Composition Over Time',\n",
    "            xlabel='Time',\n",
    "            ylabel='Relative group size',\n",
    "            ylim=(None,0.63),\n",
    "        )\n",
    "        ax.yaxis.set_major_formatter(matplotlib.ticker.PercentFormatter(xmax=1,decimals=0))\n",
    "        for s,pos in zip(state_colors, [\n",
    "            (420,0.39),\n",
    "            (420,0.52),\n",
    "            (150,0.12),\n",
    "        ]):\n",
    "            ax.annotate(\n",
    "                text=f'$p_\\\\mathrm{{({s})}}$',\n",
    "                xy=pos,\n",
    "                color=darken(state_colors[s]),\n",
    "            )\n",
    "        threshold_linestyle = dict(color='tab:red',alpha=0.3,zorder=-1,linestyle=':')\n",
    "        ax.axhline(disparity_population_threshold, **threshold_linestyle)\n",
    "        ax.axvline(t_disparity, **threshold_linestyle)\n",
    "        ax.annotate('$\\\\le1\\\\%$',(330,0.04),c='tab:red',alpha=0.7)\n",
    "\n",
    "    #\n",
    "    # Group fitness\n",
    "    #\n",
    "    if fitness_ax is not None:\n",
    "        ax = fitness_ax\n",
    "        Ft = SimplexInterpolatedFitness(P=P,F=F).interp((PROJ@res.T).T)\n",
    "        for s in state_colors:\n",
    "            if s=='NY':\n",
    "                continue\n",
    "            ax.plot(\n",
    "                Ft[:,selected_states.index(s)],\n",
    "                color=state_colors[s],\n",
    "            )\n",
    "        ax.plot(\n",
    "            np.arange(t_disparity),\n",
    "            Ft[:t_disparity,selected_states.index('NY')],\n",
    "            label='$p_\\\\mathrm{(NY)}$',\n",
    "            color=state_colors['NY'],\n",
    "        )\n",
    "        ax.plot(\n",
    "            np.arange(T-t_disparity)+t_disparity,\n",
    "            Ft[t_disparity:,selected_states.index('NY')],\n",
    "            label='$p_\\\\mathrm{(NY)}$',\n",
    "            color=state_colors['NY'],\n",
    "            linestyle='--'\n",
    "        )\n",
    "        ax.plot(\n",
    "            (Ft*res).sum(axis=1),\n",
    "            linestyle=':',\n",
    "        )\n",
    "        ax.set(\n",
    "            title='Group Fitness Over Time',\n",
    "            xlabel='Time',\n",
    "            ylabel='Classification accuracy',\n",
    "        )\n",
    "        ax.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(0.01))\n",
    "        ax.yaxis.set_major_formatter(matplotlib.ticker.PercentFormatter(xmax=1,decimals=0))\n",
    "        for s,pos in zip(state_colors, [\n",
    "            (390,0.757),\n",
    "            (390,0.7635),\n",
    "            (170,0.74),\n",
    "        ]):\n",
    "            ax.annotate(\n",
    "                text=f'$\\\\mathrm{{acc}}_\\\\mathrm{{({s})}}$',\n",
    "                xy=pos,\n",
    "                color=darken(state_colors[s]),\n",
    "            )\n",
    "        ax.annotate(\n",
    "            text=f'$\\\\mathrm{{acc}}_{{\\\\vp}}$',\n",
    "            xy=(60,0.757),\n",
    "            color='tab:blue',\n",
    "        )\n",
    "        ax.set_ylim((0.735,0.768))\n",
    "        ax.axvline((group_disparity<0.01).argmax(), **threshold_linestyle)\n",
    "        return {\n",
    "            'disparity_time': t_disparity,\n",
    "            'group_disparity': group_disparity,\n",
    "            'res': res,\n",
    "            'fit': fit,\n",
    "            'ternary_ax': ternary_ax,\n",
    "            'composition_ax': composition_ax,\n",
    "            'fitness_ax': fitness_ax,\n",
    "        }\n",
    "\n",
    "fig=plt.figure(figsize=(9.5,2.7),constrained_layout={'w_pad':0.15})\n",
    "tern_ax = fig.add_subplot(1,3,1,projection='ternary')\n",
    "analysis_res = analyze(\n",
    "    F=F,\n",
    "    ternary_ax=tern_ax,\n",
    "    composition_ax=fig.add_subplot(1,3,2),\n",
    "    fitness_ax=fig.add_subplot(1,3,3),\n",
    "    convergence=convergence,\n",
    "    p0=normalize(1,1,1),\n",
    ")\n",
    "param(analysis_res['disparity_time'], 'svm_disparity_time')\n",
    "tern_ax.set_position([-0.2, 0.17, 0.7, 0.675])\n",
    "tern_ax.annotate(\n",
    "    'Population composition',\n",
    "    (0.15,0.05),\n",
    "    ha='center',\n",
    "    xycoords='figure fraction',\n",
    ")\n",
    "\n",
    "evoml.analysis.save_and_download_fig(fig, './output/folktables.pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "d27e1e46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x32ccd4dc0>]"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 17344 characters]",
      "text/plain": [
       "<Figure size 300x200 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 196,
       "width": 293
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "(\n",
    "    analysis_res['group_disparity'][0],\n",
    "    analysis_res['group_disparity'][analysis_res['disparity_time']],\n",
    ")\n",
    "fig,ax=plt.subplots(figsize=(3,2))\n",
    "ax.plot(analysis_res['group_disparity'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "fe565673",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/210 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 210/210 [00:02<00:00, 84.63it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "simulation_length=600\n",
      "disparity_proportion_threshold=0.01\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 210/210 [00:02<00:00, 84.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "simulation_length=600\n",
      "disparity_proportion_threshold=0.01\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 210/210 [00:09<00:00, 23.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "simulation_length=600\n",
      "disparity_proportion_threshold=0.01\n",
      "Figure saved as ./output/folktables_alg_comparison.pdf\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<a href=\"data:text/plain;base64,[encoded data, 127156 characters]\" download=\"folktables_alg_comparison.pdf\">Download folktables_alg_comparison.pdf</a>"
      ],
      "text/plain": [
       "<evoml.analysis.DownloadableIO at 0x32ca7c130>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 271304 characters]",
      "text/plain": [
       "<Figure size 950x270 with 3 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 281,
       "width": 932
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig=plt.figure(figsize=(9.5,2.7),constrained_layout={'w_pad':0.15})\n",
    "\n",
    "alg_names = {\n",
    "    'svm': 'Linear SVM',\n",
    "    'logreg': 'Logistic Regression',\n",
    "    'xgb': 'XGBoost',\n",
    "}\n",
    "\n",
    "for i, alg in enumerate(alg_names):\n",
    "    ax=fig.add_subplot(1,3,i+1,projection='ternary')\n",
    "    F_alg = algs_f_vec[alg].T\n",
    "    alg_convergence = np.zeros_like(P.T)\n",
    "    for i,p0 in enumerate(tqdm(P.T)):\n",
    "        res = evoml.dynamics.discrete_replicator(\n",
    "            p0=p0,\n",
    "            fitness=SimplexInterpolatedFitness(P=P,F=F_alg),\n",
    "            T=1000 if alg!='xgb' else 4000,\n",
    "        )\n",
    "        alg_convergence[i] = res[-1]\n",
    "    analyze(\n",
    "        F=F_alg,\n",
    "        ternary_ax=ax,\n",
    "        composition_ax=None,\n",
    "        fitness_ax=None,\n",
    "        convergence=alg_convergence,\n",
    "        p0=normalize(3,4,1),\n",
    "    )\n",
    "    ax.set(\n",
    "        title=alg_names[alg],\n",
    "    )\n",
    "\n",
    "evoml.analysis.save_and_download_fig(fig, './output/folktables_alg_comparison.pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "f011b0b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.375, 0.5  , 0.125])"
      ]
     },
     "execution_count": 132,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normalize(3,4,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "6a6d12ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameters saved as: ./output/folktables.json\n"
     ]
    }
   ],
   "source": [
    "param_tracker.save('./output/folktables.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
}
