{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <center>Plot the figures in paper</center>\n",
    "\n",
    "##### <center><font color='gray'>_Xuelong Sun_, _2020-04-01_</font></center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as mpatches\n",
    "from matplotlib.collections import PatchCollection\n",
    "import numpy as np\n",
    "import scipy.io as sio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize = 30\n",
    "r_folder = 'results/'\n",
    "d_folder = 'data/'\n",
    "\n",
    "# the simulated 3D world\n",
    "world = sio.loadmat(d_folder + 'world.mat')\n",
    "# load the data of mesh sampled ZMs\n",
    "fre_world_data = sio.loadmat(d_folder+'MeshSampled_ZMs_fixedHeading.mat')\n",
    "\n",
    "# the max order of ZM\n",
    "zm_n_max = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <font color='red'> visual homin (VH) </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_sub_folder = 'visual_homing/'\n",
    "\n",
    "from visual_homing import VisualHomingAgent\n",
    "\n",
    "# the route memory\n",
    "route_memory = sio.loadmat(d_folder + 'ArcRouteMemory.mat')\n",
    "# the home memory\n",
    "home_memory = sio.loadmat(d_folder + 'HomeMemory_X0Y0.mat')\n",
    "\n",
    "# set up parameters\n",
    "num_pn = 81\n",
    "num_kc = 4000\n",
    "vh_learning_rate = 0.1\n",
    "vh_kc_thr = 0.04\n",
    "# create an instance\n",
    "vh = VisualHomingAgent(world, route_memory, home_memory, zm_n_max, vh_learning_rate, vh_kc_thr, num_pn, num_kc)\n",
    "# training\n",
    "en = vh.train_mb_network()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _visual novelty heatmap_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48728 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# calculate the map\n",
    "map_x = fre_world_data['x'][0]\n",
    "map_y = fre_world_data['y'][0]\n",
    "map_ZM_As_random = fre_world_data['A']\n",
    "e_map = np.zeros([len(map_x), len(map_y)]) \n",
    "for xi in range(len(map_x)):\n",
    "    for yi in range(len(map_y)):\n",
    "        e_map[xi, yi] = vh.mb.run(map_ZM_As_random[xi, yi, :vh.zm_coeff_num])\n",
    "\n",
    "# draw the map\n",
    "fig,ax = plt.subplots(figsize=(8,8))\n",
    "\n",
    "ax.pcolormesh(map_x, np.flipud(map_x), np.rot90(e_map), cmap='RdPu')\n",
    "ax.set_aspect(1)\n",
    "ax.set_axis_off()\n",
    "ax.set_xlim(-10,5)\n",
    "ax.set_ylim(-10,5);\n",
    "# plt.savefig('MB_Familiar_Map_l7.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _visual homing results_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 82132 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,ax = plt.subplots(figsize=(8,8))\n",
    "# get the data generated from the simulation\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'VH_ONLY.mat')\n",
    "pos = data['pos']\n",
    "h = data['h']\n",
    "\n",
    "# the index of the homing route to plot\n",
    "plot_ind = [8]\n",
    "# the indexes of the agents to caculate the initial heading\n",
    "caculate_ind = [0,1,2,3,5,6,8,10,11]\n",
    "\n",
    "# plot the homing route\n",
    "for num,i in enumerate(plot_ind):\n",
    "    if num == 0:\n",
    "        ax.plot(pos[i][:,0]/100,pos[i][:,1]/100,color='r',lw=1.5,alpha = (num+1)*(1/len(plot_ind)),label='Homing Routes')\n",
    "    ax.plot(pos[i][:,0]/100,pos[i][:,1]/100,color='r',lw=1.5,alpha = (num+1)*(1/len(plot_ind)))\n",
    "    ax.quiver(pos[i][0,0]/100,pos[i][0,1]/100,np.cos(h[i][0]),np.sin(h[i][0]),\n",
    "              color='r',scale=15, width=0.004, alpha = (num+1)*(1/len(plot_ind)))\n",
    "\n",
    "# analysis the initial headings\n",
    "import circular_statistics as cs\n",
    "# simulated ant\n",
    "cp_ind = cs.get_check_points_index(250, pos)\n",
    "cp_h = [h[i,cp_ind[i]] for i in caculate_ind]\n",
    "acc = 360 # 5deg goniometer\n",
    "d_ref = np.linspace(0,2*np.pi,acc)\n",
    "h_r, h_mean, h_sd0, h_ci95, r_num = cs.circular_statistics(cp_h,acc=acc)\n",
    "# add the polar plot\n",
    "insert_ax = fig.add_axes([0.2,0.6,0.2,0.2],projection='polar')\n",
    "insert_ax.quiver(0,0,h_r*np.cos(h_mean),h_r*np.sin(h_mean),color='r', scale=0.5,scale_units='y',width=0.025,label='Model')\n",
    "insert_ax.plot(np.linspace(h_mean-h_ci95,h_mean+h_ci95,20),np.repeat(h_r*4,20),color='r',lw=5)\n",
    "insert_ax.scatter(h_mean-h_ci95,h_r*4,s=50,marker='.',color='r',edgecolor='r')\n",
    "insert_ax.scatter(h_mean+h_ci95,h_r*4,s=50,marker='.',color='r',edgecolor='r')\n",
    "# real ant\n",
    "real_ant_data = sio.loadmat(r_folder + r_sub_folder + 'Wystrach2012_OnOffRoute.mat')\n",
    "real_cp_h = real_ant_data['head_N1_RX_P'][0]\n",
    "h_r, h_mean, h_sd0, h_ci95, r_num = cs.circular_statistics(real_cp_h,acc=acc)\n",
    "insert_ax.quiver(0,0,h_r*np.cos(h_mean),h_r*np.sin(h_mean),color='k', scale=0.5,scale_units='y',width=0.025,label='Real Ant')\n",
    "insert_ax.plot(np.linspace(h_mean-h_ci95,h_mean+h_ci95,20),np.repeat(h_r*4,20),color='k',lw=5)\n",
    "insert_ax.scatter(h_mean-h_ci95,h_r*4,s=50,marker='.',color='k',edgecolor='k')\n",
    "insert_ax.scatter(h_mean+h_ci95,h_r*4,s=50,marker='.',color='k',edgecolor='k')\n",
    "\n",
    "# habitual route\n",
    "ax.plot(route_memory['pos'][:,0]/100,route_memory['pos'][:,1]/100,color='gray',lw=10,alpha=0.6,label='Habitual Route')\n",
    "interval = 4\n",
    "ax.quiver(route_memory['pos'][::interval,0]/100,route_memory['pos'][::interval,1]/100, \n",
    "          np.cos(route_memory['h'][0][::interval]),np.sin(route_memory['h'][0][::interval]),\n",
    "          width=0.005,scale=30,color='k')\n",
    "\n",
    "# RP\n",
    "ax.text((pos[i][0,0]+20)/100,(pos[i][0,1]-20)/100,'RP',fontsize=fontsize,color='k')\n",
    "ax.scatter(pos[i][0,0]/100,pos[i][0,1]/100, color='k',s=100)\n",
    "\n",
    "# Nest\n",
    "ax.scatter(0,0,color='r',marker='*',s=300)\n",
    "ax.text(0.2,-0.2,'Nest',fontsize=fontsize,color='red')\n",
    "\n",
    "# Feeder\n",
    "ax.scatter(-7,-7,color='k',marker='s',s=300)\n",
    "ax.text(-7.2,-8,'Feeder',fontsize=30,color='k')\n",
    "\n",
    "# format adjust\n",
    "insert_ax.set_yticklabels('')\n",
    "ticks = insert_ax.get_yticklabels()\n",
    "[tick.set_fontsize(fontsize-5) for tick in ticks]\n",
    "ticks = insert_ax.get_xticklabels()\n",
    "[tick.set_fontsize(fontsize-5) for tick in ticks]\n",
    "insert_ax.legend(fontsize=fontsize-8)\n",
    "\n",
    "ticks = ax.get_yticklabels()\n",
    "[tick.set_fontsize(fontsize) for tick in ticks]\n",
    "ticks = ax.get_xticklabels()\n",
    "[tick.set_fontsize(fontsize) for tick in ticks]\n",
    "ax.set_title('Visual Homing',color='k',fontsize=fontsize)\n",
    "ax.set_xlabel('X / m',fontsize=fontsize)\n",
    "ax.set_ylabel('Y / m',fontsize=fontsize)\n",
    "ax.legend(fontsize=fontsize-5)\n",
    "ax.set_xlim(-10,5)\n",
    "ax.set_ylim(-10,5)\n",
    "ax.grid(1)\n",
    "ax.set_aspect(1)\n",
    "# plt.savefig('viusal_homing_results.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='gray'> optimal integration </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_sub_folder = 'optimal_integration/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _results of tuning PI uncertainty_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65300 characters]\n",
      "text/plain": [
       "<Figure size 432x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# initial heading analysis of the agents\n",
    "import circular_statistics as cs\n",
    "\n",
    "PI_l = [10,100,300,700]\n",
    "pos = []\n",
    "h = []\n",
    "# load data\n",
    "for i in range(4):\n",
    "    data = sio.loadmat(r_folder + r_sub_folder + 'PI' + str(PI_l[i]) + '_VH_SP5.mat')\n",
    "    pos.append(data['pos'])\n",
    "    h.append(data['h'])\n",
    "# calculate the initial headings\n",
    "ck_ind = []\n",
    "ck_h = []\n",
    "for j in range(4):\n",
    "    ck_ind.append(cs.get_check_points_index(60,pos[j]))\n",
    "    ck_ind[j] = list(filter(lambda x: x!=0, ck_ind[j]))\n",
    "    ck_h.append([h[j][i][ck_ind[j][i]] for i in range(len(ck_ind[j]))])\n",
    "# calculate the circular statistics\n",
    "r = []\n",
    "mean = []\n",
    "ci_95 = []\n",
    "num = []\n",
    "for i in range(4):\n",
    "    r_i, mean_i, sd0_i, ci95_i, num_i = cs.circular_statistics(ck_h[i],acc=100)\n",
    "    r.append(r_i)\n",
    "    mean.append(mean_i)\n",
    "    ci_95.append(ci95_i)\n",
    "    num.append(num_i)\n",
    "\n",
    "# plot the circular data\n",
    "fig = plt.figure(figsize=(6,6))\n",
    "fontsize=30\n",
    "ax = fig.add_subplot(111,projection='polar')\n",
    "color_list = ['gray','gray','gray','gray']\n",
    "for i in range(4):\n",
    "    ax.quiver(0,0,r[i]*np.cos(mean[i]),r[i]*np.sin(mean[i]),color=color_list[i], width=0.012, scale=0.5,scale_units='y')\n",
    "    ax.plot(np.linspace(mean[i]-ci_95[i],mean[i]+ci_95[i],20),np.repeat(4 * r[i],20),color=color_list[i],lw=4)\n",
    "    ax.scatter(mean[i]-ci_95[i],4 * r[i],s=220,marker='.',color='gray',edgecolor=color_list[i])\n",
    "    ax.scatter(mean[i]+ci_95[i],4 * r[i],s=220,marker='.',color='gray',edgecolor=color_list[i])\n",
    "    ax.text(mean[i]+ci_95[i],4 * r[i],'%sm'%(PI_l[i]/100),fontsize=24,color='orange')\n",
    "# the PI and Visual Diretion\n",
    "pi_d = 3*np.pi/2\n",
    "pi_v = np.deg2rad(140)\n",
    "\n",
    "tri_len = 0.8\n",
    "tri_pos_l = 4.9\n",
    "tri_pi = plt.Polygon([[pi_d,tri_pos_l],\n",
    "                      [pi_d+np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)],\n",
    "                      [pi_d-np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)]],\n",
    "                    color='orange')\n",
    "ax.add_patch(tri_pi)\n",
    "tri_pi = plt.Polygon([[pi_v,tri_pos_l],\n",
    "                      [pi_v+np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)],\n",
    "                      [pi_v-np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)]],\n",
    "                    color='r')\n",
    "ax.add_patch(tri_pi)\n",
    "ax.set_ylim(0,5)\n",
    "\n",
    "ax.set_yticklabels('')\n",
    "ticks = ax.get_xticklabels()\n",
    "[tick.set_fontsize(fontsize-5) for tick in ticks];\n",
    "\n",
    "# plt.savefig('optimal_tune_pi.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _results of tuning VH uncertainty_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69364 characters]\n",
      "text/plain": [
       "<Figure size 432x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# combining the tuning visual experiments\n",
    "tune_visual_r = [0,0,0]\n",
    "tune_visual_mean = [0,0,0]\n",
    "tune_visual_ci95 = [0,0,0]\n",
    "tune_visual_num = [0,0,0]\n",
    "\n",
    "tune_visual_vf = [0,0,0]\n",
    "tune_visual_dis = [0,0,0]\n",
    "\n",
    "# RP1\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'PI100_VH_SP5.mat')\n",
    "pos = data['pos']\n",
    "h = data['h']\n",
    "tune_visual_dis[0] = np.sqrt(np.sum(pos[0,0]**2 + pos[0,1]**2))/100.0\n",
    "ck_h = cs.get_check_points_h(60,pos,h)\n",
    "tune_visual_r[0], tune_visual_mean[0], x, tune_visual_ci95[0], tune_visual_num[0] = cs.circular_statistics(ck_h,acc=100)\n",
    "\n",
    "# RP2\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'PI_VH_SP15.mat')\n",
    "pos = data['pos']\n",
    "h = data['h']\n",
    "tune_visual_dis[1] = np.sqrt(np.sum(pos[0,0]**2 + pos[0,1]**2))/100.0\n",
    "ck_h = cs.get_check_points_h(30,pos,h)\n",
    "tune_visual_r[1], tune_visual_mean[1], x, tune_visual_ci95[1], tune_visual_num[1] = cs.circular_statistics(ck_h,acc=100)\n",
    "\n",
    "# RP3\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'PI_VH_SP25.mat')\n",
    "pos = data['pos']\n",
    "h = data['h']\n",
    "tune_visual_dis[2] = np.sqrt(np.sum(pos[0,0]**2 + pos[0,1]**2))/100.0\n",
    "ck_h = cs.get_check_points_h(30,pos,h)\n",
    "tune_visual_r[2], tune_visual_mean[2], x, tune_visual_ci95[2], tune_visual_num[2] = cs.circular_statistics(ck_h,acc=100)\n",
    "\n",
    "# plot the data\n",
    "fig = plt.figure(figsize=(6,6))\n",
    "fontsize=30\n",
    "ax = fig.add_subplot(111,projection='polar')\n",
    "\n",
    "for i in range(3):\n",
    "    ax.quiver(0,0,tune_visual_r[i]*np.cos(tune_visual_mean[i]),tune_visual_r[i]*np.sin(tune_visual_mean[i]),color='gray', \n",
    "              width=0.012, scale=0.5,scale_units='y')\n",
    "    ax.plot(np.linspace(tune_visual_mean[i]-tune_visual_ci95[i],tune_visual_mean[i]+tune_visual_ci95[i],20),\n",
    "            np.repeat(4 * tune_visual_r[i],20),color='gray',lw=4)\n",
    "    ax.scatter(tune_visual_mean[i]-tune_visual_ci95[i],4 * tune_visual_r[i],s=220,marker='.',color='gray',edgecolor='gray')\n",
    "    ax.scatter(tune_visual_mean[i]+tune_visual_ci95[i],4 * tune_visual_r[i],s=220,marker='.',color='gray',edgecolor='gray')\n",
    "    ax.text(tune_visual_mean[i]+tune_visual_ci95[i],4 * tune_visual_r[i],'%.4sm'%(tune_visual_dis[i]),fontsize=24,color='r')\n",
    "\n",
    "pi_d = 3*np.pi/2\n",
    "pi_v = np.deg2rad(140)\n",
    "\n",
    "tri_len = 0.6\n",
    "tri_pos_l = 2.9\n",
    "tri_pi = plt.Polygon([[pi_d,tri_pos_l],\n",
    "                      [pi_d+np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)],\n",
    "                      [pi_d-np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)]],\n",
    "                    color='orange')\n",
    "ax.add_patch(tri_pi)\n",
    "tri_pi = plt.Polygon([[pi_v,tri_pos_l],\n",
    "                      [pi_v+np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)],\n",
    "                      [pi_v-np.arctan2((tri_len/2),tri_pos_l-tri_len*np.cos(np.pi/6)),np.sqrt((tri_len/2)**2+(tri_pos_l-tri_len*np.cos(np.pi/6))**2)]],\n",
    "                    color='r')\n",
    "ax.add_patch(tri_pi)\n",
    "\n",
    "ax.set_ylim(0,3)\n",
    "ax.set_yticklabels('')\n",
    "ticks = ax.get_xticklabels()\n",
    "[tick.set_fontsize(fontsize-5) for tick in ticks];\n",
    "\n",
    "# plt.savefig('optimal_tune_v.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _typical homing path_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x11501b0>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 39972 characters]\n",
      "text/plain": [
       "<Figure size 432x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# pi_route\n",
    "pi_route_mem_data = sio.loadmat(d_folder + 'PiRouteMemory.mat')\n",
    "\n",
    "\n",
    "fontsize = 30\n",
    "fig,ax = plt.subplots(figsize=(6,8))\n",
    "\n",
    "# pi\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'PI_Route.mat')\n",
    "homing_pos = data['pos'][:500]\n",
    "ax.plot(homing_pos[:,0]/100,homing_pos[:,1]/100,color='orange',lw=2,label='PI')\n",
    "# vh\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'VH_Route.mat')\n",
    "homing_pos = data['pos']\n",
    "ax.plot(homing_pos[:,0]/100,homing_pos[:,1]/100,color='r',lw=2,label='VH')\n",
    "# pi+vh\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'PI_VH_Route.mat')\n",
    "homing_pos = data['pos']\n",
    "ax.plot(homing_pos[:,0]/100,homing_pos[:,1]/100,color='gray',lw=2,label='PI+VH')\n",
    "\n",
    "ax.plot(pi_route_mem_data['pos'][:,0]/100,pi_route_mem_data['pos'][:,1]/100,color='gray',lw=10,alpha=0.6)\n",
    "interval = 4\n",
    "ax.quiver(pi_route_mem_data['pos'][::interval,0]/100,pi_route_mem_data['pos'][::interval,1]/100, \n",
    "          np.cos(pi_route_mem_data['h'][0][::interval]),np.sin(pi_route_mem_data['h'][0][::interval]),\n",
    "          width=0.01,scale=20,color='k')\n",
    "# nest\n",
    "ax.scatter(0,0,color='r',marker='*',s=300)\n",
    "ax.text(0.2,-0.2,'Nest',fontsize=fontsize-10,color='red')\n",
    "\n",
    "# feeder\n",
    "ax.scatter(pi_route_mem_data['pos'][-1,0]/100,pi_route_mem_data['pos'][-1,1]/100,color='k',marker='s',s=300)\n",
    "ax.text(pi_route_mem_data['pos'][-1,0]/100+0.2,pi_route_mem_data['pos'][-1,1]/100+0.2,'Feeder',fontsize=fontsize-10,color='k')\n",
    "\n",
    "ax.scatter(homing_pos[0,0]/100,homing_pos[0,1]/100,color='k',marker='.',s=200)\n",
    "ax.text(homing_pos[0,0]/100,homing_pos[0,1]/100,'RP',fontsize=fontsize-5)\n",
    "ax.set_xlim(-5,4)\n",
    "# ax.set_ylim(-3,13)\n",
    "ax.set_xlabel('X / m',fontsize=fontsize-15)\n",
    "ax.set_ylabel('Y / m',fontsize=fontsize-15)\n",
    "plt.xticks(fontsize=fontsize-15)\n",
    "plt.yticks(fontsize=fontsize-15)\n",
    "ax.grid(1)\n",
    "ax.set_aspect(1)\n",
    "ax.legend(fontsize=fontsize-10)\n",
    "# plt.savefig('pi_vh_route.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _extented homing path: supplementary figure_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 114604 characters]\n",
      "text/plain": [
       "<Figure size 1728x432 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot the extended path of the PI+VH - varies PI length\n",
    "fig,ax = plt.subplots(1,4,figsize=(24,6))\n",
    "PI_length = [10,100,300,700]\n",
    "for i in range(4):\n",
    "    data = sio.loadmat(r_folder + r_sub_folder + 'PI' + str(PI_l[i]) + '_VH_SP5.mat')\n",
    "    pos = data['pos']\n",
    "    h = data['h']\n",
    "    ax[i].scatter(pos[0,0,0]/100,pos[0,0,1]/100,color='k',marker='s',s=100)\n",
    "    ax[i].scatter(0,0,color='r',marker='*',s=300)\n",
    "\n",
    "    for j in range(len(pos)):\n",
    "        ax[i].plot(pos[j,:,0]/100,pos[j,:,1]/100,color='k',lw=0.5)\n",
    "    ax[i].set_xlim([-1,2.5])\n",
    "    ax[i].set_ylim([-2.5,1])\n",
    "    ax[i].set_aspect(1)\n",
    "    ax[i].grid(1)\n",
    "    ax[i].set_xlabel('X / m',fontsize=fontsize-10)\n",
    "    ax[i].set_ylabel('Y / m',fontsize=fontsize-10)\n",
    "    ax[i].set_xticks([-1,0,1,2])\n",
    "    ax[i].set_yticks([-2,-1,0,1])\n",
    "    ax[i].set_title('home vector length = %sm'%(PI_length[i]/100),fontsize=fontsize-10)\n",
    "    ticks = ax[i].get_yticklabels()\n",
    "    [tick.set_fontsize(fontsize-10) for tick in ticks]\n",
    "    ticks = ax[i].get_xticklabels()\n",
    "    [tick.set_fontsize(fontsize-10) for tick in ticks]\n",
    "# plt.savefig('supp_optimal_path1.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 41232 characters]\n",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot the extended path of the PI+VH - varies VISUAL\n",
    "data_rp1 = sio.loadmat(r_folder + r_sub_folder + 'PI100_VH_SP5.mat')\n",
    "pos_rp1 = data_rp1['pos']\n",
    "h_rp1 = data_rp1['h']\n",
    "\n",
    "data_rp2 = sio.loadmat(r_folder + r_sub_folder + 'PI_VH_SP15.mat')\n",
    "pos_rp2 = data_rp2['pos']\n",
    "h_rp2 = data_rp2['h']\n",
    "\n",
    "data_rp3 = sio.loadmat(r_folder + r_sub_folder + 'PI_VH_SP25.mat')\n",
    "pos_rp3 = data_rp3['pos']\n",
    "h_rp3 = data_rp3['h']\n",
    "\n",
    "fig,ax = plt.subplots(1,2,figsize=(12,6))\n",
    "for i in range(len(h_rp2)):\n",
    "    ax[0].plot(pos_rp2[i,:,0]/100,pos_rp2[i,:,1]/100,color='k',lw=0.5)\n",
    "    ax[1].plot(pos_rp3[i,:,0]/100,pos_rp3[i,:,1]/100,color='k',lw=0.5)\n",
    "\n",
    "ax[0].scatter(pos_rp2[0,0,0]/100,pos_rp2[0,0,1]/100,color='k',marker='s',s=100)\n",
    "ax[1].scatter(pos_rp3[0,0,0]/100,pos_rp3[0,0,1]/100,color='k',marker='s',s=100)\n",
    "\n",
    "rp_dis = [6.18,10.3]\n",
    "for i in [0,1]:\n",
    "    ax[i].scatter(0,0,color='r',marker='*',s=300)\n",
    "    ax[i].set_title('RP distance = %2.4sm'%(rp_dis[i]),fontsize=fontsize-10)\n",
    "    ax[i].set_xlim([-1,10])\n",
    "    ax[i].set_ylim([-6,1])\n",
    "    ax[i].set_aspect(1)\n",
    "    ax[i].grid(1)\n",
    "    ax[i].set_xlabel('X / m',fontsize=fontsize-10)\n",
    "    ax[i].set_ylabel('Y / m',fontsize=fontsize-10)\n",
    "    ticks = ax[i].get_yticklabels()\n",
    "    [tick.set_fontsize(fontsize-10) for tick in ticks]\n",
    "    ticks = ax[i].get_xticklabels()\n",
    "    [tick.set_fontsize(fontsize-10) for tick in ticks]\n",
    "# plt.savefig('supp_optimal_path2.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <font color='blue'> route following (RF) </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_sub_folder = 'route_following/'\n",
    "\n",
    "# the route memory\n",
    "route_memory = sio.loadmat(d_folder + 'ArcRouteMemory.mat')\n",
    "# the home memory\n",
    "home_memory = sio.loadmat(d_folder + 'HomeMemory_X0Y0.mat')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _ZM phase tracking performances_ </font>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <font color='black'> _single point ZM tracking_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from image_processing import visual_sense\n",
    "# bad(-2.5,-3.5) /good(-8,2) tracking\n",
    "h_s = [np.arange(s_h,s_h+np.pi*2,19/180*np.pi) for s_h in [0, np.pi/2, np.pi, 3*np.pi/2.0]]\n",
    "zm_a_s = []\n",
    "zm_p_s = []\n",
    "for h_ii in h_s:\n",
    "    for h_i in h_ii:\n",
    "        # good tracking\n",
    "#         zm_a, zm_p = visual_sense(world, -8, 2, h_i, nmax=zm_n_max)\n",
    "        # bad tracking\n",
    "        zm_a, zm_p = visual_sense(world, -2.5, -3.5, h_i, nmax=zm_n_max)\n",
    "        zm_a_s.append(zm_a)\n",
    "        zm_p_s.append(zm_p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 42120 characters]\n",
      "text/plain": [
       "<Figure size 576x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot bad(-2.5,-3.5) / good(1.5,-3) tracking\n",
    "fig,ax = plt.subplots(figsize=(8,4))\n",
    "ax.set_xlabel('Heading / rad',fontsize=30)\n",
    "ax.set_ylabel('ZM-phase / rad',fontsize=30)\n",
    "zm_phase = [zm_p_s[i][16] for i in range(len(zm_p_s))]\n",
    "ax.set_ylim(-np.pi,np.pi)\n",
    "ax.set_xlim(-np.pi,np.pi)\n",
    "for i in range(2,3):\n",
    "    perfect_h = zm_phase[i*19]/180*np.pi - h_s[i]\n",
    "    perfect_h = ((perfect_h) + np.pi)%(np.pi*2.0) - np.pi\n",
    "    ax.scatter((h_s[i] + np.pi) % (2.0 * np.pi)-np.pi,perfect_h,color='r',label='Animal orientation',s=60,alpha=0.8)\n",
    "    ax.plot((h_s[i] + np.pi) % (2.0 * np.pi)-np.pi,perfect_h,color='r')\n",
    "    ax.scatter((h_s[i] + np.pi) % (2.0 * np.pi) - np.pi,np.array(zm_phase[i*19:(i+1)*19])/180*np.pi,color='b',\n",
    "               label='ZM-Phase tracking',s=60,alpha=0.8)\n",
    "    ax.plot((h_s[i] + np.pi) % (2.0 * np.pi) - np.pi,np.array(zm_phase[i*19:(i+1)*19])/180*np.pi,color='b')\n",
    "ax.grid(1)\n",
    "ax.legend(fontsize=20)\n",
    "plt.xticks(fontsize=30)\n",
    "plt.yticks(fontsize=30);\n",
    "# plt.savefig('zm_p_tracking_bad.pdf')\n",
    "# plt.savefig('zm_p_tracking_good.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <font color='black'> _cross world ZM tracking_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data from the simulation\n",
    "rf_data = sio.loadmat(r_folder + r_sub_folder + 'QuiverPlotData_X-10_2_Y-8_2_SH20.mat')\n",
    "\n",
    "# sampled num in x-y direction\n",
    "sample_num = 20\n",
    "# sampled locations\n",
    "pos_x = np.linspace(-10,2,sample_num)\n",
    "pos_y = np.linspace(-8,2,sample_num)\n",
    "\n",
    "# sampled headings\n",
    "h = np.linspace(-np.pi,np.pi,10)\n",
    "\n",
    "# check the data for RF\n",
    "# 1.RF memory , 2.the phase-tracking, 3.RF suggested\n",
    "\n",
    "# current zm p\n",
    "current_zm_p = rf_data['current_zm_p']\n",
    "current_zm_p_rad = current_zm_p/180.0 *np.pi\n",
    "diff_zm_p = (np.diff(current_zm_p_rad) + np.pi) %(np.pi*2) - np.pi\n",
    "for p in diff_zm_p:\n",
    "    p[np.where(p>0)] = 0\n",
    "zm_p_ratio = diff_zm_p/(h[1]-h[0])\n",
    "# inverse the ratio<-1.0\n",
    "for p in zm_p_ratio:\n",
    "    p[np.where(p<-1.0)] = 1.0/p[np.where(p<-1.0)]\n",
    "zm_p_ratio_mean = np.mean(zm_p_ratio,axis=1)\n",
    "\n",
    "import circular_statistics as cs\n",
    "# ann output\n",
    "ann_output = rf_data['ann_output']\n",
    "ann_mean = list(map(cs.circular_statistics, ann_output))\n",
    "ann_cir_statistic = {'r':[ann_mean[i][0] for i in range(sample_num**2)],'mean':[ann_mean[i][1] for i in range(sample_num**2)],\n",
    "                     'ci95':[ann_mean[i][3] for i in range(sample_num**2)]}\n",
    "# RF angle \n",
    "rf_angle = list(map(cs.calculate_rf_motor, ann_output, np.tile(h,[sample_num**2,1]), current_zm_p_rad))\n",
    "rf_mean = list(map(cs.circular_statistics,rf_angle))\n",
    "rf_cir_statistic = {'r':[rf_mean[i][0] for i in range(sample_num**2)],'mean':[rf_mean[i][1] for i in range(sample_num**2)],\n",
    "                    'ci95':[rf_mean[i][3] for i in range(sample_num**2)]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 111344 characters]\n",
      "text/plain": [
       "<Figure size 576x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot\n",
    "fig,ax = plt.subplots(figsize=(8,10))\n",
    "\n",
    "h_plot_s = (np.array(rf_cir_statistic['mean']) - np.array(rf_cir_statistic['ci95'])).reshape(sample_num,sample_num) + np.pi\n",
    "h_plot_e = (np.array(rf_cir_statistic['mean']) + np.array(rf_cir_statistic['ci95'])).reshape(sample_num,sample_num) + np.pi\n",
    "\n",
    "h_plot_s = np.rad2deg(h_plot_s)\n",
    "h_plot_e = np.rad2deg(h_plot_e)\n",
    "\n",
    "wedges = []\n",
    "for i in range(sample_num):\n",
    "    for j in range(sample_num):\n",
    "        wedge = mpatches.Wedge([pos_x[j],pos_y[i]], 0.3, h_plot_s[i,j], h_plot_e[i,j], ec=\"none\")\n",
    "        wedges.append(wedge)\n",
    "\n",
    "collection = PatchCollection(wedges,  color='blue', alpha=0.3)\n",
    "ax.add_collection(collection)\n",
    "\n",
    "h_plot = np.array(rf_cir_statistic['mean']).reshape(sample_num,sample_num) + np.pi\n",
    "r_plot = np.array(rf_cir_statistic['r']).reshape(sample_num,sample_num)\n",
    "\n",
    "# sampled positions\n",
    "X,Y = np.meshgrid(pos_x,pos_y)\n",
    "ax.scatter(X,Y,color='gray',marker='o',s=10,alpha=0.5)\n",
    "\n",
    "ax.set_title('Route Following',fontsize=30)\n",
    "plt.xticks(fontsize=30)\n",
    "plt.yticks(fontsize=30)\n",
    "ax.set_aspect(1)\n",
    "# plt.savefig('Whole_RouteFollowingResults.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _results and homing path_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 143408 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,ax=plt.subplots(figsize=(12,12))\n",
    "insert_ax = fig.add_axes([0.2,0.6,0.2,0.2],projection='polar')\n",
    "\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'RF_ONLY.mat')\n",
    "pos = data['pos']\n",
    "h = data['h']\n",
    "# plot will end at the points nearest to the nest for the all homing path\n",
    "dis = [pos[i][:,0]**2 + pos[i][:,1]**2 for i in range(len(pos))]\n",
    "end_t = [np.argmin(dis[i])+1 for i in range(len(pos))]\n",
    "\n",
    "# plot homing route\n",
    "for i in range(len(pos)):   \n",
    "#     ax.scatter(InsectNaviAgent.homing_pos[mask][:,0]/100,InsectNaviAgent.homing_pos[mask][:,1]/100,c='red')\n",
    "    ax.plot(pos[i][:end_t[i]-1,0]/100,pos[i][:end_t[i]-1,1]/100,lw=1.5,color='blue')\n",
    "    ax.scatter(pos[i][0,0]/100,pos[i][0,1]/100,marker='o',color='k',s=100,alpha=0.5)\n",
    "\n",
    "# initial heading analysis of the agents released on route\n",
    "import circular_statistics as cs\n",
    "\n",
    "## simualted ant\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'RF_released_on_route.mat')\n",
    "pos = data['pos']\n",
    "h = data['h']\n",
    "ck_h = cs.get_check_points_h(60,pos,h)\n",
    "r_i, mean_i, sd0_i, ci95_i, num_i = cs.circular_statistics(ck_h,acc=100)\n",
    "insert_ax.quiver(0,0,r_i*np.cos(mean_i),r_i*np.sin(mean_i),color='blue', width=0.012, scale=0.5,scale_units='y',label='Model')\n",
    "insert_ax.plot(np.linspace(mean_i-ci95_i,mean_i+ci95_i,20),np.repeat(4 * r_i,20),color='blue',lw=5,alpha=0.6)\n",
    "insert_ax.scatter(mean_i-ci95_i,4 * r_i,s=50,marker='.',color='blue',edgecolor='blue',alpha=0.6)\n",
    "insert_ax.scatter(mean_i+ci95_i,4 * r_i,s=50,marker='.',color='blue',edgecolor='blue',alpha=0.6)\n",
    "\n",
    "## real ant\n",
    "real_ant_data = sio.loadmat(r_folder + r_sub_folder + 'Wystrach2012_OnOffRoute.mat')\n",
    "real_cp_h = real_ant_data['head_N1_RX_R'][0]\n",
    "h_r, h_mean, h_sd0, h_ci95, r_num = cs.circular_statistics(real_cp_h,acc=acc)\n",
    "insert_ax.quiver(0,0,h_r*np.cos(h_mean),h_r*np.sin(h_mean),color='r', scale=0.5,scale_units='y',width=0.025,label='Real Ant')\n",
    "insert_ax.plot(np.linspace(h_mean-h_ci95,h_mean+h_ci95,20),np.repeat(h_r*4,20),color='r',lw=5)\n",
    "insert_ax.scatter(h_mean-h_ci95,h_r*4,s=50,marker='.',color='r',edgecolor='r')\n",
    "insert_ax.scatter(h_mean+h_ci95,h_r*4,s=50,marker='.',color='r',edgecolor='r')\n",
    "\n",
    "# habitual route\n",
    "ax.plot(route_memory['pos'][:,0]/100,route_memory['pos'][:,1]/100,color='gray',lw=10,alpha=0.6)\n",
    "interval = 4\n",
    "ax.quiver(route_memory['pos'][::interval,0]/100,route_memory['pos'][::interval,1]/100, \n",
    "          np.cos(route_memory['h'][0][::interval]),np.sin(route_memory['h'][0][::interval]),\n",
    "          width=0.005,scale=30,color='k')\n",
    "\n",
    "# nest\n",
    "ax.scatter(0,0,color='r',marker='*',s=300)\n",
    "ax.text(0.5,-0.2,'Nest',fontsize=30,color='red')\n",
    "\n",
    "# RPs\n",
    "ax.text(-8,-7.5,'Release pionts',fontsize=30,color='k')\n",
    "\n",
    "plt.xticks(fontsize=30)\n",
    "plt.yticks(fontsize=30)\n",
    "# format adjust\n",
    "insert_ax.set_yticklabels('')\n",
    "ticks = insert_ax.get_yticklabels()\n",
    "[tick.set_fontsize(fontsize-5) for tick in ticks]\n",
    "ticks = insert_ax.get_xticklabels()\n",
    "[tick.set_fontsize(fontsize-5) for tick in ticks]\n",
    "insert_ax.legend(fontsize=fontsize-8)\n",
    "\n",
    "ax.grid(1)\n",
    "ax.set_xlabel('X / m',fontsize=30)\n",
    "ax.set_ylabel('Y / m',fontsize=30)\n",
    "ax.set_aspect(1)\n",
    "# plt.savefig('route_following_results.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <font color='black'> whole model </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_sub_folder = 'whole_model/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _homing path of ZV and FV_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x22a5d470>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 40044 characters]\n",
      "text/plain": [
       "<Figure size 576x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# homing path of the whole model\n",
    "route_mem_data = sio.loadmat(d_folder + 'ArcRouteMemory.mat')\n",
    "fig,ax = plt.subplots(figsize=(8,10))\n",
    "\n",
    "# nest\n",
    "ax.scatter(route_mem_data['pos'][-1,0]/100,route_mem_data['pos'][-1,1]/100,color='r',marker='*',s=300)\n",
    "ax.text(route_mem_data['pos'][-1,0]/100,route_mem_data['pos'][-1,1]/100,'Nest',fontsize=fontsize-10,color='red')\n",
    "\n",
    "# feeder\n",
    "ax.scatter(route_mem_data['pos'][1,0]/100,route_mem_data['pos'][1,1]/100,color='k',marker='s',s=300)\n",
    "ax.text(route_mem_data['pos'][1,0]/100+0.2,route_mem_data['pos'][1,1]/100+0.2,'Feeder',fontsize=fontsize-10,color='k')\n",
    "\n",
    "# route\n",
    "ax.plot(route_mem_data['pos'][:,0]/100,route_mem_data['pos'][:,1]/100,color='gray',lw=10,alpha=0.6)\n",
    "interval = 4\n",
    "ax.quiver(route_mem_data['pos'][::interval,0]/100,route_mem_data['pos'][::interval,1]/100, \n",
    "          np.cos(route_mem_data['h'][0][::interval]),np.sin(route_mem_data['h'][0][::interval]),\n",
    "          width=0.01,scale=25,color='k',alpha=0.5)\n",
    "\n",
    "# ZV\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'final_ZV.mat')\n",
    "homing_pos = data['homing_pos'][:180]\n",
    "ax.plot(homing_pos[:,0]/100,homing_pos[:,1]/100,color='k',lw=1.5,ls='--',label='ZV')\n",
    "ax.scatter(homing_pos[96,0]/100,homing_pos[96,1]/100,color='orange',marker='D',s=100)\n",
    "\n",
    "# FV\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'final_FV.mat')\n",
    "homing_pos = data['homing_pos']\n",
    "ax.plot(homing_pos[:,0]/100,homing_pos[:,1]/100,color='k',lw=1.5,ls='-',label='FV')\n",
    "ax.scatter(homing_pos[78,0]/100,homing_pos[78,1]/100,color='orange',marker='D',s=100)\n",
    "\n",
    "# RP\n",
    "ax.scatter(homing_pos[0,0]/100,homing_pos[0,1]/100,color='k',marker='.',s=200)\n",
    "ax.text(homing_pos[0,0]/100,homing_pos[0,1]/100,'RP',fontsize=fontsize-5)\n",
    "\n",
    "ax.set_xlim(-8,3)\n",
    "ax.set_ylim(-8,2)\n",
    "ax.set_xlabel('X / m',fontsize=fontsize-15)\n",
    "ax.set_ylabel('Y / m',fontsize=fontsize-15)\n",
    "plt.xticks(fontsize=fontsize-15)\n",
    "plt.yticks(fontsize=fontsize-15)\n",
    "ax.grid(1)\n",
    "ax.set_aspect(1)\n",
    "ax.legend(fontsize=fontsize-15)\n",
    "\n",
    "# plt.savefig('whole_homing_path.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='black'> _activation of TUN, SN1 and SN2_ </font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36984 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# ZV\n",
    "# data = sio.loadmat(r_folder + r_sub_folder + 'final_ZV.mat')\n",
    "\n",
    "# FV\n",
    "data = sio.loadmat(r_folder + r_sub_folder + 'final_FV.mat')\n",
    "\n",
    "fig,ax = plt.subplots(3,1,figsize=(8,6))\n",
    "\n",
    "end_t = 160\n",
    "\n",
    "SN1 = data['homing_state'][0][:end_t]\n",
    "\n",
    "# SN1\n",
    "ax[0].plot(SN1, c='r',lw=2)\n",
    "# ax[0].set_xticks([])\n",
    "ax[0].text(10,0.75,'SN1',fontsize=fontsize-10,color='r')\n",
    "\n",
    "# SN2\n",
    "ax[1].plot(1-SN1,c='cyan',lw=2)\n",
    "# ax[1].set_xticks([])\n",
    "ax[1].text(10,0.75,'SN2',fontsize=fontsize-10,color='cyan')\n",
    "\n",
    "# TUN\n",
    "ax[2].plot(data['homing_vh_sensory'][0][:end_t] / 80, c='blue',lw=2)\n",
    "ax[2].text(10,0.75,'TUN',fontsize=fontsize-10,color='blue')\n",
    "ax[2].set_xlabel('Time step',fontsize=fontsize)\n",
    "ax[2].set_ylabel('Activation',fontsize=fontsize)\n",
    "\n",
    "for a in ax:\n",
    "    # dimond marker\n",
    "    a.scatter(78,1,c='orange',s=100,marker='D')\n",
    "    a.set_xlim([0,end_t+1])\n",
    "    a.set_ylim([-0.1,1.1])\n",
    "    a.grid(1)\n",
    "    ticks = a.get_yticklabels()\n",
    "    [tick.set_fontsize(fontsize-10) for tick in ticks]\n",
    "    ticks = a.get_xticklabels()\n",
    "    [tick.set_fontsize(fontsize-10) for tick in ticks]\n",
    "\n",
    "# plt.savefig('SN1_SN2_TUN_ZV.pdf')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
