Current injection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Script to run current injection experiments. All results are saved to pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python2.7/site-packages/matplotlib/font_manager.py:273: UserWarning: Matplotlib is building the font cache using fc-list. This may take a moment.\n",
      "  warnings.warn('Matplotlib is building the font cache using fc-list. This may take a moment.')\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "    \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import central_complex\n",
    "import cx_rate\n",
    "import trials\n",
    "import plotter\n",
    "\n",
    "save_figs=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def update_cells_w_inject(heading, velocity, tb1, memory, cx, filtered_steps=0.0, inject=None, inject_val=1.0):\n",
    "    \"\"\"Generate activity for all cells, based on previous activity and current\n",
    "    motion.\"\"\"\n",
    "    tl2 = cx.tl2_output(heading)\n",
    "    if 'tl2' in inject:\n",
    "        tl2[inject['tl2']] = inject_val\n",
    "\n",
    "    cl1 = cx.cl1_output(tl2)\n",
    "    if 'cl1' in inject:\n",
    "        cl1[inject['cl1']] = inject_val\n",
    "\n",
    "    tb1 = cx.tb1_output(cl1, tb1)\n",
    "    if 'tb1' in inject:\n",
    "        tb1[inject['tb1']] = inject_val\n",
    "\n",
    "    # Speed\n",
    "    flow = cx.get_flow(heading, velocity)\n",
    "    tn1 = cx.tn1_output(flow)\n",
    "    if 'tn1' in inject:\n",
    "        tn1[inject['tn1']] = inject_val\n",
    "    tn2 = cx.tn2_output(flow)\n",
    "    if 'tn2' in inject:\n",
    "        tn2[inject['tn2']] = inject_val\n",
    "\n",
    "    # Update memory for distance just travelled\n",
    "    memory = cx.cpu4_update(memory, tb1, tn1, tn2)\n",
    "    if 'memory' in inject:\n",
    "        memory[inject['memory']] = inject_val\n",
    "    cpu4 = cx.cpu4_output(memory)\n",
    "    if 'cpu4' in inject:\n",
    "        cpu4[inject['cpu4']] = inject_val\n",
    "\n",
    "    # Steer based on memory and direction\n",
    "    cpu1 = cx.cpu1_output(tb1, cpu4)\n",
    "    if 'cpu1' in inject:\n",
    "        cpu1[inject['cpu1']] = inject_val\n",
    "    motor = cx.motor_output(cpu1)\n",
    "\n",
    "    return tl2, cl1, tb1, tn1, tn2, memory, cpu4, cpu1, motor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import bee_simulator\n",
    "def homing_w_inject(T, tb1, memory, cx, acceleration=0.15, drag=0.15,\n",
    "           current_heading=0.0, current_velocity=np.array([0.0, 0.0]),\n",
    "           turn_sharpness=1.0, inject=None, start_inject=250, stop_inject=1000,\n",
    "           inject_val=1.0):\n",
    "    \"\"\"Based on current state, return home. First is duplicate\"\"\"\n",
    "    # TODO: Check some serious off by one errors here! on a small course best\n",
    "    headings = np.empty(T + 1)\n",
    "    headings[0] = current_heading\n",
    "    velocity = np.empty([T + 1, 2])\n",
    "    velocity[0, :] = current_velocity\n",
    "\n",
    "    cx_log = trials.CXLogger(0, T + 1, cx)\n",
    "\n",
    "    for t in range(1, T + 1):\n",
    "        tl2, cl1, tb1, tn1, tn2, memory, cpu4, cpu1, motor = update_cells_w_inject(\n",
    "            heading=headings[t - 1],  # Remove sign to use proportionate shift\n",
    "            velocity=velocity[t - 1, :],\n",
    "            tb1=tb1,\n",
    "            memory=memory,\n",
    "            cx=cx,\n",
    "            inject=inject if t>start_inject and t<stop_inject else {},\n",
    "            inject_val=inject_val)\n",
    "\n",
    "        \n",
    "        cx_log.update_log(t, tl2, cl1, tb1, tn1, tn2, memory, cpu4, cpu1, motor)\n",
    "\n",
    "        rotation = turn_sharpness * motor\n",
    "\n",
    "        headings[t], velocity[t, :] = bee_simulator.get_next_state(\n",
    "            headings[t - 1], velocity[t - 1, :], rotation, acceleration, drag)\n",
    "    return headings, velocity, cx_log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "T_outbound = 500\n",
    "T_inbound = 1500\n",
    "\n",
    "# TEST CODE\n",
    "memories = []\n",
    "tb1_snapshots = []\n",
    "\n",
    "cx = cx_rate.CXRatePontin()\n",
    "\n",
    "start_inject_out = 50\n",
    "stop_inject_out = 250\n",
    "start_inject_in = 550\n",
    "stop_inject_in = 750\n",
    "v_out = np.ones([T_outbound, 2]) * -0.5\n",
    "v_out[:,0] = 0\n",
    "h_out = np.ones(T_outbound) * np.pi\n",
    "inject_val = 1.0\n",
    "\n",
    "labels = ['TB1:1', 'TB1:2', 'TB1:3', 'TB1:4', 'TB1:5', 'TB1:6', 'TB1:7', 'TB1:8']\n",
    "\n",
    "experiments = [{'name':'tb1_injection_out', 'start':start_inject_out, 'stop':stop_inject_out,\n",
    "                'injects':[{'tb1':[0]}, {'tb1':[1]}, {'tb1':[2]}, {'tb1':[3]},\n",
    "                           {'tb1':[4]}, {'tb1':[5]}, {'tb1':[6]}, {'tb1':[7]}],\n",
    "                'labels':['TB1:1', 'TB1:2', 'TB1:3', 'TB1:4', 'TB1:5', 'TB1:6', 'TB1:7', 'TB1:8']},\n",
    "               {'name':'tb1_injection_in', 'start':start_inject_in, 'stop':stop_inject_in,\n",
    "                'injects':[{'tb1':[0]}, {'tb1':[1]}, {'tb1':[2]}, {'tb1':[3]},\n",
    "                           {'tb1':[4]}, {'tb1':[5]}, {'tb1':[6]}, {'tb1':[7]}],\n",
    "                'labels':['TB1:1', 'TB1:2', 'TB1:3', 'TB1:4', 'TB1:5', 'TB1:6', 'TB1:7', 'TB1:8']},\n",
    "               {'name':'tn2_injection_out', 'start':start_inject_out, 'stop':stop_inject_out,\n",
    "                'injects':[{'tn2':[0]}, {'tn2':[1]}],\n",
    "                'labels':['TN2:1', 'TN2:2']},\n",
    "               {'name':'tn2_injection_in', 'start':start_inject_in, 'stop':stop_inject_in,\n",
    "                'injects':[{'tn2':[0]}, {'tn2':[1]}],\n",
    "                'labels':['TN2:1', 'TN2:2']},\n",
    "               {'name':'cpu1_injection_out', 'start':start_inject_out, 'stop':stop_inject_out,\n",
    "                'injects':[{'cpu1':[0]}, {'cpu1':[0, 1]}, {'cpu1':[0, 1, 2]}, {'cpu1':[0, 1, 2, 3]}],\n",
    "                'labels':['CPU1:1', 'CPU1:1,2', 'CPU1:1-3', 'CPU1:1-4']},\n",
    "               {'name':'cpu1_injection_in', 'start':start_inject_in, 'stop':stop_inject_in,\n",
    "                'injects':[{'cpu1':[0]}, {'cpu1':[0, 1]}, {'cpu1':[0, 1, 2]}, {'cpu1':[0, 1, 2, 3]}],\n",
    "                'labels':['CPU1:1', 'CPU1:1,2', 'CPU1:1-3', 'CPU1:1-4']}]\n",
    "\n",
    "line_colors = ['purple', 'orange', 'gold', 'olivedrab', 'lightgreen', 'green', 'turquoise', 'blue', 'darkblue']\n",
    "\n",
    "for experiment in experiments:\n",
    "    fig, ax = plt.subplots()\n",
    "    injects = experiment['injects']\n",
    "    start_inject = experiment['start']\n",
    "    stop_inject = experiment['stop']\n",
    "    labels = experiment['labels']\n",
    "    handles = []\n",
    "    \n",
    "    for i, inject in enumerate(injects):\n",
    "        #Initialize memory stuff\n",
    "        tb1 = np.zeros(central_complex.N_TB1)\n",
    "        memory = 0.5 * np.ones(central_complex.N_CPU4)\n",
    "        log_out = trials.CXLogger(T_outbound, 0, cx)\n",
    "\n",
    "        for t in range(T_outbound):\n",
    "            tl2, cl1, tb1, tn1, tn2, memory, cpu4, cpu1, motor = update_cells_w_inject(\n",
    "                heading=h_out[t], velocity=v_out[t], tb1=tb1, memory=memory,\n",
    "                cx=cx, inject=inject if t>start_inject and t<stop_inject else {},\n",
    "                inject_val=inject_val)\n",
    "            log_out.update_log(t, tl2, cl1, tb1, tn1, tn2, memory, cpu4, cpu1, motor)\n",
    "\n",
    "        tb1 = log_out.tb1[:, -1]\n",
    "        memory = log_out.memory[:, -1]\n",
    "        cpu4_snapshot = memory.copy()\n",
    "        memories.append(cpu4_snapshot)\n",
    "\n",
    "        h_in, v_in, log_in = homing_w_inject(T=T_inbound, tb1=tb1, memory=cpu4_snapshot, cx=cx, acceleration=0.1,\n",
    "                                             current_heading=h_out[-1], current_velocity=v_out[-1], drag=0.15,\n",
    "                                             inject=inject, start_inject=start_inject-T_outbound,\n",
    "                                             stop_inject=stop_inject-T_outbound, inject_val=inject_val)\n",
    "        h = np.hstack([h_out, h_in])\n",
    "        v = np.vstack([v_out, v_in])\n",
    "\n",
    "        log = log_out + log_in\n",
    "\n",
    "        # Zoomed for spiral observation\n",
    "        ax.set_aspect('equal')\n",
    "        ax.set_xlim([-200, 200])\n",
    "        ax.set_aspect('equal')\n",
    "        ax.set_xticks([-200, -100, 0, 100, 200])\n",
    "        \n",
    "        xy = np.vstack([np.array([0.0, 0.0]), np.cumsum(v, axis=0)])\n",
    "        x, y = xy[:, 0], xy[:, 1]\n",
    "        \n",
    "        # Plot line out\n",
    "        if i==0:\n",
    "            line_out, = ax.plot(x[0:T_outbound+1], y[0:T_outbound+1], lw=0.5,\n",
    "                            color='purple', label='Outbound')\n",
    "            handles.append(line_out)\n",
    "        \n",
    "        # Plot line in\n",
    "        line_in, = ax.plot(x[T_outbound:], y[T_outbound:], lw=0.5,\n",
    "                            color=line_colors[i+1], label=labels[i])\n",
    "        handles.append(line_in)\n",
    "\n",
    "\n",
    "        # Plot injection\n",
    "        ax.plot(x[start_inject:stop_inject], y[start_inject:stop_inject], color='cyan')\n",
    "        \n",
    "        if experiment['name'] == 'tb1_injection_in':\n",
    "            tb1_snapshots.append(log.tb1[:, start_inject_in+50])\n",
    "        \n",
    "    \n",
    "        if experiment['name'] == 'tn2_injection_out' or experiment['name'] == 'tn2_injection_in':\n",
    "            fig2, ax2 = plotter.plot_traces(\n",
    "                    log, include=['TN1', 'TN2', 'TB1', 'CPU4', 'CPU1', 'motor'])\n",
    "            plotter.save_plot(fig2, 'current_injection_{0}{1}'.format(experiment['name'], i))\n",
    "                \n",
    "    l = ax.legend(handles=handles,\n",
    "                  loc='best',\n",
    "                  fontsize=10,\n",
    "                  handlelength=0,\n",
    "                  handletextpad=0)\n",
    "    for handle in l.legendHandles:\n",
    "        handle.set_visible(False)\n",
    "    for i, text in enumerate(l.get_texts()):\n",
    "        text.set_color(line_colors[i])\n",
    "    l.draw_frame(False)\n",
    "    \n",
    "    if save_figs:\n",
    "        plotter.save_plot(fig, 'current_injection_{0}'.format(experiment['name']))\n",
    "        \n",
    "    if experiment['name'] == 'cpu1_injection_in':\n",
    "        ax.set_xlim([-50, 50])\n",
    "        ax.set_ylim([-300, -200])\n",
    "        plotter.save_plot(fig, 'current_injection_cpu1_injection_in_zoomed')\n",
    "        \n",
    "plt.close('all')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}