three_figure_widget.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dbb52a0024164bb0b831e565b1739e55",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=1.0, continuous_update=False, description='wind_mag', max=4.0, readout…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import time\n",
    "import scipy\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import sys\n",
    "import itertools\n",
    "import json\n",
    "import numpy as np\n",
    "from odor_tracking_sim import utility\n",
    "from pompy import models\n",
    "# from matplotlib.widgets import Slider,Button\n",
    "# from matplotlib.transforms import Bbox\n",
    "from extras import UpdatingVPatch,plot_wedges\n",
    "from core_functions import f0,f1,f1_wedge,f2,f3,f4,f5\n",
    "from ipywidgets import interactive\n",
    "import ipywidgets as widgets\n",
    "from matplotlib.gridspec import GridSpec\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "#Constants that don't change with drag bars\n",
    "num_flies = 20000\n",
    "fly_speed = 1.6\n",
    "\n",
    "number_sources = 8\n",
    "radius_sources = 1000.0\n",
    "source_locations, _ = utility.create_circle_of_sources(number_sources,\n",
    "                radius_sources,None)\n",
    "source_pos = scipy.array([scipy.array(tup) for tup in source_locations])\n",
    "release_location = np.zeros(2)\n",
    "\n",
    "intended_heading_angles = np.random.uniform(0,2*np.pi,num_flies)\n",
    "intended_heading_angles = np.linspace(0,2*np.pi,num_flies)\n",
    "\n",
    "initial_cone_angle = np.radians(10.)\n",
    "\n",
    "windmag_slider = widgets.FloatSlider(\n",
    "    value=1.,\n",
    "    min=0,\n",
    "    max=4.0,\n",
    "    step=0.1,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    ")\n",
    "\n",
    "cone_angle_slider = widgets.FloatSlider(\n",
    "    value=np.degrees(initial_cone_angle),\n",
    "    min=0.,\n",
    "    max=40.0,\n",
    "    step=1.,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    ")\n",
    "\n",
    "K_slider = widgets.FloatSlider(\n",
    "    value=0.4,\n",
    "    min=0.,\n",
    "    max=1.0,\n",
    "    step=0.1,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    ")\n",
    "\n",
    "x_0_slider = widgets.FloatSlider(\n",
    "    value=300.,\n",
    "    min=0.,\n",
    "    max=1000.0,\n",
    "    step=10.,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    ")\n",
    "\n",
    "toggle = widgets.ToggleButton(\n",
    "    value=False,\n",
    "    description='Headings | Intersections',\n",
    "    disabled=False,\n",
    "    button_style='', # 'success', 'info', 'warning', 'danger' or ''\n",
    "    tooltip='Description')\n",
    "\n",
    "\n",
    "def f(wind_mag,toggle,cone_angle,K,x_0):\n",
    "\n",
    "    #----OBTAIN NEW VALUES--------\n",
    "    wind_angle = 7*scipy.pi/8.\n",
    "    cone_angle = np.radians(cone_angle)\n",
    "\n",
    "    release_times=0.\n",
    "\n",
    "    K = -1*K\n",
    "#     x_0 = 300\n",
    "\n",
    "    #------RECOMPUTE OUTPUTS------\n",
    "    track_heading_angles,dispersing_speeds = f0(intended_heading_angles,wind_mag,\n",
    "        wind_angle)\n",
    "    intersection_distances,dispersal_distances = f1_wedge(\n",
    "    track_heading_angles,source_pos,wind_angle,cone_angle)\n",
    "    success_probabilities = f2(intersection_distances,K,x_0,source_pos,wind_angle)\n",
    "    plume_assignments = f3(success_probabilities,dispersal_distances)\n",
    "    dispersal_travel_times,release_to_chosen_plume_distances = f4(\n",
    "            plume_assignments,dispersal_distances,dispersing_speeds)\n",
    "    arrival_times,chasing_times,\\\n",
    "    which_flies,which_traps = f5(plume_assignments,dispersal_travel_times,\n",
    "        intersection_distances,fly_speed,release_times)\n",
    "\n",
    "  \n",
    "    #plot scaffolding\n",
    "    gs = GridSpec(9, 9)\n",
    "    fig = plt.figure(figsize=(10,10))\n",
    "    \n",
    "    #------\"FIGURE\" 1 \n",
    "    ax =  fig.add_subplot(gs[0:4,0:4])   \n",
    "    xlim = (-1500., 1500.)\n",
    "    ylim = (-1500., 1500.)\n",
    "    im_extents = xlim[0], xlim[1], ylim[0], ylim[1]\n",
    "\n",
    "    ax.set_ylim(list(ylim))\n",
    "    ax.set_xlim(list(xlim))\n",
    "    \n",
    "    if toggle:\n",
    "        x,y = release_to_chosen_plume_distances*np.cos(track_heading_angles), \\\n",
    "            release_to_chosen_plume_distances*np.sin(track_heading_angles)\n",
    "        \n",
    "    else: \n",
    "        time = 5*60.\n",
    "        mag = time*dispersing_speeds \n",
    "        x,y = mag*np.cos(track_heading_angles), \\\n",
    "            mag*np.sin(track_heading_angles)\n",
    "            \n",
    "    plt.scatter(x,y,alpha=0.02,color='r')\n",
    "    \n",
    "    wedge_points = plot_wedges(source_pos,wind_angle,cone_angle)\n",
    "\n",
    "    plume_wedges = [matplotlib.patches.Polygon(\n",
    "        wedge_points[:,i,:],color='black',alpha=0.2) for i in range(number_sources)]\n",
    "\n",
    "    for plume_wedge in plume_wedges:\n",
    "        ax.add_patch(plume_wedge)\n",
    "\n",
    "    for x,y in source_locations:\n",
    "        plt.scatter(x,y,marker='x',s=50,c='k')\n",
    "    \n",
    "    ax.set_aspect('equal')\n",
    "    \n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    \n",
    "    #------------\"FIGURE\" 2 : ARRIVAL CDFs---------------# \n",
    "    num_bins = 50\n",
    "\n",
    "    trap_counts = scipy.zeros(8)\n",
    "    rasters = []\n",
    "    labels = ['N','NE','E','SE','S','SW','W','NW']\n",
    "    sim_reorder = scipy.array([3,2,1,8,7,6,5,4])\n",
    "    axes = []\n",
    "    lines = []\n",
    "    cdf_patches = []\n",
    "    cdf_steepnesses = np.zeros(8)\n",
    "    first_hit_times = np.full(8,np.nan)\n",
    "    new_maxes = 400*np.ones(8)\n",
    "    for i in range(8):\n",
    "\n",
    "        row = sim_reorder[i]-1\n",
    "        ax =  fig.add_subplot(gs[row,5:])    \n",
    "        t_sim = arrival_times[which_traps==i]\n",
    "\n",
    "        if len(t_sim)==0:\n",
    "            ax.set_xticks([0,10,20,30,40,50])\n",
    "            trap_total = 0\n",
    "            pass\n",
    "        else:\n",
    "            t_sim = t_sim/60.\n",
    "            (n, bins) = np.histogram(t_sim,bins=num_bins,\n",
    "                range=(0,max(t_sim)))\n",
    "            cum_n = np.cumsum(n)\n",
    "            line, = plt.step(bins,np.hstack((np.array([0,]),cum_n)))\n",
    "            lines.append(line)\n",
    " \n",
    "            patch_object = UpdatingVPatch(min(t_sim),max(t_sim)-min(t_sim))\n",
    "            ax.add_patch(patch_object.rectangle)\n",
    "            cdf_patches.append(patch_object)\n",
    "            try:\n",
    "                trap_counts[i]=max(cum_n)\n",
    "            except(IndexError):\n",
    "                trap_counts[i]=0\n",
    "\n",
    "            cdf_steepnesses[i] = trap_counts[i]/(max(t_sim)-min(t_sim))\n",
    "            first_hit_times[i] = min(t_sim)\n",
    "            new_maxes[i] = max(400.,50*np.ceil(max(cum_n)/50.))\n",
    "    \n",
    "\n",
    "\n",
    "        if sim_reorder[i]-1==0:\n",
    "             ax.set_title('Cumulative Trap Arrivals')\n",
    "\n",
    "        ax.set_xlim([0,50])\n",
    "        plt.tick_params(\n",
    "        axis='x',          # changes apply to the x-axis\n",
    "        which='both',      # both major and minor ticks are affected\n",
    "        bottom=True,      # ticks along the bottom edge are off\n",
    "        top=False,         # ticks along the top edge are off\n",
    "        labelbottom=True)\n",
    "        ax.text(1.1,0.5,str(labels[sim_reorder[i]-1]),transform=ax.transAxes,fontsize=20,\n",
    "            horizontalalignment='center',verticalalignment='center')\n",
    "        if sim_reorder[i]-1==7:\n",
    "            ax.set_xlabel('Time (min)',x=0.5,horizontalalignment='center',fontsize=20)\n",
    "            plt.tick_params(axis='x', which='major', labelsize=15)\n",
    "        else:\n",
    "            ax.set_xticklabels('')\n",
    "        axes.append(ax)\n",
    "\n",
    "    for i,ax in enumerate(axes):\n",
    "        ax.set_yticks([0,200,400,600,800])\n",
    "        ax.set_ylim([0,np.max(new_maxes)])\n",
    "        patch_object = cdf_patches[i]\n",
    "        patch_object.rectangle.set_height(ax.get_ylim()[1])\n",
    "\n",
    "  \n",
    "    #------------\"FIGURE\" 3 : Trap Histograms---------------# \n",
    "    \n",
    "    ax =  fig.add_subplot(gs[4:8,0:4])   \n",
    "    ax.set_aspect('equal')\n",
    "    steepness_max = 300.\n",
    "\n",
    "    num_traps = np.shape(source_pos)[0]\n",
    "    trap_locs = (2*np.pi/num_traps)*np.array(range(num_traps))\n",
    "    #Set 0s to 1 for plotting purposes\n",
    "    trap_counts[trap_counts==0] = .5\n",
    "    radius_scale = 0.3\n",
    "    plot_size = 1.5\n",
    "    trap_locs_2d = [(scipy.cos(trap_loc),scipy.sin(trap_loc)) for trap_loc in trap_locs]\n",
    "    trap_patches = [plt.Circle(center, size,\n",
    "        alpha=min(cdf_steepnesses[i]/steepness_max,1.)) for center, size, i in zip(\n",
    "            trap_locs_2d, radius_scale*trap_counts/np.max(new_maxes),range(8))]\n",
    "    for trap_patch in trap_patches:\n",
    "        ax.add_patch(trap_patch)\n",
    "\n",
    "    vmin = 5.;vmax = 20.\n",
    "    trap_cmap_vals = (first_hit_times-vmin)/vmax\n",
    "    trap_cmap  = matplotlib.cm.get_cmap('plasma_r')\n",
    "  \n",
    "    for trap_cmap_val,trap_patch in zip(trap_cmap_vals,trap_patches):\n",
    "        # trap_patch.set_color(trap_cmap(trap_cmap_val)[:-1])\n",
    "        color = tuple(np.array((trap_cmap(trap_cmap_val)[:-1])).astype(float).tolist())\n",
    "        trap_patch.set_color(color)\n",
    "\n",
    "    ax.set_ylim([-plot_size,plot_size]);ax.set_xlim([-plot_size,plot_size])\n",
    "    ax.set_xticks([])\n",
    "    ax.set_xticklabels('')\n",
    "    ax.set_yticks([])\n",
    "    ax.set_yticklabels('')\n",
    "\n",
    "    coll = matplotlib.collections.PatchCollection(trap_patches)#, facecolors=colors,edgecolors=colors)\n",
    "    coll.set(cmap=trap_cmap,array=[])\n",
    "    coll.set_clim(vmin=vmin,vmax=vmax)\n",
    "    fig.colorbar(coll, ax=ax,pad=0.2)\n",
    "    ax.text(2.1,.1,'First Arrival Time (min)',horizontalalignment='center',\n",
    "        rotation=-90,verticalalignment='center',fontsize=15)\n",
    "    #Wind arrow\n",
    "    plt.arrow(0.5, 0.5, 0.1*scipy.cos(wind_angle), 0.1*scipy.sin(wind_angle),transform=ax.transAxes,color='b',\n",
    "        width=0.01,head_width=0.05)\n",
    "    fontsize=15\n",
    "    ax.text(0,1.5,'N',horizontalalignment='center',verticalalignment='center',fontsize=fontsize)\n",
    "    ax.text(0,-1.5,'S',horizontalalignment='center',verticalalignment='center',fontsize=fontsize)\n",
    "    ax.text(1.5,0,'E',horizontalalignment='center',verticalalignment='center',fontsize=fontsize)\n",
    "    ax.text(-1.5,0,'W',horizontalalignment='center',verticalalignment='center',fontsize=fontsize)\n",
    "    ax.axis('off')\n",
    "\n",
    "    #Tidy up plot\n",
    "    plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9,wspace=0.5,hspace=0.5)\n",
    "    \n",
    "    \n",
    "interactive(f, wind_mag=windmag_slider,cone_angle=cone_angle_slider,K=K_slider,x_0=x_0_slider,toggle=toggle)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "be103",
   "language": "python",
   "name": "be103"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}