{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <center>Simulations of Insect Navigation Model</center>\n",
    "\n",
    "##### <center><font color='gray'>_Xuelong Sun_, _2020-04-01_</font></center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Dependance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Packages and settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.io as sio\n",
    "import numpy as np\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the simulated 3D world\n",
    "world = sio.loadmat('data/world.mat')\n",
    "# the route memory\n",
    "route_memory = sio.loadmat('data/ArcRouteMemory.mat')\n",
    "# the home memory\n",
    "home_memory = sio.loadmat('data/HomeMemory_X0Y0.mat')\n",
    "# the max order of ZM\n",
    "zm_n_max = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Basic simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='orange'>_path integration (PI)_</font>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### * generate PI memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from path_integration import PathIntegrationAgent\n",
    "# set parameters\n",
    "initial_memory = 0.5\n",
    "# create an agent\n",
    "pi = PathIntegrationAgent(initial_memory)\n",
    "# generate the PI memory\n",
    "pi_len = 3.0 # m\n",
    "pi_dir = 45 # deg\n",
    "pi_memory = pi.generate_pi_memory(pi_len, pi_dir, initial_memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### * homing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# homing\n",
    "start_pos = [-700,-700]\n",
    "start_h = 0\n",
    "time_out = 100\n",
    "motor_k = 0.5\n",
    "pos, h, velocity, pi_memory = pi.homing(start_pos, start_h, time_out, motor_k, step_size=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='red'>_visual homin (VH)_</font>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### * training the MB network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from visual_homing import VisualHomingAgent\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": [
    "#### * homing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_pos = [0,-700]\n",
    "start_h = 0\n",
    "time_out = 100\n",
    "vh_k = 0.5\n",
    "motor_k = 1.5 * 0.25\n",
    "pos, h, velocity, mb_out, mb_delta = vh.homing(start_pos, start_h, time_out, vh_k, motor_k, step_size=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='blue'>_route following (RF)_</font>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### * training the ANN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from route_following import RouteFollowingAgent\n",
    "# set parameters\n",
    "rf_learning_rate = 0.1\n",
    "rf_learning_step = 30000\n",
    "# create an instance\n",
    "rf = RouteFollowingAgent(world, route_memory, home_memory, zm_n_max, num_neurons=30)\n",
    "# train the ANN\n",
    "rf.train_nn_network(rf_learning_step, rf_learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### * homing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_pos = [-700,-700]\n",
    "start_h = 0\n",
    "time_out = 100\n",
    "motor_k = 1.5\n",
    "pos, h, velocity, nn_output = rf.homing(start_pos, start_h, time_out, motor_k, step_size=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ___whole model___"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### * create the agent of insect navigation and train the networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from insect_navigation import InsectNavigationAgent\n",
    "\n",
    "# set PI parameters\n",
    "pi_initial_memory = 0.1\n",
    "pi_len = 3.0 # m\n",
    "pi_dir = 90 # deg\n",
    "\n",
    "# set VH parameters\n",
    "num_pn = 81\n",
    "num_kc = 4000\n",
    "vh_learning_rate = 0.1\n",
    "vh_kc_thr = 0.04\n",
    "\n",
    "# set RF parameters\n",
    "rf_learning_rate = 0.1\n",
    "rf_learning_step = 30000\n",
    "ann_num_neurons = 30\n",
    "\n",
    "# set SMP neuron parameters\n",
    "tun_k = 0.0125\n",
    "sn_thr = 5.0\n",
    "\n",
    "# create the insect navigation agent\n",
    "agent = InsectNavigationAgent(world, route_memory, home_memory, zm_n_max, \n",
    "                              vh_learning_rate, vh_kc_thr, num_pn, num_kc, \n",
    "                              tun_k, sn_thr,\n",
    "                              ann_num_neurons,\n",
    "                              pi_initial_memory)\n",
    "# training the MB network\n",
    "en = agent.train_mb_network()\n",
    "# training the ANN network\n",
    "err = agent.train_ann_network(rf_learning_step, rf_learning_rate)\n",
    "# generate PI memory\n",
    "pi = agent.generate_pi_memory(pi_len, pi_dir, pi_initial_memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### * homing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_pos = [-200,-700]\n",
    "start_h = 90\n",
    "time_out = 10\n",
    "motor_k = 1.5\n",
    "\n",
    "# set PI parameters\n",
    "pi_initial_memory = 0.1\n",
    "pi_len = 3.0 # m\n",
    "pi_dir = 225 # deg\n",
    "# generate PI memory\n",
    "pi = agent.generate_pi_memory(pi_len, pi_dir, pi_initial_memory)\n",
    "\n",
    "# VH tuning scalar\n",
    "vh_k = 0.5\n",
    "\n",
    "# start homing\n",
    "end_t, pos, h, velocity, mb_out, mb_delta, ann_out, vh_memory, pi_memory, ra_memory, tn, sn1, sn2 = agent.homing(start_pos, start_h, time_out,\n",
    "                                                                                                               vh_k, sn_thr, tun_k, motor_k,\n",
    "                                                                                                               step_size=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Reproducing behavioural data\n",
    "\n",
    "Using the seperated model for VH and RF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visual_navigation on and off route (Wystrach2012)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='red'>_visual homin (VH)_</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from visual_homing import VisualHomingAgent\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 MB\n",
    "en = vh.train_mb_network()\n",
    "\n",
    "# trials setting\n",
    "start_pos = [0,-700]\n",
    "start_h_s = np.linspace(0, 2 * np.pi, 12, endpoint=False)\n",
    "time_out = 100\n",
    "vh_k = 2.0\n",
    "motor_k = 0.125\n",
    "\n",
    "pos_s = []\n",
    "h_s = []\n",
    "\n",
    "# run the trial\n",
    "for start_h in start_h_s:\n",
    "    pos, h, velocity, mb_out, mb_delta = vh.homing(start_pos, start_h, time_out, vh_k, motor_k, step_size=4)\n",
    "    # store the homing data\n",
    "    pos_s.append(pos)\n",
    "    h_s.append(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='blue'>_route following (RF)_</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from route_following import RouteFollowingAgent\n",
    "# set parameters\n",
    "rf_learning_rate = 0.1\n",
    "rf_learning_step = 30000\n",
    "# create an instance\n",
    "rf = RouteFollowingAgent(world, route_memory, home_memory, zm_n_max, num_neurons=30)\n",
    "# train the ANN\n",
    "rf.train_nn_network(rf_learning_step, rf_learning_rate)\n",
    "\n",
    "start_pos = [-700,-700]\n",
    "start_h_s = np.linspace(0, 2 * np.pi, 2, endpoint=False)\n",
    "time_out = 5\n",
    "motor_k = 0.125\n",
    "\n",
    "pos_s = []\n",
    "h_s = []\n",
    "\n",
    "# run the trial\n",
    "for start_h in start_h_s:\n",
    "    pos, h, velocity, nn_output = rf.homing(start_pos, start_h, time_out, motor_k, step_size=4)\n",
    "    # store the homing data\n",
    "    pos_s.append(pos)\n",
    "    h_s.append(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <font color='gray'>__optimal cue integration__</font>\n",
    "\n",
    "Using the whole model, but turn-off route following by setting ___sn_thr = 0.0___"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### _tuning PI uncertainty (Wystrach2015)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from insect_navigation import InsectNavigationAgent\n",
    "\n",
    "# set PI parameters\n",
    "pi_initial_memory = 0.1\n",
    "pi_len = 3.0 # m\n",
    "pi_dir = 90 # deg\n",
    "\n",
    "# set VH parameters\n",
    "num_pn = 81\n",
    "num_kc = 4000\n",
    "vh_learning_rate = 0.1\n",
    "vh_kc_thr = 0.04\n",
    "\n",
    "# set RF parameters\n",
    "rf_learning_rate = 0.1\n",
    "rf_learning_step = 30000\n",
    "ann_num_neurons = 30\n",
    "\n",
    "# set SMP neuron parameters\n",
    "tun_k = 0.1\n",
    "sn_thr = 0.0\n",
    "\n",
    "# create the insect navigation agent\n",
    "agent = InsectNavigationAgent(world, route_memory, home_memory, zm_n_max, \n",
    "                              vh_learning_rate, vh_kc_thr, num_pn, num_kc, \n",
    "                              tun_k, sn_thr,\n",
    "                              ann_num_neurons,\n",
    "                              pi_initial_memory)\n",
    "# training the MB network\n",
    "en = agent.train_mb_network()\n",
    "\n",
    "# trials setting\n",
    "# generate PI memory\n",
    "pi_len_s = [0.1, 1.0, 3.0, 7.0] # m\n",
    "pi_dir = 90 # deg\n",
    "\n",
    "start_pos = []\n",
    "start_h_s = np.linspace(0, 2 * np.pi, 2, endpoint=False)\n",
    "time_out = 2\n",
    "motor_k = 0.125\n",
    "\n",
    "# VH tuning scalar\n",
    "vh_k = 2.0\n",
    "\n",
    "pos_s = []\n",
    "h_s = []\n",
    "\n",
    "# run the trials \n",
    "for start_h in start_h_s:\n",
    "    for pi_len in pi_len_s:\n",
    "        # generate different PI home vector length\n",
    "        pi = agent.generate_pi_memory(pi_len, pi_dir, pi_initial_memory)\n",
    "        # start homing\n",
    "        end_t, pos, h, velocity, mb_out, mb_delta, ann_out, vh_memory, pi_memory, ra_memory, tn, sn1, sn2 = agent.homing(start_pos, start_h, time_out,\n",
    "                                                                                                                   vh_k, sn_thr, tun_k, motor_k,\n",
    "                                                                                                                   step_size=8)\n",
    "        pos_s.append(pos)\n",
    "        h_s.append(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### _tuning VH uncertainty (Legge2014)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from insect_navigation import InsectNavigationAgent\n",
    "\n",
    "# set PI parameters\n",
    "pi_initial_memory = 0.1\n",
    "pi_len = 3.0 # m\n",
    "pi_dir = 90 # deg\n",
    "\n",
    "# set VH parameters\n",
    "num_pn = 81\n",
    "num_kc = 4000\n",
    "vh_learning_rate = 0.1\n",
    "vh_kc_thr = 0.04\n",
    "\n",
    "# set RF parameters\n",
    "rf_learning_rate = 0.1\n",
    "rf_learning_step = 30000\n",
    "ann_num_neurons = 30\n",
    "\n",
    "# set SMP neuron parameters\n",
    "tun_k = 0.1\n",
    "sn_thr = 0.0\n",
    "\n",
    "# create the insect navigation agent\n",
    "agent = InsectNavigationAgent(world, route_memory, home_memory, zm_n_max, \n",
    "                              vh_learning_rate, vh_kc_thr, num_pn, num_kc, \n",
    "                              tun_k, sn_thr,\n",
    "                              ann_num_neurons,\n",
    "                              pi_initial_memory)\n",
    "# training the MB network\n",
    "en = agent.train_mb_network()\n",
    "\n",
    "# generate PI memory\n",
    "pi = agent.generate_pi_memory(pi_len, pi_dir, pi_initial_memory)\n",
    "        \n",
    "# trials setting\n",
    "# generate PI memory\n",
    "pi_len_s = [0.1, 1.0, 3.0, 7.0] # m\n",
    "pi_dir = 90 # deg\n",
    "\n",
    "start_pos_s = []\n",
    "start_h_s = np.linspace(0, 2 * np.pi, 2, endpoint=False)\n",
    "time_out = 2\n",
    "motor_k = 0.125\n",
    "\n",
    "# VH tuning scalar\n",
    "vh_k = 2.0\n",
    "\n",
    "pos_s = []\n",
    "h_s = []\n",
    "\n",
    "# run the trials \n",
    "for start_h in start_h_s:\n",
    "    for start_pos in start_pos_s:\n",
    "        # start homing\n",
    "        end_t, pos, h, velocity, mb_out, mb_delta, ann_out, vh_memory, pi_memory, ra_memory, tn, sn1, sn2 = agent.homing(start_pos, start_h, time_out,\n",
    "                                                                                                                   vh_k, sn_thr, tun_k, motor_k,\n",
    "                                                                                                                   step_size=8)\n",
    "        pos_s.append(pos)\n",
    "        h_s.append(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ___whole model for all the properties we want___\n",
    "\n",
    "\n",
    "\n",
    "This can be done by using the ___whole model___ cell in the __Basic simulation__ section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Other analysis tools\n",
    "\n",
    "This section contains some code to generate some analysis data, such as the ZM encoding of the 3D world."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### _Check the frequency phase tracking across the world_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from model import *\n",
    "# check the data for RF\n",
    "# 1.RF memory , 2.the phase-tracking, 3.RF suggested\n",
    "\n",
    "# sampled num\n",
    "sample_num = 20\n",
    "\n",
    "# sampled locations\n",
    "pos_x = np.linspace(-10,2,sample_num)\n",
    "pos_y = np.linspace(-8,2,sample_num)\n",
    "# sampled heading\n",
    "h = np.linspace(-np.pi,np.pi,10)\n",
    "\n",
    "# stored data\n",
    "ann_output = np.zeros([sample_num**2,len(h)])\n",
    "current_zm_p = np.zeros([sample_num**2,len(h)])\n",
    "\n",
    "vc_phase_prefs = np.linspace(-np.pi,np.pi,8,endpoint=False)\n",
    "\n",
    "for i in range(sample_num**2):\n",
    "    for k,h_i in enumerate(h):\n",
    "        A,P = visual_sense(InsectNaviAgent.world_data, pos_x[i%sample_num],pos_y[i//sample_num],h_i,nmax=InsectNaviAgent.nmax)\n",
    "        current_zm_p[i,k] = P[16]\n",
    "        nn_input = A.copy()\n",
    "        nn_input = (nn_input - np.min(nn_input))/np.max(nn_input)\n",
    "        nn_res = InsectNaviAgent.VC.run_test(nn_input)\n",
    "        ann_output[i,k] = np.arctan2(np.sum(nn_res*np.sin(vc_phase_prefs)), \n",
    "                                     np.sum(nn_res*np.cos(vc_phase_prefs)))          \n",
    "        \n",
    "# sio.savemat('QuiverPlotData_X-10_2_Y-8_2_SH20.mat',{'ann_output':ann_output, 'current_zm_p':current_zm_p})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### _generate frequency coding memory of specific locations and heading defined by __pos and h___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate visual memory along PI route\n",
    "from image_processing import get_img_view\n",
    "from zernike_moment import zernike_moment\n",
    "\n",
    "home_pos_new = np.array([0, 0])\n",
    "sample_len = 20\n",
    "pos = np.array([[home_pos_new[0] - sample_len, home_pos_new[1]],\n",
    "                [home_pos_new[0] - sample_len, home_pos_new[1] - sample_len],\n",
    "                [home_pos_new[0], home_pos_new[1] - sample_len],\n",
    "                [home_pos_new[0] + sample_len, home_pos_new[1] - sample_len],\n",
    "                [home_pos_new[0] + sample_len, home_pos_new[1]],\n",
    "                [home_pos_new[0] + sample_len, home_pos_new[1] + sample_len],\n",
    "                [home_pos_new[0], home_pos_new[1] + sample_len],\n",
    "                [home_pos_new[0] - sample_len, home_pos_new[1] + sample_len],\n",
    "                [home_pos_new[0], home_pos_new[1]],\n",
    "                ])\n",
    "h = np.array([0, 0.25, 0.5, 0.75, 1.0, -0.75, -0.5, -0.25, 0.5]) * np.pi\n",
    "        \n",
    "memory_ZM_As = np.zeros([len(pos), 81])\n",
    "memory_ZM_Ps = np.zeros([len(pos), 81])\n",
    "memory_imgs =  np.zeros([len(pos), 208, 208])\n",
    "for i in range(len(pos)):\n",
    "    memory_imgs[i, :, :] = get_img_view(InsectNaviAgent.world_data, pos[i, 0] / 100.0, pos[i, 1] / 100.0, 0.01,\n",
    "                                        h[i], hfov_d=360, wrap=True,\n",
    "                                        blur=False, blur_kernel_size=3)\n",
    "    index = 0\n",
    "    for n in range(InsectNaviAgent.nmax + 1):\n",
    "        for m in range(n + 1):\n",
    "            if (n - m) % 2 == 0:\n",
    "                M, memory_ZM_As[i, index], memory_ZM_Ps[i, index] = zernike_moment(255 - memory_imgs[i, :, :], n, m)\n",
    "                index += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### _compute the frequecny encoding of locations across the world_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "from image_processing import visual_sense\n",
    "# generate the ZM coefficients for the world\n",
    "sample_num = 20\n",
    "map_x = np.linspace(-10,10,sample_num)\n",
    "map_y = np.linspace(-10,10,sample_num)\n",
    "h = np.zeros([len(map_x),len(map_y)])\n",
    "world_zm_a = np.zeros([len(map_x),len(map_y),81])\n",
    "world_zm_p = np.zeros([len(map_x),len(map_y),81])\n",
    "for i,y in enumerate(map_y):\n",
    "    for j,x in enumerate(map_x):\n",
    "        A,P = visual_sense(InsectNaviAgent.world_data,x,y,h[i,j],nmax=16)\n",
    "        world_zm_a[j,i] = A\n",
    "        world_zm_p[j,i] = P"
   ]
  }
 ],
 "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
}
