{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "id": "EvuPYzXTcPgl"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from math import *\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "import warnings\n",
    "import matplotlib.colors as mcolors\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "from matplotlib import colors\n",
    "warnings.simplefilter(\"ignore\", UserWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure_dir = Path('../figures')\n",
    "figure_dir.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "empirics_dir = 'empirics'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [],
   "source": [
    "# different colors for attention values\n",
    "c_semantic='deeppink'\n",
    "c_positional='royalblue'\n",
    "c_att = 'rebeccapurple'\n",
    "\n",
    "# different colors for the lienar student\n",
    "c_lin = 'orange'\n",
    "\n",
    "# colors for the phase transition borders\n",
    "c_attlin = 'crimson' \n",
    "c_spinodal = 'forestgreen'\n",
    "\n",
    "# color for changing specific parameters\n",
    "c_no_col = 'black'\n",
    "\n",
    "\n",
    "cmap_uninf = LinearSegmentedColormap.from_list('INF-UNINF',\n",
    "                                                   [mcolors.to_rgba(c_semantic)[:3], (1, 1, 1), mcolors.to_rgba(c_positional)[:3]], N=100)\n",
    "cmap_attlin = LinearSegmentedColormap.from_list('INF-UNINF',\n",
    "                                                   [mcolors.to_rgba(c_lin)[:3], (1, 1, 1), mcolors.to_rgba(c_att)[:3]], N=100)\n",
    "cmap_pos = LinearSegmentedColormap.from_list('INF-UNINF',\n",
    "                                                   [mcolors.to_rgba('#8DE0A8')[:3],\n",
    "                                                    mcolors.to_rgba('#93FFE0')[:3],\n",
    "                                                    mcolors.to_rgba('#85EFFF')[:3],\n",
    "                                                    mcolors.to_rgba('#6BBFFF')[:3],\n",
    "                                                    mcolors.to_rgba(c_positional)[:3]], N=100)\n",
    "cmap_att = LinearSegmentedColormap.from_list('INF-UNINF',\n",
    "                                                   [mcolors.to_rgba('#FFE0B6')[:3],\n",
    "                                                    mcolors.to_rgba('#FFB48C')[:3],\n",
    "                                                    mcolors.to_rgba('#FF8166')[:3],\n",
    "                                                    mcolors.to_rgba('#FF3E3B')[:3],\n",
    "                                                    mcolors.to_rgba(c_semantic)[:3]], N=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(65.0, 105.0, 225.0)"
      ]
     },
     "execution_count": 215,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#6BBFFF\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Standard A "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "theory = {\n",
    "    'sigma-0.5': {\n",
    "        'spinodal': {\n",
    "            'alpha_cross': np.sort([0.5480511155124588,0.7071599112677189,0.9277860822926854,1.2410177405149914,1.06,1.4541109691229464,1.93929723421423,1.7505718608649372,0.8180015676221762]),\n",
    "            'omegas':np.sort([0.,.1,.2,.3,.25,.35,.425,.4,.15]),\n",
    "        },\n",
    "        'att-linear': {\n",
    "            'alpha_cross': np.sort([0.7224524901475151,0.8505675450846425, 1.0714067387463304, 1.4505799369795134,1.2254411566145098,1.749529526524848,1.9932320281583897,0.941534866004223,0.7652803696599021,1.5843728192149278,1.32705219011143,0.8076328324972342]),\n",
    "            'omegas':np.sort([0.,.1,.2,.3,.25,.35,.38,.15,.05,0.325,.275,0.075]),\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "df_nothing = pd.read_csv('theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.3_agno.csv') #fixed point that learns nothing\n",
    "df_positional = pd.read_csv('theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.3_para.csv') #positional\n",
    "df_semantic = pd.read_csv('theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.3.csv') #semantic\n",
    "\n",
    "#df_positional = pd.read_csv('results/theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.5_para.csv') #positional\n",
    "#df_semantic = pd.read_csv('results/theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.5.csv') #semantic\n",
    "\n",
    "\n",
    "df_spinodal = pd.DataFrame(theory['sigma-0.5']['spinodal']).set_index('omegas')\n",
    "df_attlin = pd.DataFrame(theory['sigma-0.5']['att-linear']).set_index('omegas')\n",
    "\n",
    "# todo: para and ferromagnetic fixed point for omega=0.3, their theta and their free energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_spinodal(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['spinodal']['alpha_cross'], theory['sigma-0.5']['spinodal']['omegas'], label=label, color=c_spinodal, linestyle='--', linewidth=2)\n",
    "    \n",
    "def plot_att_linear(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['att-linear']['alpha_cross'], theory['sigma-0.5']['att-linear']['omegas'], label=label, color=c_attlin, linestyle='--', linewidth=2)\n",
    "    \n",
    "def plot_spinodal_faint(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['spinodal']['alpha_cross'], theory['sigma-0.5']['spinodal']['omegas'], label=label, color=c_spinodal, linestyle='--', linewidth=2,alpha=0.3)\n",
    "    \n",
    "def plot_att_linear_faint(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['att-linear']['alpha_cross'], theory['sigma-0.5']['att-linear']['omegas'], label=label, color=c_attlin, linestyle='--', linewidth=2,alpha=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('empirics/results_mean_standardA.csv')\n",
    "df_std = pd.read_csv('empirics/results_std_standardA.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "d=1000\n",
    "L=2\n",
    "r=1\n",
    "delta=0.4\n",
    "lmbda_linear = 0.0001\n",
    "sigma=0.5\n",
    "lmbda=0.01\n",
    "omega=0.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28464 characters]",
      "text/plain": [
       "<Figure size 400x350 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "25 25\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36480 characters]",
      "text/plain": [
       "<Figure size 400x350 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def phase_diagrams(df,name=''):\n",
    "\n",
    "    fig=plt.figure(figsize=(4,3.5))\n",
    "    d=1000\n",
    "    L=2\n",
    "    r=1\n",
    "    sigma=0.5\n",
    "    delta=0.4\n",
    "    lmbda_linear = 0.0001\n",
    "    lmbda=0.01\n",
    "\n",
    "    lin_error = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==True][df.optim=='GD'][df.omega<=0.5][df.sigma==sigma].set_index(['alpha','omega']).linear_gen_error_mean\n",
    "    lin_loss = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==True][df.optim=='GD'][df.omega<=0.5][df.sigma==sigma].set_index(['alpha','omega']).linear_gen_error_mean\n",
    "    att_error_informed = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==True][df.optim=='GD'][df.omega<=0.5][df.sigma==sigma].set_index(['alpha','omega']).attention_gen_error_mean\n",
    "    att_error_positional = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==False][df.optim=='GD'][df.omega<=0.5][df.sigma==sigma].set_index(['alpha','omega']).attention_gen_error_mean\n",
    "    att_train_loss_informed = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==True][df.optim=='GD'][df.omega<=0.5][df.sigma==sigma].set_index(['alpha','omega']).attention_train_error_mean\n",
    "    att_train_loss_positional = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==False][df.optim=='GD'][df.omega<=0.5][df.sigma==sigma].set_index(['alpha','omega']).attention_train_error_mean\n",
    "\n",
    "    mask=(att_train_loss_informed<= att_train_loss_positional)*1\n",
    "    mask2=(att_train_loss_informed>= att_train_loss_positional)*1\n",
    "    \n",
    "    att_error_informed=mask*att_error_informed+mask2*att_error_positional\n",
    "    \n",
    "    loss_diff = (att_train_loss_informed - att_train_loss_positional).reset_index()#.unstack()\n",
    "    error_diff2 = (lin_error-att_error_informed).reset_index()\n",
    "    \n",
    "    \n",
    "\n",
    "    vmax=0.0015\n",
    "    alphas = sorted(loss_diff.alpha.unique())\n",
    "    omegas = sorted(loss_diff.omega.unique())\n",
    "    plt.pcolor(alphas,omegas,loss_diff.attention_train_error_mean.values.reshape(len(alphas),len(omegas))/1000,vmin=-vmax,vmax=vmax,cmap=cmap_uninf)\n",
    "\n",
    "\n",
    "    #plt.scatter(loss_diff.alpha,loss_diff.omega,c=loss_diff.attention_train_error_mean/1000,vmin=-vmax,vmax=vmax,s=100,cmap=cmap_uninf)\n",
    "    plt.colorbar(label=r'$\\Delta \\epsilon_t$')\n",
    "    plt.xlabel(r'$\\alpha$')\n",
    "    plt.ylabel(r'$\\omega$')\n",
    "    if name == '':\n",
    "        plot_spinodal(plt.gca())\n",
    "    #plot_att_linear(plt.gca())\n",
    "    plt.ylim(0.01,0.5)\n",
    "    plt.xlim(0,2.0)\n",
    "    plt.savefig(figure_dir / f'attention_train_error_pd{name}.pdf', bbox_inches='tight')\n",
    "    plt.show()\n",
    "\n",
    "    fig=plt.figure(figsize=(4,3.5))\n",
    "    vmax = 0.002\n",
    "    alphas = sorted(error_diff2.alpha.unique())\n",
    "    omegas = sorted(error_diff2.omega.unique())\n",
    "    print(len(alphas),len(omegas))\n",
    "    plt.pcolor(alphas,omegas,error_diff2[0].values.reshape(len(alphas),len(omegas)),vmin=-vmax,vmax=vmax,cmap=cmap_attlin)\n",
    "\n",
    "    #plt.scatter(error_diff2.alpha,error_diff2.omega,c=error_diff2[0],s=100,vmin=-vmax,vmax=vmax,cmap=cmap_attlin)\n",
    "    plt.colorbar(label=r'$\\epsilon^{\\mathrm{lin}}_g-\\epsilon_g$')\n",
    "    plt.xlabel(r'$\\alpha$')\n",
    "    plt.ylabel(r'$\\omega$')\n",
    "    if name == '':\n",
    "        plot_spinodal(plt.gca())\n",
    "        plot_att_linear(plt.gca())\n",
    "    plt.ylim(0.01,0.5)\n",
    "    plt.xlim(0,2.0)\n",
    "    plt.savefig(figure_dir / f'attention_gen_error_pd{name}.pdf', bbox_inches='tight')\n",
    "    plt.show()\n",
    "\n",
    "phase_diagrams(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "d=1000\n",
    "L=2\n",
    "r=1\n",
    "delta=0.4\n",
    "omega=0.34\n",
    "lmbda_linear = 0.0001\n",
    "sigma=0.5\n",
    "lmbda=0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23892 characters]",
      "text/plain": [
       "<Figure size 400x380 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig=plt.figure(figsize=(4,3.8))\n",
    "omega=0.3\n",
    "condition = (\n",
    "    (df['omega'] == omega) &\n",
    "    (df['sigma'] == sigma) &\n",
    "    (df['attention_lmbda'] == lmbda) &\n",
    "    (df['optim'] == 'GD')\n",
    ")\n",
    "\n",
    "dfi = df.loc[condition & (df['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu = df.loc[condition & (df['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "condition_std = (\n",
    "    (df_std['omega'] == omega) &\n",
    "    (df_std['sigma'] == sigma) &\n",
    "    (df_std['attention_lmbda'] == lmbda) &\n",
    "    (df_std['optim'] == 'GD')\n",
    ")\n",
    "\n",
    "dfi_std = df_std.loc[condition_std & (df_std['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu_std = df_std.loc[condition_std & (df_std['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "dfx = (dfu.attention_train_error_mean-dfi.attention_train_error_mean).reset_index()\n",
    "dfx_std = (dfu_std.attention_train_error_mean+dfi_std.attention_train_error_mean).reset_index()\n",
    "\n",
    "dfi.reset_index(inplace=True)\n",
    "dfu.reset_index(inplace=True)\n",
    "\n",
    "# todo add theory\n",
    "\n",
    "d = 1000\n",
    "\n",
    "dfi_std.reset_index(inplace = True)\n",
    "dfu_std.reset_index(inplace = True)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#plt.plot(df_semantic.alpha,df_semantic.theta/sigma**2, \n",
    "#         label='semantic', color=c_semantic,ls=\"--\")\n",
    "plt.plot(df_semantic[df_semantic[\"alpha\"]>=1.24].alpha,df_semantic[df_semantic[\"alpha\"]>=1.24].theta/sigma**2, \n",
    "          color=c_semantic)\n",
    "plt.plot(df_positional[df_positional[\"alpha\"]>=1.24].alpha,df_positional[df_positional[\"alpha\"]>=1.24].theta/sigma**2, color=c_positional)\n",
    "#plt.plot(df_positional[df_positional[\"alpha\"]>=1.24].alpha,df_positional[df_positional[\"alpha\"]>=1.24].theta/sigma**2,\n",
    "#         label='positional', color=c_positional)\n",
    "plt.scatter(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_theta_mean,label=r'$\\theta/\\sigma^2$', color=c_semantic)\n",
    "plt.errorbar(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_theta_mean,yerr=dfi_std[dfi_std[\"alpha\"]>=1.24].attention_theta_mean, color=c_semantic, linestyle='None')\n",
    "plt.scatter(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_theta_mean, color=c_semantic)\n",
    "plt.errorbar(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_theta_mean,yerr=dfu_std[dfu_std[\"alpha\"]<=1.24].attention_theta_mean, color=c_semantic, linestyle='None')\n",
    "\n",
    "\n",
    "\n",
    "plt.plot(df_semantic[df_semantic[\"alpha\"]<=1.24].alpha,df_semantic[df_semantic[\"alpha\"]<=1.24].m/sigma**2,  color=c_semantic)\n",
    "#plt.plot(df_positional.alpha,df_positional.m/sigma**2, label='positional', color=c_positional,linestyle=\"--\")\n",
    "plt.plot(df_positional[df_positional[\"alpha\"]<=1.24].alpha,df_positional[df_positional[\"alpha\"]<=1.24].m/sigma**2,\n",
    "    color=c_positional)\n",
    "plt.scatter(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_mag_mean/sigma**2, color=c_positional)\n",
    "plt.errorbar(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_mag_mean/sigma**2,yerr=dfu_std[dfu_std[\"alpha\"]<=1.24].attention_mag_mean/sigma**2, color=c_positional, linestyle='None')\n",
    "plt.scatter(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_mag_mean/sigma**2,label=r'$m/\\sigma^2$', color=c_positional)\n",
    "plt.errorbar(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_mag_mean/sigma**2,yerr=dfi_std[dfi_std[\"alpha\"]>=1.24].attention_mag_mean/sigma**2, color=c_positional, linestyle='None')\n",
    "\n",
    "plt.xlabel(r'$\\alpha$')\n",
    "\n",
    "plt.axvline(df_spinodal.loc[0.3].alpha_cross,ymin=0.0, ymax=1, color=\"g\", ls=\"--\")\n",
    "plt.xticks(ticks=list(plt.xticks()[0]) + [df_spinodal.loc[0.3].alpha_cross],labels=[str(i) for i in plt.xticks()[0]]+[r\"$\\alpha_c$\"])\n",
    "\n",
    "\n",
    "\n",
    "plt.legend()\n",
    "plt.xlim(0.0,2.0)\n",
    "plt.savefig(figure_dir / 'attention_overlap_m_theta.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30172 characters]",
      "text/plain": [
       "<Figure size 400x380 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig=plt.figure(figsize=(4,3.8))\n",
    "omega=0.3\n",
    "condition = (\n",
    "    (df['omega'] == omega) &\n",
    "    (df['sigma'] == sigma) &\n",
    "    (df['attention_lmbda'] == lmbda) &\n",
    "    (df['optim'] == 'GD')\n",
    ")\n",
    "\n",
    "dfi = df.loc[condition & (df['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu = df.loc[condition & (df['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "condition_std = (\n",
    "    (df_std['omega'] == omega) &\n",
    "    (df_std['sigma'] == sigma) &\n",
    "    (df_std['attention_lmbda'] == lmbda) &\n",
    "    (df_std['optim'] == 'GD')\n",
    ")\n",
    "\n",
    "dfi_std = df_std.loc[condition_std & (df_std['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu_std = df_std.loc[condition_std & (df_std['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "dfx = (dfu.attention_train_error_mean-dfi.attention_train_error_mean).reset_index()\n",
    "dfx_std = (dfu_std.attention_train_error_mean+dfi_std.attention_train_error_mean).reset_index()\n",
    "\n",
    "dfi.reset_index(inplace=True)\n",
    "dfu.reset_index(inplace=True)\n",
    "\n",
    "dfi_std.reset_index(inplace=True)\n",
    "dfu_std.reset_index(inplace=True)\n",
    "\n",
    "# todo add theory\n",
    "\n",
    "d = 1000\n",
    "\n",
    "c_neutral = c_att#'black'\n",
    "\n",
    "lin_error = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==False][df.optim=='GD'][df.omega==omega][df.sigma==sigma]#.mean()\n",
    "plt.axhline(lin_error.linear_gen_error_mean.mean(), color=c_lin, label=r'$\\epsilon_g^{\\mathrm{lin}}$')#,ls=\"--\")\n",
    "# fill between the std\n",
    "plt.fill_between([0,2],lin_error.linear_gen_error_mean.mean()-lin_error.linear_gen_error_mean.std(),lin_error.linear_gen_error_mean.mean()+lin_error.linear_gen_error_mean.std(),color=c_lin,alpha=0.2)\n",
    "\n",
    "plt.scatter(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_gen_error_mean, color=c_neutral)\n",
    "plt.scatter(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_gen_error_mean,color=c_neutral)\n",
    "plt.errorbar(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_gen_error_mean,yerr=dfi_std[dfi_std[\"alpha\"]>=1.24].attention_gen_error_mean, color=c_neutral, linestyle='None')\n",
    "plt.errorbar(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_gen_error_mean,yerr=dfu_std[dfu_std[\"alpha\"]<=1.24].attention_gen_error_mean, color=c_neutral, linestyle='None')\n",
    "\n",
    "#print(df_nothing.columns)\n",
    "#plt.plot(df_nothing.alpha,df_nothing.mse/2, label='uninformed dot-p attention', color='black')\n",
    "plt.plot(df_positional[df_positional[\"alpha\"]<1.25].alpha,df_positional[df_positional[\"alpha\"]<1.25].mse/2, color=c_neutral,label=r'$\\epsilon_g^{\\mathrm{minloss}}$')\n",
    "plt.plot(df_semantic[df_semantic[\"alpha\"]>1.24].alpha,df_semantic[df_semantic[\"alpha\"]>1.24].mse/2, color=c_neutral)\n",
    "\n",
    "plt.axvline(df_spinodal.loc[0.3].alpha_cross,ymin=0, ymax=1, color=\"g\", ls=\"--\")\n",
    "plt.axvline(1.45,ymin=0, ymax=1, color=\"r\", ls=\"--\")\n",
    "\n",
    "plt.xticks(ticks=[0.,.5,1.,2.] + [df_spinodal.loc[0.3].alpha_cross,1.45],labels=[str(i) for i in [0.,.5,1.,2.]]+[r\"$\\alpha_c$\",r\"$\\alpha_l$\"])\n",
    "\n",
    "\n",
    "plt.ylabel(r'$\\epsilon_g$')\n",
    "plt.xlabel(r'$\\alpha$')\n",
    "plt.legend()\n",
    "plt.xlim(0.0,2.0)\n",
    "plt.savefig(figure_dir / 'attention_gen_error_global_only.pdf', bbox_inches='tight')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34652 characters]",
      "text/plain": [
       "<Figure size 400x380 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "fig=plt.figure(figsize=(4,3.8))\n",
    "omega=0.3\n",
    "condition = (\n",
    "    (df['omega'] == omega) &\n",
    "    (df['sigma'] == sigma) &\n",
    "    (df['attention_lmbda'] == lmbda) &\n",
    "    (df['optim'] == 'GD')\n",
    ")\n",
    "\n",
    "dfi = df.loc[condition & (df['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu = df.loc[condition & (df['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "condition_std = (\n",
    "    (df_std['omega'] == omega) &\n",
    "    (df_std['sigma'] == sigma) &\n",
    "    (df_std['attention_lmbda'] == lmbda) &\n",
    "    (df_std['optim'] == 'GD')\n",
    ")\n",
    "\n",
    "dfi_std = df_std.loc[condition_std & (df_std['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu_std = df_std.loc[condition_std & (df_std['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "dfx = (dfu.attention_train_error_mean-dfi.attention_train_error_mean).reset_index()\n",
    "dfx_std = (dfu_std.attention_train_error_mean+dfi_std.attention_train_error_mean).reset_index()\n",
    "\n",
    "dfi.reset_index(inplace=True)\n",
    "dfu.reset_index(inplace=True)\n",
    "\n",
    "# todo add theory\n",
    "\n",
    "d = 1000\n",
    "\n",
    "\n",
    "\n",
    "vmax=0.0015\n",
    "plt.plot(df_semantic.alpha,-df_semantic.Phi+df_positional.Phi,c='grey')\n",
    "plt.errorbar(dfx.alpha[::2],(-dfx.attention_train_error_mean/d)[::2],yerr=(dfx_std.attention_train_error_mean/d)[::2], color='grey', linestyle='None')\n",
    "plt.axhline(0,linestyle='--',color='grey')\n",
    "plt.scatter(dfx.alpha[::2], (-dfx.attention_train_error_mean/d)[::2], c=(dfx.attention_train_error_mean/d)[::2], vmin=-vmax, vmax=vmax, cmap=cmap_uninf, s=80, zorder=100, linewidths=1, edgecolors='grey')\n",
    "plt.axvline(df_spinodal.loc[0.3].alpha_cross, color=c_spinodal,ls=\"--\")\n",
    "plt.xlabel(r'$\\alpha$')\n",
    "plt.ylabel(r'$\\Delta \\epsilon_t}$')\n",
    "plt.xlim(0.0,2.0)\n",
    "\n",
    "#print(plt.xticks())\n",
    "plt.xticks(ticks=list(plt.xticks()[0]) + [df_spinodal.loc[0.3].alpha_cross],labels=[str(i) for i in plt.xticks()[0]]+[r\"$\\alpha_c$\"])\n",
    "\n",
    "plt.savefig(figure_dir / 'attention_free_entropy.pdf', bbox_inches='tight')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.5  0.   0.5  1.   1.5  2.   2.5]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 59940 characters]",
      "text/plain": [
       "<Figure size 1275x285 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "factor = 0.75\n",
    "c_neutral = c_att#'black'\n",
    "\n",
    "fig, axes = plt.subplots(1,3,figsize=(17*factor,3.8*factor))\n",
    "\n",
    "omega=0.3\n",
    "condition = (\n",
    "    (df['omega'] == omega) &\n",
    "    (df['sigma'] == sigma) &\n",
    "    (df['attention_lmbda'] == lmbda) &\n",
    "    (df['optim'] == 'GD')\n",
    ")\n",
    "\n",
    "dfi = df.loc[condition & (df['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu = df.loc[condition & (df['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "\n",
    "dfi_std = df_std.loc[condition & (df_std['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu_std = df_std.loc[condition & (df_std['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "dfx = (dfu.attention_train_error_mean-dfi.attention_train_error_mean).reset_index()\n",
    "dfx_std = (dfu_std.attention_train_error_mean+dfi_std.attention_train_error_mean).reset_index()\n",
    "\n",
    "dfi.reset_index(inplace=True)\n",
    "dfu.reset_index(inplace=True)\n",
    "\n",
    "# todo add theory\n",
    "\n",
    "d = 1000\n",
    "\n",
    "# first axes\n",
    "ax = axes[0]\n",
    "vmax=0.0015\n",
    "ax.plot(df_semantic.alpha,-df_semantic.Phi+df_positional.Phi,c='grey')\n",
    "ax.errorbar(dfx.alpha[::2],(-dfx.attention_train_error_mean/d)[::2],yerr=(dfx_std.attention_train_error_mean/d)[::2], color='grey', linestyle='None')\n",
    "ax.axhline(0,linestyle='--',color='grey')\n",
    "ax.scatter(dfx.alpha[::2], (-dfx.attention_train_error_mean/d)[::2], c=(-dfx.attention_train_error_mean/d)[::2], vmin=-vmax, vmax=vmax, cmap=cmap_uninf, s=80, zorder=100, linewidths=1, edgecolors='grey')\n",
    "ax.axvline(df_spinodal.loc[0.3].alpha_cross, color=c_spinodal,ls=\"--\")\n",
    "ax.set_xlabel(r'$\\alpha$')\n",
    "ax.set_ylabel(r'$\\Delta \\epsilon_t}$')\n",
    "\n",
    "#print(plt.xticks())\n",
    "print(ax.get_xticks())\n",
    "ax.set_ylim(-0.00125,0.00075)\n",
    "#ax.set_yticks(np.arange( -0.00125,0.00075, 0.00025))\n",
    "\n",
    "ax.set_xticks(ticks=list(ax.get_xticks()) + [df_spinodal.loc[0.3].alpha_cross],labels=[str(i) for i in ax.get_xticks()]+[r\"$\\alpha_c$\"])\n",
    "ax.set_xlim(0.0,2.0)\n",
    "# second axes \n",
    "\n",
    "dfi = df.loc[condition & (df['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu = df.loc[condition & (df['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "dfi_std = df_std.loc[condition & (df_std['informed'] == True)].set_index(['alpha', 'omega'])\n",
    "dfu_std = df_std.loc[condition & (df_std['informed'] == False)].set_index(['alpha', 'omega'])\n",
    "\n",
    "dfx = (dfu.attention_train_error_mean-dfi.attention_train_error_mean).reset_index()\n",
    "dfx_std = (dfu_std.attention_train_error_mean+dfi_std.attention_train_error_mean).reset_index()\n",
    "\n",
    "dfi.reset_index(inplace=True)\n",
    "dfu.reset_index(inplace=True)\n",
    "ax = axes[1]\n",
    "\n",
    "\n",
    "dfi_std.reset_index(inplace = True)\n",
    "dfu_std.reset_index(inplace = True)\n",
    "\n",
    "\n",
    "#plt.plot(df_semantic.alpha,df_semantic.theta/sigma**2, \n",
    "#         label='semantic', color=c_semantic,ls=\"--\")\n",
    "ax.plot(df_semantic[df_semantic[\"alpha\"]>=1.24].alpha,df_semantic[df_semantic[\"alpha\"]>=1.24].theta/sigma**2, \n",
    "          color=c_semantic)\n",
    "ax.plot(df_positional[df_positional[\"alpha\"]>=1.24].alpha,df_positional[df_positional[\"alpha\"]>=1.24].theta/sigma**2, color=c_positional)\n",
    "#plt.plot(df_positional[df_positional[\"alpha\"]>=1.24].alpha,df_positional[df_positional[\"alpha\"]>=1.24].theta/sigma**2,\n",
    "#         label='positional', color=c_positional)\n",
    "ax.scatter(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_theta_mean,label=r'$\\theta/\\sigma^2$', color=c_semantic)\n",
    "ax.errorbar(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_theta_mean,yerr=dfi_std[dfi_std[\"alpha\"]>=1.24].attention_theta_mean, color=c_semantic, linestyle='None')\n",
    "ax.scatter(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_theta_mean, color=c_semantic)\n",
    "ax.errorbar(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_theta_mean,yerr=dfu_std[dfu_std[\"alpha\"]<=1.24].attention_theta_mean, color=c_semantic, linestyle='None')\n",
    "\n",
    "\n",
    "\n",
    "ax.plot(df_semantic[df_semantic[\"alpha\"]<=1.24].alpha,df_semantic[df_semantic[\"alpha\"]<=1.24].m/sigma**2,  color=c_semantic)\n",
    "#plt.plot(df_positional.alpha,df_positional.m/sigma**2, label='positional', color=c_positional,linestyle=\"--\")\n",
    "ax.plot(df_positional[df_positional[\"alpha\"]<=1.24].alpha,df_positional[df_positional[\"alpha\"]<=1.24].m/sigma**2,\n",
    "    color=c_positional)\n",
    "ax.scatter(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_mag_mean/sigma**2, color=c_positional)\n",
    "ax.errorbar(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_mag_mean/sigma**2,yerr=dfu_std[dfu_std[\"alpha\"]<=1.24].attention_mag_mean/sigma**2, color=c_positional, linestyle='None')\n",
    "ax.scatter(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_mag_mean/sigma**2,label=r'$m/\\sigma^2$', color=c_positional)\n",
    "ax.errorbar(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_mag_mean/sigma**2,yerr=dfi_std[dfi_std[\"alpha\"]>=1.24].attention_mag_mean/sigma**2, color=c_positional, linestyle='None')\n",
    "\n",
    "ax.set_xlabel(r'$\\alpha$')\n",
    "\n",
    "ax.axvline(df_spinodal.loc[0.3].alpha_cross,ymin=0.0, ymax=1, color=\"g\", ls=\"--\")\n",
    "ax.set_xticks(ticks=list(ax.get_xticks()) + [df_spinodal.loc[0.3].alpha_cross],labels=[str(i) for i in ax.get_xticks()]+[r\"$\\alpha_c$\"])\n",
    "\n",
    "\n",
    "\n",
    "ax.legend()\n",
    "ax.set_xlim(0.0,2.0)\n",
    "\n",
    "# third axes\n",
    "c_neutral = c_att#'black'\n",
    "ax =axes[2]\n",
    "lin_error = df[df.L==L][df.r==r][df.d==d][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==False][df.optim=='GD'][df.omega==omega][df.sigma==sigma]#.mean()\n",
    "ax.axhline(lin_error.linear_gen_error_mean.mean(), color=c_lin, label=r'$\\epsilon_g^{\\mathrm{lin}}$')#,ls=\"--\")\n",
    "# fill between the std\n",
    "ax.fill_between([0,2],lin_error.linear_gen_error_mean.mean()-lin_error.linear_gen_error_mean.std(),lin_error.linear_gen_error_mean.mean()+lin_error.linear_gen_error_mean.std(),color=c_lin,alpha=0.2)\n",
    "\n",
    "ax.scatter(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_gen_error_mean, color=c_neutral)\n",
    "ax.scatter(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_gen_error_mean,color=c_neutral)\n",
    "ax.errorbar(dfi[dfi[\"alpha\"]>=1.24].alpha,dfi[dfi[\"alpha\"]>=1.24].attention_gen_error_mean,yerr=dfi_std[dfi_std[\"alpha\"]>=1.24].attention_gen_error_mean, color=c_neutral, linestyle='None')\n",
    "ax.errorbar(dfu[dfu[\"alpha\"]<=1.24].alpha,dfu[dfu[\"alpha\"]<=1.24].attention_gen_error_mean,yerr=dfu_std[dfu_std[\"alpha\"]<=1.24].attention_gen_error_mean, color=c_neutral, linestyle='None')\n",
    "\n",
    "#print(df_nothing.columns)\n",
    "#plt.plot(df_nothing.alpha,df_nothing.mse/2, label='uninformed dot-p attention', color='black')\n",
    "ax.plot(df_positional[df_positional[\"alpha\"]<1.25].alpha,df_positional[df_positional[\"alpha\"]<1.25].mse/2, color=c_neutral,label=r'$\\epsilon_g^{\\mathrm{min}}$')\n",
    "ax.plot(df_semantic[df_semantic[\"alpha\"]>1.24].alpha,df_semantic[df_semantic[\"alpha\"]>1.24].mse/2, color=c_neutral)\n",
    "\n",
    "ax.axvline(df_spinodal.loc[0.3].alpha_cross,ymin=0, ymax=1, color=\"g\", ls=\"--\")\n",
    "ax.axvline(1.45,ymin=0, ymax=1, color=\"r\", ls=\"--\")\n",
    "\n",
    "ax.set_xticks(ticks=[0.,.5,1.,2.] + [df_spinodal.loc[0.3].alpha_cross,1.45],labels=[str(i) for i in [0.,.5,1.,2.]]+[r\"$\\alpha_c$\",r\"$\\alpha_l$\"])\n",
    "\n",
    "ax.set_xlabel(r'$\\alpha$')\n",
    "ax.legend()\n",
    "ax.set_xlim(0.0,2.0)\n",
    "\n",
    "fig.subplots_adjust(wspace=0.25)\n",
    "\n",
    "plt.savefig(figure_dir / 'all_omega_0.3.pdf', bbox_inches='tight')\n",
    "\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scaling behaviour"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('empirics/results_scaling.csv')\n",
    "df_std = pd.read_csv('empirics/results_std_scaling.csv')\n",
    "df_orig = pd.read_csv('empirics/results_scaling_orig.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39728 characters]",
      "text/plain": [
       "<Figure size 640x480 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df_orig = df_orig[df_orig.alpha==1.5]\n",
    "df_orig = df_orig.groupby('d').head(30)\n",
    "sigma=0.5\n",
    "plt.scatter(df_orig[df_orig.informed].attention_mag_mean/sigma**2, df_orig[df_orig.informed].attention_theta_mean, marker='+',c=df_orig[df_orig.informed].d, cmap=cmap_att, norm=colors.LogNorm())\n",
    "plt.colorbar(label=r'$d$')\n",
    "plt.scatter(df_orig[~df_orig.informed].attention_mag_mean/sigma**2,df_orig[~df_orig.informed].attention_theta_mean,  marker='+',c=df_orig[~df_orig.informed].d, cmap=cmap_pos\n",
    "            , norm=colors.LogNorm())\n",
    "plt.colorbar()\n",
    "plt.scatter([],[],marker='+',c=c_semantic,label='semantic')\n",
    "plt.scatter([],[],marker='+',c=c_positional,label='positional')\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(r'$m/\\sigma^2$')\n",
    "plt.ylabel(r'$\\theta/\\sigma^2$')\n",
    "plt.savefig(figure_dir / 'd_scaling_omega_m_theta.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 26072 characters]",
      "text/plain": [
       "<Figure size 900x285 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "alpha = 1.5\n",
    "\n",
    "fig, axes = plt.subplots(1,2,figsize=(12*factor,3.8*factor),sharey=True)\n",
    "\n",
    "condition = (\n",
    "        (df['omega'] == omega) &\n",
    "        (df['sigma'] == sigma) &\n",
    "        (df['attention_lmbda'] == lmbda) &\n",
    "        (df['optim'] == 'GD') &\n",
    "        (df['alpha'] == alpha)\n",
    "    )\n",
    "\n",
    "dfi = df.loc[condition & (df['informed'] == True)].set_index(['d','omega'])\n",
    "dfu = df.loc[condition & (df['informed'] == False)].set_index(['d','omega'])\n",
    "\n",
    "dfi_std = df_std.loc[condition & (df_std['informed'] == True)].set_index(['d','omega'])\n",
    "dfu_std = df_std.loc[condition & (df_std['informed'] == False)].set_index(['d','omega'])\n",
    "\n",
    "dfi.reset_index(inplace=True)\n",
    "dfu.reset_index(inplace=True)\n",
    "dfi_std.reset_index(inplace = True)\n",
    "dfu_std.reset_index(inplace = True)\n",
    "\n",
    "mse_s = (df_semantic.mse.values[4]+df_semantic.mse.values[5])/2\n",
    "mse_p = (df_positional.mse.values[4]+df_positional.mse.values[5])/2\n",
    "    \n",
    "    \n",
    "# third axes\n",
    "c_neutral = c_att#'black'\n",
    "ax =  axes[0]\n",
    "\n",
    "ax.scatter(dfi.d,dfi.attention_gen_error_mean, color=c_semantic)\n",
    "\n",
    "ax.errorbar(dfi.d,dfi.attention_gen_error_mean,yerr=dfi_std.attention_gen_error_mean, color=c_semantic, linestyle='None')\n",
    "ax.set_ylabel(r'$\\epsilon_g$')\n",
    "ax.axhline(mse_s/2, color=c_semantic,linestyle='dashed')\n",
    "ax.axhline(mse_p/2, color=c_positional,linestyle='dashed')\n",
    "ax.set_xlabel(r'$d$')\n",
    "ax.set_xscale('log')\n",
    "\n",
    "ax =  axes[1]\n",
    "\n",
    "ax.scatter(dfu.d,dfu.attention_gen_error_mean,color=c_positional)\n",
    "ax.errorbar(dfu.d,dfu.attention_gen_error_mean,yerr=dfu_std.attention_gen_error_mean, color=c_positional, linestyle='None')\n",
    "ax.axhline(mse_s/2, color=c_semantic,linestyle='dashed')\n",
    "ax.axhline(mse_p/2, color=c_positional,linestyle='dashed')\n",
    "\n",
    "ax.scatter([],[], color=c_positional,label='positional',marker='s')\n",
    "ax.scatter([],[], color=c_semantic,label='semantic',marker='s')\n",
    "ax.scatter([],[], color='black',label='empirircs')\n",
    "ax.plot([],[], color='black',label='theory',linestyle='dashed')\n",
    "ax.set_xlabel(r'$d$')\n",
    "ax.set_xscale('log')\n",
    "ax.legend()\n",
    "plt.savefig(figure_dir / 'd_scaling_omega_0.3.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Alternative A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('empirics/results_mean_otherA.csv')\n",
    "df_std = pd.read_csv('empirics/results_std_otherA.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "theory = {\n",
    "    'sigma-0.5': {\n",
    "        'spinodal': {\n",
    "            'alpha_cross': np.sort([0.5480511155124588,0.7071599112677189,0.9277860822926854,1.2410177405149914,1.06,1.4541109691229464,1.93929723421423,1.7505718608649372,0.8180015676221762]),\n",
    "            'omegas':np.sort([0.,.1,.2,.3,.25,.35,.425,.4,.15]),\n",
    "        },\n",
    "        'att-linear': {\n",
    "            'alpha_cross': np.sort([0.7224524901475151,0.8505675450846425, 1.0714067387463304, 1.4505799369795134,1.2254411566145098,1.749529526524848,1.9932320281583897,0.941534866004223,0.7652803696599021,1.5843728192149278,1.32705219011143,0.8076328324972342]),\n",
    "            'omegas':np.sort([0.,.1,.2,.3,.25,.35,.38,.15,.05,0.325,.275,0.075]),\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "#df_nothing = pd.read_csv('theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.3_agno.csv') #fixed point that learns nothing\n",
    "df_positional = pd.read_csv('theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.14_para_Aalt.csv') #positional\n",
    "df_semantic = pd.read_csv('theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.14_Aalt.csv') #semantic\n",
    "df_nothing = df_semantic\n",
    "#df_positional = pd.read_csv('results/theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.5_para.csv') #positional\n",
    "#df_semantic = pd.read_csv('results/theory/repl_softmax_Mixed_sig0.5_lamb0.01_omega0.5.csv') #semantic\n",
    "\n",
    "\n",
    "df_spinodal = pd.DataFrame(theory['sigma-0.5']['spinodal']).set_index('omegas')\n",
    "df_attlin = pd.DataFrame(theory['sigma-0.5']['att-linear']).set_index('omegas')\n",
    "\n",
    "# todo: para and ferromagnetic fixed point for omega=0.3, their theta and their free energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_spinodal(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['spinodal']['alpha_cross'], theory['sigma-0.5']['spinodal']['omegas'], label=label, color=c_spinodal, linestyle='--', linewidth=2)\n",
    "    \n",
    "def plot_att_linear(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['att-linear']['alpha_cross'], theory['sigma-0.5']['att-linear']['omegas'], label=label, color=c_attlin, linestyle='--', linewidth=2)\n",
    "    \n",
    "def plot_spinodal_faint(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['spinodal']['alpha_cross'], theory['sigma-0.5']['spinodal']['omegas'], label=label, color=c_spinodal, linestyle='--', linewidth=2,alpha=0.3)\n",
    "    \n",
    "def plot_att_linear_faint(ax,label=''):\n",
    "    ax.plot(theory['sigma-0.5']['att-linear']['alpha_cross'], theory['sigma-0.5']['att-linear']['omegas'], label=label, color=c_attlin, linestyle='--', linewidth=2,alpha=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22252 characters]",
      "text/plain": [
       "<Figure size 400x350 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "25 25\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25668 characters]",
      "text/plain": [
       "<Figure size 400x350 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "phase_diagrams(df,'otherA')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Appendix slices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('empirics/results_mean_standardA.csv')\n",
    "df_std = pd.read_csv('empirics/results_std_standardA.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65648 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "d=1000\n",
    "L=2\n",
    "r=1\n",
    "sigma=0.5\n",
    "delta=0.4\n",
    "lmbda_linear = 0.0001\n",
    "lmbda=0.01\n",
    "colors = ['royalblue','deeppink','orange','rebeccapurple','crimson']\n",
    "for i, omega in enumerate([0.2,0.3,0.7]):\n",
    "    att_error_informed = df[df.L==L][df.r==r][df.d==d][df.sigma==sigma][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==True][df.optim=='GD'][df.omega==omega].set_index(['alpha']).attention_gen_error_mean\n",
    "    att_error_positional = df[df.L==L][df.r==r][df.d==d][df.sigma==sigma][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==False][df.optim=='GD'][df.omega==omega].set_index(['alpha']).attention_gen_error_mean\n",
    "    # todo: here we need to take the min of the train error actually, not the gen error!\n",
    "    #att_error = np.min([att_error_informed, att_error_positional], axis=0)\n",
    "    #att_error_informed = att_error_informed.reset_index()\n",
    "    att_error_informed = att_error_informed.reset_index()\n",
    "    att_error_positional = att_error_positional.reset_index()\n",
    "    plt.plot(att_error_informed.alpha[::2],att_error_informed.attention_gen_error_mean[::2],lw=1,marker='o',color=colors[i])\n",
    "    plt.plot(att_error_positional.alpha[::2],att_error_positional.attention_gen_error_mean[::2],lw=1,marker='x',color=colors[i])\n",
    "    plt.plot([],[],label=omega,lw=1,color=colors[i])\n",
    "plt.plot([],[],label='semantic',lw=1,color='black',marker='o')\n",
    "plt.plot([],[],label='positional',lw=1,color='black',marker='x')\n",
    "plt.legend(title=r'$\\omega$')\n",
    "plt.xlabel(r'$\\alpha$')\n",
    "plt.ylabel(r'$\\epsilon_g$')\n",
    "plt.savefig(figure_dir / 'omega_slices.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49396 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "d=1000\n",
    "L=2\n",
    "r=1\n",
    "sigma=0.5\n",
    "delta=0.4\n",
    "omega=0.3\n",
    "lmbda_linear = 0.0001\n",
    "lmbda=0.01\n",
    "for i, sigma in enumerate([0.3,0.5,0.7]):\n",
    "    att_error_informed = df[df.L==L][df.r==r][df.d==d][df.sigma==sigma][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==True][df.optim=='GD'][df.omega==omega].set_index(['alpha']).attention_gen_error_mean\n",
    "    att_error_positional = df[df.L==L][df.r==r][df.d==d][df.sigma==sigma][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==False][df.optim=='GD'][df.omega==omega].set_index(['alpha']).attention_gen_error_mean\n",
    "    # todo: here we need to take the min of the train error actually, not the gen error!\n",
    "    #att_error = np.min([att_error_informed, att_error_positional], axis=0)\n",
    "    #att_error_informed = att_error_informed.reset_index()\n",
    "    att_error_informed = att_error_informed.reset_index()\n",
    "    att_error_positional = att_error_positional.reset_index()\n",
    "    plt.plot(att_error_informed.alpha[::2],att_error_informed.attention_gen_error_mean[::2],lw=1,marker='o',color=colors[i])\n",
    "    plt.plot(att_error_positional.alpha[::2],att_error_positional.attention_gen_error_mean[::2],lw=1,marker='x',color=colors[i])\n",
    "    plt.plot([],[],label=sigma,lw=1,color=colors[i])\n",
    "plt.plot([],[],label='semantic',lw=1,color='black',marker='o')\n",
    "plt.plot([],[],label='positional',lw=1,color='black',marker='x')\n",
    "plt.legend(title=r'$\\sigma$')\n",
    "plt.xlabel(r'$\\alpha$')\n",
    "plt.ylabel(r'$\\epsilon_g$')\n",
    "plt.savefig(figure_dir / 'sigma_slices.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54636 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "d=1000\n",
    "L=2\n",
    "r=1\n",
    "sigma=0.5\n",
    "delta=0.4\n",
    "omega=0.3\n",
    "lmbda_linear = 0.0001\n",
    "sigma=0.5\n",
    "\n",
    "for i,lmbda in enumerate([0.1,0.01,0.001]):\n",
    "    att_error_informed = df[df.L==L][df.r==r][df.d==d][df.sigma==sigma][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==True][df.optim=='GD'][df.omega==omega].set_index(['alpha']).attention_gen_error_mean\n",
    "    att_error_positional = df[df.L==L][df.r==r][df.d==d][df.sigma==sigma][df.attention_lmbda==lmbda][df.linear_lmbda==lmbda_linear][df.delta==delta][df.informed==False][df.optim=='GD'][df.omega==omega].set_index(['alpha']).attention_gen_error_mean\n",
    "    # todo: here we need to take the min of the train error actually, not the gen error!\n",
    "    #att_error = np.min([att_error_informed, att_error_positional], axis=0)\n",
    "    #att_error_informed = att_error_informed.reset_index()\n",
    "    att_error_informed = att_error_informed.reset_index()\n",
    "    att_error_positional = att_error_positional.reset_index()\n",
    "    plt.plot(att_error_informed.alpha[::2],att_error_informed.attention_gen_error_mean[::2],lw=1,marker='o',color=colors[i])\n",
    "    plt.plot(att_error_positional.alpha[::2],att_error_positional.attention_gen_error_mean[::2],lw=1,marker='x',color=colors[i])\n",
    "    plt.plot([],[],label=lmbda,lw=1,color=colors[i])\n",
    "plt.plot([],[],label='semantic',lw=1,color='black',marker='o')\n",
    "plt.plot([],[],label='positional',lw=1,color='black',marker='x')\n",
    "plt.legend(title=r'$\\lambda$')\n",
    "plt.xlabel(r'$\\alpha$')\n",
    "plt.ylabel(r'$\\epsilon_g$')\n",
    "plt.savefig(figure_dir / 'lmbda_slices.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# R = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('empirics/results_mean_r=2.csv')\n",
    "df_std = pd.read_csv('empirics/results_std_r=2.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(49, 19)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32480 characters]",
      "text/plain": [
       "<Figure size 800x350 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig=plt.figure(figsize=(8,3.5))\n",
    "vmax = 0.002\n",
    "df = df[df.omega < 0.4][df.omega < 2.01]\n",
    "error_diff = df.linear_gen_error_mean-df.attention_gen_error_mean\n",
    "alphas = sorted(df.alpha.unique())\n",
    "omegas = sorted(df.omega.unique())\n",
    "#plt.scatter(df.alpha,df.omega,c=error_diff, vmin=-vmax,vmax=vmax,label='attention',cmap=cmap_attlin)\n",
    "print(error_diff.values.reshape(len(alphas),len(omegas)).shape)\n",
    "plt.pcolor(alphas,omegas ,error_diff.values.reshape(len(omegas),len(alphas)),vmin=-vmax,vmax=vmax,cmap=cmap_attlin)\n",
    "plt.colorbar(label=r'$\\epsilon^{\\mathrm{lin}}_g-\\epsilon_g$')\n",
    "plt.xlabel(r'$\\alpha$')\n",
    "plt.ylabel(r'$\\omega$')\n",
    "plt.savefig(figure_dir / f'attention_gen_error_crossover_r=2.pdf', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison of GD and Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "def group(df):\n",
    "    return df.groupby(['d','L','r','delta','sigma','omega','attention_lmbda','linear_lmbda','alpha','informed','optim'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 26996 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df = pd.read_csv('empirics/results_samples_standardA.csv')\n",
    "df = group(df).head(24).reset_index()\n",
    "sigma=0.5\n",
    "plt.scatter(df[df.informed].attention_mag_mean/sigma**2, df[df.informed].attention_theta_mean,marker='+',color=c_semantic,label='semantic')\n",
    "plt.scatter(df[~df.informed].attention_mag_mean/sigma**2, df[~df.informed].attention_theta_mean,marker='+',color=c_positional,label='positional')\n",
    "plt.xlabel(r'$m/\\sigma^2$')\n",
    "plt.ylabel(r'$\\theta/\\sigma^2$')\n",
    "plt.xlim(0,1.2)\n",
    "plt.ylim(0,0.85)\n",
    "plt.legend()\n",
    "plt.title('informed initialization, GD, d=1000')\n",
    "plt.savefig(figure_dir / f'init-vary-2.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "141\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26704 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df = pd.read_csv('empirics/results_adam.csv')\n",
    "print(len(df))\n",
    "df = group(df).head(140).reset_index()\n",
    "plt.scatter(df.attention_mag_mean/sigma**2, df.attention_theta_mean,marker='+',color='black',label='random')\n",
    "plt.xlabel(r'$m/\\sigma^2$')\n",
    "plt.ylabel(r'$\\theta/\\sigma^2$')\n",
    "plt.xlim(0,1.2)\n",
    "plt.ylim(0,0.85)\n",
    "plt.title('random initialization, Adam, d=100')\n",
    "plt.savefig(figure_dir / f'init-vary-1.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# First Figure Landscape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2956388/1450812528.py:11: RuntimeWarning: covariance is not symmetric positive-semidefinite.\n",
      "  return np.random.multivariate_normal(mean, covariance, num_points)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 59984 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import gaussian_kde\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "# Set random seed for reproducibility\n",
    "np.random.seed(42)\n",
    "\n",
    "# Function to generate points from a Gaussian distribution\n",
    "def generate_gaussian_points(mean, covariance, num_points):\n",
    "    return np.random.multivariate_normal(mean, covariance, num_points)\n",
    "\n",
    "# Parameters for Gaussian distributions\n",
    "mean1 = np.array([10, 10])\n",
    "covariance1 =  np.array([[3, 2], [2, 1]])\n",
    "\n",
    "mean2 = np.array([0, 0])  # Increase the spacing\n",
    "covariance2 = np.array([[1, -0.5], [-0.5, 1]])\n",
    "covariance2 = np.array([[3, 2], [2, 1]])\n",
    "\n",
    "# Adjustable parameters\n",
    "spacing = 15  # Control the spacing between Gaussian distributions\n",
    "random_points_ratio = 4  # Control the ratio of random points to Gaussian points\n",
    "\n",
    "# Generate points from Gaussian distributions\n",
    "a = 30\n",
    "gaussian_points1 = generate_gaussian_points(mean1, covariance1, a)\n",
    "gaussian_points2 = generate_gaussian_points(mean2, covariance2, a)\n",
    "\n",
    "# Generate random uniformly sampled points\n",
    "num_random_points = len(gaussian_points1) * random_points_ratio\n",
    "uniform_points = np.random.uniform(low=-20, high=2, size=(int(num_random_points), 2))\n",
    "\n",
    "# Concatenate all points\n",
    "all_points = np.vstack((gaussian_points1, gaussian_points2, uniform_points))\n",
    "\n",
    "# Calculate KDE for the points\n",
    "kde = gaussian_kde(all_points.T)\n",
    "\n",
    "# Create a meshgrid for 3D plotting\n",
    "x_grid, y_grid = np.meshgrid(np.linspace(-10, 15, 100), np.linspace(-10, 15, 100))\n",
    "z_values = -kde(np.vstack([x_grid.ravel(), y_grid.ravel()]))\n",
    "\n",
    "# Reshape z_values to match the shape of x_grid and y_grid\n",
    "z_values = z_values.reshape(x_grid.shape)\n",
    "\n",
    "# Plot the 3D surface plot of KDE\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.plot_surface(x_grid, y_grid, z_values, cmap='Purples_r', alpha=0.8)\n",
    "\n",
    "# Remove grid and labels\n",
    "ax.xaxis.pane.fill = False\n",
    "ax.yaxis.pane.fill = False\n",
    "ax.zaxis.pane.fill = False\n",
    "ax.grid(False)\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.set_zticks([])\n",
    "ax.set_xlabel('')\n",
    "ax.set_ylabel('')\n",
    "ax.set_zlabel('')\n",
    "ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))\n",
    "ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))\n",
    "ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))\n",
    "# make the grid lines transparent\n",
    "ax.xaxis._axinfo[\"grid\"]['color'] =  (1,1,1,0)\n",
    "ax.yaxis._axinfo[\"grid\"]['color'] =  (1,1,1,0)\n",
    "ax.zaxis._axinfo[\"grid\"]['color'] =  (1,1,1,0)\n",
    "ax.set_axis_off()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
