trap_arrival_widget.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "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 ipywidgets import interactive\n",
    "from ipywidgets import Layout, Button, Box, VBox,Label\n",
    "import ipywidgets as widgets\n",
    "from matplotlib.gridspec import GridSpec\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_circle_of_sources(number_sources,radius,strength):\n",
    "    location_list = []\n",
    "    for i in range(number_sources):\n",
    "        angle = i*(2.0*scipy.pi)/number_sources\n",
    "        x = radius*scipy.cos(angle)\n",
    "        y = radius*scipy.sin(angle)\n",
    "        location_list.append((x,y))\n",
    "    strength_list = [strength for x in location_list]\n",
    "    return location_list, strength_list\n",
    "\n",
    "def sigmoid(x,x_0,L,y_0,k):\n",
    "    return (x-x_0)+(L/2)+y_0 - L/(np.exp(-k*(x-x_0))+1)\n",
    "\n",
    "\n",
    "def speed_sigmoid_func(x):\n",
    "    x_0a = -0.4\n",
    "    x_0b = 3.6\n",
    "    L = 0.8\n",
    "    k = 4.\n",
    "    y_0 = 1.6\n",
    "    output = np.zeros_like(x)\n",
    "    output[(x>=x_0a)&(x<=x_0b)] = 1.6\n",
    "    output[x<x_0a] = sigmoid(x[x<x_0a],x_0a,L,y_0,k)\n",
    "    output[x>x_0b] = sigmoid(x[x>x_0b],x_0b,L,y_0,k)\n",
    "    return output\n",
    "def find_nearest(array, value):\n",
    "    #For each element in value, returns the index of array it is closest to.\n",
    "    #array should be 1 x n and value should be m x 1\n",
    "    idx = (np.abs(array - value)).argmin(axis=1) #this rounds up and down (\n",
    "    #of the two values in array closest to value, picks the closer. (not the larger or the smaller)\n",
    "    return idx\n",
    "\n",
    "def sigmoidm(x,x_0,L,y_0,k,m):\n",
    "    return m*(x-x_0)+(L/2)+y_0 - L/(np.exp(-k*(x-x_0))+1)\n",
    "\n",
    "def speed_sigmoid_func_kwarg(x,x_0a=-0.4,x_0b=3.6,\n",
    "    L=0.8,k=4.,y_0=1.6,m=1.):\n",
    "    output = np.zeros_like(x)\n",
    "    output[(x>=x_0a)&(x<=x_0b)] = y_0\n",
    "    output[x<x_0a] = sigmoidm(x[x<x_0a],x_0a,L,y_0,k,m)\n",
    "    output[x>x_0b] = sigmoidm(x[x>x_0b],x_0b,L,y_0,k,m)\n",
    "    return output\n",
    "\n",
    "def f_rotated(inputs,x_0a,x_0b,L,k,y_0,m,theta):\n",
    "    yl,yu = -4,8\n",
    "    buffer = 10\n",
    "    num_points = len(inputs)\n",
    "    outputs = speed_sigmoid_func_kwarg(inputs,x_0a,x_0b,L,k,y_0,m)\n",
    "    rot_mat = np.array([[np.cos(theta),-1.*np.sin(theta)],[np.sin(theta),np.cos(theta)]])\n",
    "    rotation_origin = np.array([x_0a+(x_0b-x_0a)/2,y_0])\n",
    "    rotation_origin_ones = np.repeat(rotation_origin[:,None],num_points,axis=1)\n",
    "    inputs1,outputs1 = np.dot(rot_mat,np.vstack((inputs,outputs))-rotation_origin_ones)+rotation_origin_ones\n",
    "    which_inputs = find_nearest(inputs1,inputs[:,None])\n",
    "    return outputs1[which_inputs]\n",
    "\n",
    "def new_speed_sigmoid_func_perfect_controller(x):\n",
    "    x_0a = -0.4\n",
    "    x_0b= 1.45\n",
    "    L=0.8\n",
    "    k=4.\n",
    "    y_0=1.6\n",
    "    m=1.\n",
    "    theta=0.\n",
    "    return f_rotated(x,x_0a,x_0b,L,k,y_0,m,theta)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f0(intended_heading_angles,wind_mag,wind_angle,\n",
    "    x_0a,x_0b,L,y_0,plot=False,c1=0,c2=1):\n",
    "    \n",
    "    #Converts intended heading angles to track heading angles\n",
    "    #Currently only have computation for c1 = 0, c2 = 1\n",
    "    n = len(intended_heading_angles)\n",
    "    # intended_heading_angles = np.linspace(360./n,360,n)*np.pi/180\n",
    "    signed_wind_par_mags = wind_mag*np.cos(wind_angle-intended_heading_angles)\n",
    "    adjusted_mag = speed_sigmoid_func_kwarg(signed_wind_par_mags,x_0a,x_0b,L,4.,y_0)\n",
    "    \n",
    "    intended_heading_unit_vectors = np.vstack((\n",
    "        np.cos(intended_heading_angles),np.sin(intended_heading_angles)))\n",
    "    intended_heading_vectors = adjusted_mag*intended_heading_unit_vectors\n",
    "\n",
    "    w_vec = np.array([wind_mag*np.cos(wind_angle),wind_mag*np.sin(wind_angle)])\n",
    "    wind_par_vec = (np.dot(\n",
    "        w_vec,intended_heading_unit_vectors))*intended_heading_unit_vectors\n",
    "\n",
    "    w_perp_vec = w_vec[:,None] - wind_par_vec\n",
    "    heading_final_vec = intended_heading_vectors+c1*wind_par_vec+c2*w_perp_vec\n",
    "    dispersing_speeds = np.sqrt(np.sum(heading_final_vec**2,axis=0))\n",
    "    track_heading_angles = np.arctan2(heading_final_vec[1],heading_final_vec[0])\n",
    "\n",
    "    if plot:\n",
    "        global fig,gs\n",
    "        ax =  fig.add_subplot(gs[:5,10:])   \n",
    "        ax.set_aspect('equal')\n",
    "        ax.spines['left'].set_position('center')\n",
    "        ax.spines['bottom'].set_position('center')\n",
    "\n",
    "        # Eliminate upper and right axes\n",
    "        ax.spines['right'].set_color('none')\n",
    "        ax.spines['top'].set_color('none')\n",
    "\n",
    "        # Show ticks in the left and lower axes only\n",
    "        ax.xaxis.set_ticks_position('bottom')\n",
    "        ax.yaxis.set_ticks_position('left')\n",
    "\n",
    "        ax.spines['bottom'].set_position('zero')\n",
    "        ax.spines['left'].set_position('zero')\n",
    "\n",
    "        ax.set_xlim([-4,7])\n",
    "        ax.set_ylim([-3,6])\n",
    "        ax.set_xlabel('Parallel Windspeed Magnitude')\n",
    "\n",
    "        markersize=0.2\n",
    "\n",
    "        plt.scatter(signed_wind_par_mags,adjusted_mag,\n",
    "        label='intended groundspeed',s=markersize,color='orange')\n",
    "\n",
    "        wind_perp_mags = np.sqrt(np.sum(w_perp_vec**2,axis=0))\n",
    "\n",
    "        plt.scatter(signed_wind_par_mags,wind_perp_mags,\n",
    "        label='wind perp mag',s=markersize,color='blue')\n",
    "\n",
    "        sc = plt.scatter(signed_wind_par_mags,dispersing_speeds,\n",
    "            label='final groundspeed',\n",
    "                cmap='cool',s=markersize,c=dispersing_speeds,vmin=0.,vmax=4.)\n",
    "        lgnd = plt.legend()\n",
    "        lgnd.legendHandles[0]._sizes = [30]\n",
    "        lgnd.legendHandles[1]._sizes = [30]\n",
    "        lgnd.legendHandles[2]._sizes = [30]\n",
    "\n",
    "#         plt.colorbar(sc)\n",
    "    return track_heading_angles,dispersing_speeds\n",
    "\n",
    "\n",
    "def f1(track_heading_angles,source_locations,wind_angle,release_location=np.array([0,0])):\n",
    "    #cone_angle: angle of spread from directly downwind\n",
    "    #ie cone_angle = 10 degrees gives a 20 degree wedge\n",
    "\n",
    "    #Convert track_heading_angles to a list of plume distances for each fly\n",
    "    #(for example, the row for one fly would be [np.nan,500,700,800,np.nan,np.nan,np.nan,np.nan]\n",
    "    # if the fly's future path intersected only plumes 2,3,4, at those distances) )\n",
    "\n",
    "    #Using geometric algorithm here: http://geomalgorithms.com/a05-_intersect-1.html\n",
    "    #referring to diagram under 'Non-Parallel Lines'\n",
    "\n",
    "    #u - vector from source (P_0) downwind\n",
    "    #v - heading vector originating from release location (Q_0)\n",
    "    #w - vector from release location (Q_0) to source (P_0)\n",
    "    #s_1 - distance along u from source (P_0) to intersection point\n",
    "    #t_1 - distance along v from release_location (Q_0) to intersection point\n",
    "\n",
    "\n",
    "    #Turn the track heading angles into unit vectors (v)\n",
    "    fly_unit_vectors = np.vstack((np.cos(track_heading_angles),np.sin(track_heading_angles)))\n",
    "    #shape is (2 x flies)\n",
    "\n",
    "    #Turn the wind angle into a plume unit vector (u)\n",
    "    plume_unit_vector = np.array([np.cos(wind_angle),np.sin(wind_angle)])\n",
    "    #shape is (2)\n",
    "\n",
    "\n",
    "\n",
    "    #Compute vector from release location to each source\n",
    "    w = source_locations - release_location\n",
    "    #shape is (traps x 2)\n",
    "\n",
    "    #Compute s_1 & t_1\n",
    "\n",
    "    #s_1 = (v2*w1-v1*w2)/(v1*u2-v2*u1)\n",
    "    s_1 = (fly_unit_vectors[None,1,:]*w[:,0,None]-fly_unit_vectors[None,0,:]*w[:,1,None])/(\n",
    "        fly_unit_vectors[None,0,:]*plume_unit_vector[None,1,None]-\n",
    "            fly_unit_vectors[None,1,:]*plume_unit_vector[None,0,None])\n",
    "\n",
    "    #t_1 = (u1*w2-u2*w1)/(u1*v2-u2*v1)\n",
    "    t_1 = (plume_unit_vector[None,0,None]*w[:,1,None]-plume_unit_vector[None,1,None]*w[:,0,None])/(\n",
    "        fly_unit_vectors[None,1,:]*plume_unit_vector[None,0,None]-\n",
    "            fly_unit_vectors[None,0,:]*plume_unit_vector[None,1,None])\n",
    "\n",
    "    #collapse extra axis\n",
    "    s_1 = np.squeeze(s_1)\n",
    "\n",
    "    #set all locations where s_1 or t_1 is negative to np.nan\n",
    "\n",
    "    s_1[s_1<0.] = np.nan\n",
    "    s_1[t_1<0.] = np.nan\n",
    "\n",
    "    return s_1.T,t_1.T\n",
    "    #s_1.T - shape is (n flies x n plumes), entry is the distance downwind of that plume\n",
    "    #t_1.T - shape is (n flies x n plumes),\n",
    "        #entry is the distance traveled from the release point to the intersection point\n",
    "\n",
    "def f1_wedge(track_heading_angles,source_locations,wind_angle,\n",
    "    cone_angle,release_location=np.array([0,0])):\n",
    "    #Finds the intersection distances with edges of a wedge,\n",
    "    #rather than a plume line\n",
    "\n",
    "    #cone_angle: angle of spread from directly downwind\n",
    "    #ie cone_angle = 10 degrees gives a 20 degree wedge\n",
    "\n",
    "    #Convert track_heading_angles to a list of plume distances for each fly\n",
    "    #(for example, the row for one fly would be [np.nan,500,700,800,np.nan,np.nan,np.nan,np.nan]\n",
    "    # if the fly's future path intersected only plumes 2,3,4, at those distances) )\n",
    "\n",
    "    #Using geometric algorithm here: http://geomalgorithms.com/a05-_intersect-1.html\n",
    "    #referring to diagram under 'Non-Parallel Lines'\n",
    "\n",
    "    #u - vector from source (P_0) downwind\n",
    "    #v - heading vector originating from release location (Q_0)\n",
    "    #w - vector from release location (Q_0) to source (P_0)\n",
    "    #s_1 - distance along u from source (P_0) to intersection point\n",
    "    #t_1 - distance along v from release_location (Q_0) to intersection point\n",
    "\n",
    "    #Turn the track heading angles into unit vectors (v)\n",
    "    fly_unit_vectors = np.vstack((np.cos(track_heading_angles),np.sin(track_heading_angles)))\n",
    "\n",
    "    #shape is (2 x flies)\n",
    "    #Turn (wind_angle-cone_angle,wind_angle+cone_angle) into a plume unit vector (u)\n",
    "    plume_unit_vector = np.array(\n",
    "        [[np.cos(wind_angle-cone_angle),np.sin(wind_angle-cone_angle)],\n",
    "        [np.cos(wind_angle+cone_angle),np.sin(wind_angle+cone_angle)]])\n",
    "\n",
    "    #shape is (cone sides (2) x 2)\n",
    "\n",
    "    #Compute vector from release location to each source\n",
    "    w = source_locations - release_location\n",
    "    #shape is (traps x 2)\n",
    "\n",
    "    #Compute s_1 & t_1\n",
    "\n",
    "    #In the none-wedge version, the below have shape (traps x 2 x flies), and\n",
    "    #then collapse to (traps x flies) after the computation.\n",
    "\n",
    "    #For wedges, add an extra dimension to the front\n",
    "    #for 'which cone side' (-cone_angle,cone_angle)\n",
    "\n",
    "    #new shape is (2 x traps x 2 x flies)--> collapse to (2 x traps x flies)\n",
    "\n",
    "    #s_1 = (v2*w1-v1*w2)/(v1*u2-v2*u1)\n",
    "    s_1 = (fly_unit_vectors[None,None,1,:]*w[None,:,0,None]\n",
    "        -fly_unit_vectors[None,None,0,:]*w[None,:,1,None])/(\n",
    "            fly_unit_vectors[None,None,0,:]*plume_unit_vector[:,None,1,None]-\n",
    "                fly_unit_vectors[None,None,1,:]*plume_unit_vector[:,None,0,None])\n",
    "\n",
    "    #t_1 = (u1*w2-u2*w1)/(u1*v2-u2*v1)\n",
    "    t_1 = (plume_unit_vector[:,None,0,None]*w[None,:,1,None]\n",
    "        -plume_unit_vector[:,None,1,None]*w[None,:,0,None])/(\n",
    "            fly_unit_vectors[None,None,1,:]*plume_unit_vector[:,None,0,None]-\n",
    "                fly_unit_vectors[None,None,0,:]*plume_unit_vector[:,None,1,None])\n",
    "\n",
    "    #collapse extra axis\n",
    "    s_1 = np.squeeze(s_1) #now shape is (2 x traps x flies)\n",
    "\n",
    "    #set all intersection values where s_1 or t_1 is negative to np.inf\n",
    "\n",
    "    # fig10 = plt.figure(10)\n",
    "    #\n",
    "    # to_plot = np.concatenate((\n",
    "    #     track_heading_angles[(s_1[1,3,:]>0.)],#&(s_1[1,3,:]>0.)],\n",
    "    #     track_heading_angles[(s_1[1,3,:]>0.)]))  #&(s_1[1,3,:]>0.)]))\n",
    "    #\n",
    "    # n,bins = np.histogram(to_plot%(2*np.pi),bins=np.linspace(0,2*np.pi,50))\n",
    "    #\n",
    "    # try:\n",
    "    #     global hist\n",
    "    #     hist.set_ydata(n)\n",
    "    #     hist.set_xdata(bins[:-1])\n",
    "    # except(NameError):\n",
    "    #     hist, = plt.plot(bins[:-1],n,'o')\n",
    "    #\n",
    "    # fig10.canvas.draw_idle()\n",
    "\n",
    "    inds = (s_1>0.)&(t_1>0.)\n",
    "    s_1[~inds] = np.inf\n",
    "    t_1[~inds] = np.inf\n",
    "\n",
    "    #For some reason the below does NOT have the identical effect as the above\n",
    "    #for wide cone angles (>pi/8)\n",
    "\n",
    "    # s_1[s_1<0.] = np.inf\n",
    "    # t_1[s_1<0.] = np.inf\n",
    "    # # s_1[t_1<0.] = np.inf\n",
    "    # t_1[t_1<0.] = np.inf\n",
    "\n",
    "    #collapse s_1 and t_1 along the 'which cone side' dimension\n",
    "    #by selecting the cone side with the shorter (finite, positive)\n",
    "    #t_1 (distance from release_location to intersection point)\n",
    "\n",
    "    #first collapse t_1\n",
    "    collapsed_t_1 = np.min(t_1,axis=0)\n",
    "\n",
    "    where_no_intersections = np.isinf(collapsed_t_1) #value true for fly-trap pairs with no intersection\n",
    "\n",
    "    #make cone_side_mask, shape (traps x flies) which is a mask (entries (0,1))\n",
    "    #of which cone sides were intersected (nan if neither)\n",
    "    cone_side_mask_rough = (t_1==np.min(t_1,axis=0))\n",
    "    cone_side_mask = np.full(np.shape(collapsed_t_1),np.nan)\n",
    "    cone_side_mask[cone_side_mask_rough[0]] = 0.\n",
    "    cone_side_mask[cone_side_mask_rough[1]] = 1.\n",
    "    cone_side_mask[where_no_intersections] = np.nan\n",
    "\n",
    "    #now collapse s_1, the distance from source (P_0) to intersection point,\n",
    "    #by selecting the assigned side (or none of them)\n",
    "\n",
    "    collapsed_s_1 = np.full(np.shape(collapsed_t_1),np.nan)\n",
    "\n",
    "    collapsed_s_1[(cone_side_mask==0.)] = s_1[\n",
    "        0,np.where(cone_side_mask==0)[0],np.where(cone_side_mask==0)[1]]\n",
    "    collapsed_s_1[(cone_side_mask==1.)] = s_1[\n",
    "        1,np.where(cone_side_mask==1.)[0],np.where(cone_side_mask==1.)[1]]\n",
    "\n",
    "    #Check that the shape of s_1 and t_1 are as expected\n",
    "\n",
    "    return collapsed_s_1.T,collapsed_t_1.T\n",
    "    #s_1.T - shape is (n flies x n plumes), entry is the distance downwind of that plume\n",
    "    #t_1.T - shape is (n flies x n plumes),\n",
    "        #entry is the distance traveled from the release point to the intersection point\n",
    "\n",
    "\n",
    "def f1_inside_wedge(track_heading_angles,source_locations,wind_angle,\n",
    "        cone_angle,release_location=np.array([0,0])):\n",
    "\n",
    "        #Different version of f1_wedge, in which each fly draws\n",
    "        #an angle between the wedge edges (uniformly at random)\n",
    "        #which determines the location in the plume it intersects.\n",
    "\n",
    "        #That is, for each plume it intersects the fly is assigned\n",
    "        #a plume angle within the angle range specified by the wedge.\n",
    "\n",
    "        #Using geometric algorithm here: http://geomalgorithms.com/a05-_intersect-1.html\n",
    "        #referring to diagram under 'Non-Parallel Lines'\n",
    "\n",
    "        #u - vector from source (P_0) downwind\n",
    "        #v - heading vector originating from release location (Q_0)\n",
    "        #w - vector from release location (Q_0) to source (P_0)\n",
    "        #s_1 - distance along u from source (P_0) to intersection point\n",
    "        #t_1 - distance along v from release_location (Q_0) to intersection point\n",
    "\n",
    "        #Turn the track heading angles into unit vectors (v)\n",
    "        fly_unit_vectors = np.vstack((np.cos(track_heading_angles),np.sin(track_heading_angles)))\n",
    "\n",
    "        #shape is (2 x flies)\n",
    "        #Turn (wind_angle-cone_angle,wind_angle+cone_angle) into a plume unit vector (u)\n",
    "        plume_unit_vector = np.array(\n",
    "            [[np.cos(wind_angle-cone_angle),np.sin(wind_angle-cone_angle)],\n",
    "            [np.cos(wind_angle+cone_angle),np.sin(wind_angle+cone_angle)]])\n",
    "\n",
    "        #shape is (cone sides (2) x 2)\n",
    "\n",
    "        #Compute vector from release location to each source\n",
    "        w = source_locations - release_location\n",
    "        #shape is (traps x 2)\n",
    "\n",
    "        #Compute s_1 & t_1\n",
    "\n",
    "        #In the none-wedge version, the below have shape (traps x 2 x flies), and\n",
    "        #then collapse to (traps x flies) after the computation.\n",
    "\n",
    "        #For wedges, add an extra dimension to the front\n",
    "        #for 'which cone side' (-cone_angle,cone_angle)\n",
    "\n",
    "        #new shape is (2 x traps x 2 x flies)--> collapse to (2 x traps x flies)\n",
    "\n",
    "        #s_1 = (v2*w1-v1*w2)/(v1*u2-v2*u1)\n",
    "        s_1 = (fly_unit_vectors[None,None,1,:]*w[None,:,0,None]\n",
    "            -fly_unit_vectors[None,None,0,:]*w[None,:,1,None])/(\n",
    "                fly_unit_vectors[None,None,0,:]*plume_unit_vector[:,None,1,None]-\n",
    "                    fly_unit_vectors[None,None,1,:]*plume_unit_vector[:,None,0,None])\n",
    "\n",
    "        #t_1 = (u1*w2-u2*w1)/(u1*v2-u2*v1)\n",
    "        t_1 = (plume_unit_vector[:,None,0,None]*w[None,:,1,None]\n",
    "            -plume_unit_vector[:,None,1,None]*w[None,:,0,None])/(\n",
    "                fly_unit_vectors[None,None,1,:]*plume_unit_vector[:,None,0,None]-\n",
    "                    fly_unit_vectors[None,None,0,:]*plume_unit_vector[:,None,1,None])\n",
    "\n",
    "        #collapse extra axis\n",
    "        s_1 = np.squeeze(s_1) #now shape is (2 x traps x flies)\n",
    "\n",
    "\n",
    "        inds = (s_1>0.)&(t_1>0.)\n",
    "        s_1[~inds] = np.inf\n",
    "        t_1[~inds] = np.inf\n",
    "\n",
    "        #first collapse t_1\n",
    "        collapsed_t_1 = np.min(t_1,axis=0)\n",
    "\n",
    "        where_no_intersections = np.isinf(collapsed_t_1) #value true for fly-trap pairs with no intersection\n",
    "\n",
    "        #make cone_side_mask, shape (traps x flies) which is a mask (entries (0,1))\n",
    "        #of which cone sides were intersected (nan if neither)\n",
    "        cone_side_mask_rough = (t_1==np.min(t_1,axis=0))\n",
    "        cone_side_mask = np.full(np.shape(collapsed_t_1),np.nan)\n",
    "        cone_side_mask[cone_side_mask_rough[0]] = 0.\n",
    "        cone_side_mask[cone_side_mask_rough[1]] = 1.\n",
    "        cone_side_mask[where_no_intersections] = np.nan\n",
    "\n",
    "        #now collapse s_1, the distance from source (P_0) to intersection point,\n",
    "        #by selecting the assigned side (or none of them)\n",
    "\n",
    "        collapsed_s_1 = np.full(np.shape(collapsed_t_1),np.nan)\n",
    "\n",
    "        collapsed_s_1[(cone_side_mask==0.)] = s_1[\n",
    "            0,np.where(cone_side_mask==0)[0],np.where(cone_side_mask==0)[1]]\n",
    "        collapsed_s_1[(cone_side_mask==1.)] = s_1[\n",
    "            1,np.where(cone_side_mask==1.)[0],np.where(cone_side_mask==1.)[1]]\n",
    "\n",
    "        #Check that the shape of s_1 and t_1 are as expected\n",
    "\n",
    "        return collapsed_s_1.T,collapsed_t_1.T\n",
    "        #s_1.T - shape is (n flies x n plumes), entry is the distance downwind of that plume\n",
    "        #t_1.T - shape is (n flies x n plumes),\n",
    "            #entry is the distance traveled from the release point to the intersection point)\n",
    "\n",
    "def logistic_1d(x,K,x_0,wind_angle):\n",
    "        Q = 1.\n",
    "        B = 0.05\n",
    "        output = np.zeros_like(x)\n",
    "        output[x>0.]=-K+K/((1.+Q*np.exp(-B*(x[x>0.]-x_0))))\n",
    "        return output\n",
    "\n",
    "def f2(intersection_distances,K,x_0,source_locations,wind_angle):\n",
    "    #Convert intersection_distances to probabilities\n",
    "    #using a specified distance-probability function\n",
    "\n",
    "    success_probabilities = logistic_1d(intersection_distances,K,\n",
    "                                            x_0,wind_angle)\n",
    "\n",
    "    return success_probabilities\n",
    "\n",
    "def f3(success_probabilities,dispersal_distances):\n",
    "\n",
    "    #Use the success_probabilities to determine (with a random draw)\n",
    "    #which plume each fly will chase, if any\n",
    "\n",
    "\n",
    "    #the input is shape (n flies x n plumes)\n",
    "\n",
    "    #First, do a random draw for each plume it will intersect\n",
    "\n",
    "    #Choose among the positive results of the draw (among possible plumes)\n",
    "    #by **choosing the plume that drew a 1 which has the shortest dispersal\n",
    "    #distance for the fly**\n",
    "\n",
    "    # print('success_probabilities')\n",
    "    # print(np.array2string(success_probabilities,precision=2))\n",
    "\n",
    "    draws = np.random.binomial(1,success_probabilities,\n",
    "        size=np.shape(success_probabilities)).astype(bool)\n",
    "\n",
    "\n",
    "\n",
    "    distances_with_positive_draws = np.full_like(success_probabilities,np.nan)\n",
    "\n",
    "    distances_with_positive_draws[draws] = dispersal_distances[draws]\n",
    "\n",
    "    #Check that the above value either has positive floats or nan's\n",
    "    assert(np.sum(distances_with_positive_draws[~np.isnan(\n",
    "        distances_with_positive_draws)]<0.)<1)\n",
    "\n",
    "    #Take the inverse of each element to reverse the order and avoid 0 being the min\n",
    "    distances_with_positive_draws[~np.isnan(\n",
    "        distances_with_positive_draws)] = 1./(\n",
    "            distances_with_positive_draws[~np.isnan(distances_with_positive_draws)])\n",
    "\n",
    "    #Set the nan values to 0.\n",
    "    distances_with_positive_draws[np.isnan(distances_with_positive_draws)] = 0.\n",
    "\n",
    "\n",
    "    plume_assignments = distances_with_positive_draws.argmax(axis=1).astype(float)\n",
    "\n",
    "    plume_assignments[np.sum(draws,axis=1)==0] = np.nan\n",
    "\n",
    "    # print('plume_assignments')\n",
    "    # print(np.array2string(plume_assignments,precision=3))\n",
    "\n",
    "    #output is of shape (flies) which has for each fly number 0-7 of the traps\n",
    "    #it goes to, or np.nan if it does not go to a trap\n",
    "    # print('--------')\n",
    "    # print(np.sum(plume_assignments==0))\n",
    "    # print(np.sum(plume_assignments==3))\n",
    "\n",
    "\n",
    "    return plume_assignments\n",
    "\n",
    "def f4(plume_assignments,dispersal_distances,dispersing_speeds):\n",
    "    #Now that we know which plume each fly is chasing,\n",
    "    #use the intersection distances and plume ids to compute the time\n",
    "    #it took for each fly to travel from the release site to\n",
    "    #the plume it ended up detecting\n",
    "\n",
    "    release_to_chosen_plume_distances = np.full(len(plume_assignments),np.nan)\n",
    "\n",
    "    mask = ~np.isnan(plume_assignments)\n",
    "    cols = plume_assignments[mask].astype(int)\n",
    "    rows = np.where(mask)\n",
    "\n",
    "    release_to_chosen_plume_distances[mask] = dispersal_distances[rows,cols].flatten()\n",
    "\n",
    "    dispersal_travel_times = (release_to_chosen_plume_distances/dispersing_speeds)\n",
    "\n",
    "    return dispersal_travel_times,release_to_chosen_plume_distances\n",
    "\n",
    "def f5(plume_assignments,dispersal_travel_times,\\\n",
    "    intersection_distances,fly_speed,release_times):\n",
    "    #Use the intersection distances and plume ids to compute the time\n",
    "    #that each fly arrived at the source of the plume it successfully chases\n",
    "    intersection_distances_chosen_plume = np.full(len(plume_assignments),np.nan)\n",
    "\n",
    "    mask = ~np.isnan(plume_assignments)\n",
    "    cols = plume_assignments[mask].astype(int)\n",
    "    rows = np.where(mask)[0]\n",
    "\n",
    "    intersection_distances_chosen_plume[mask] = intersection_distances[rows,cols].flatten()\n",
    "\n",
    "    chasing_times = (intersection_distances_chosen_plume/fly_speed)\n",
    "\n",
    "    arrival_times = release_times+dispersal_travel_times+chasing_times\n",
    "\n",
    "    return arrival_times[~np.isnan(arrival_times)],chasing_times,rows,cols\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UpdatingVPatch(object):\n",
    "    def __init__(self,x_0,width):\n",
    "        self.rectangle = plt.Rectangle((x_0,0),width,1.,alpha=0.5,color='orange')\n",
    "    def update(self,new_x_0,new_width):\n",
    "        self.rectangle.set_x(new_x_0)\n",
    "        self.rectangle.set_width(new_width)\n",
    "        # return self.rectangle\n",
    "\n",
    "def plot_wedges(source_pos,wind_angle,cone_angle):\n",
    "    length = 2000\n",
    "    first_arms = np.array(\n",
    "        (source_pos[:,0]+length*np.cos(wind_angle+cone_angle),\n",
    "        source_pos[:,1]+length*np.sin(wind_angle+cone_angle))).T\n",
    "    second_arms = np.array(\n",
    "        (source_pos[:,0]+length*np.cos(wind_angle-cone_angle),\n",
    "        source_pos[:,1]+length*np.sin(wind_angle-cone_angle))).T\n",
    "    #shape traps x 2\n",
    "    merged = np.stack((first_arms,source_pos,second_arms))\n",
    "    return merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5511591f63cf43cabd2bdfcdad98de45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Box(children=(Box(children=(FloatSlider(value=0.0, continuous_update=False, description='Wind Speed', max=2.0,…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f67d695f97f4d23a74187aebdc4c0e9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Constants that don't change with drag bars\n",
    "num_flies = 2000\n",
    "fly_speed = 1.6\n",
    "\n",
    "number_sources = 8\n",
    "radius_sources = 1000.0\n",
    "source_locations, _ = 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",
    "\n",
    "style = {'description_width': 'initial'}\n",
    "\n",
    "\n",
    "\n",
    "windmag_slider = widgets.FloatSlider(\n",
    "    value=0.,\n",
    "    min=0,\n",
    "    max=2.0,\n",
    "    step=0.1,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    description='wind_mag',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    "    style=style\n",
    ")\n",
    "\n",
    "cone_angle_slider = widgets.FloatSlider(\n",
    "    value=np.degrees(initial_cone_angle),\n",
    "    min=0.,\n",
    "    max=25.0,\n",
    "    step=1.,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    description='cone_angle',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    "    style=style\n",
    ")\n",
    "\n",
    "x_0_slider = widgets.FloatSlider(\n",
    "    value=300.,\n",
    "    min=150.,\n",
    "    max=1000.0,\n",
    "    step=10.,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    description='x_0',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    "    style=style\n",
    ")\n",
    "\n",
    "radio = widgets.RadioButtons(options=[('Headings',0), ('Intersections',1)],disabled=False,description='toggle')                                 \n",
    "\n",
    "\n",
    "\n",
    "def slider(start,stop,step,init,description):\n",
    "    return widgets.FloatSlider(\n",
    "    value=init,\n",
    "    min=start,\n",
    "    max=stop,\n",
    "    step=step,\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    description=description,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.2f',\n",
    "    style=style\n",
    ")\n",
    "\n",
    "# def f(wind_mag,toggle,cone_angle,K,x_0,x_0a,x_0b,L,y_0,t)\n",
    "def f(wind_mag,toggle,cone_angle,x_0,x_0a,x_0b,t):\n",
    "\n",
    "    c1=0\n",
    "    c2=1\n",
    "    beta=1\n",
    "    \n",
    "    #plot scaffolding\n",
    "    global fig,gs\n",
    "    gs = GridSpec(9, 16)\n",
    "    fig = plt.figure(1,figsize=(17,10))\n",
    "    \n",
    "    \n",
    "    #----OBTAIN NEW VALUES--------\n",
    "    wind_angle = 7*scipy.pi/8.\n",
    "    cone_angle = np.radians(cone_angle)\n",
    "\n",
    "    release_times=np.random.exponential(60*beta,num_flies)\n",
    "    \n",
    "    plt.figure(1)\n",
    "               \n",
    "\n",
    "    K = -1*0.4\n",
    "    y_0 = 1.6\n",
    "    L = 0.1\n",
    "    \n",
    "    \n",
    "    #------RECOMPUTE OUTPUTS------\n",
    "    track_heading_angles,dispersing_speeds = f0(intended_heading_angles,wind_mag,\n",
    "        wind_angle,x_0a,x_0b,L,y_0,plot=False,c1=c1,c2=c2)\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",
    "    #------\"FIGURE\" 0: groundspeed curve\n",
    "    ax =  fig.add_subplot(gs[2:7,10:])   \n",
    "    ax.set_aspect('equal')\n",
    "    ax.spines['left'].set_position('center')\n",
    "    ax.spines['bottom'].set_position('center')\n",
    "\n",
    "    # Eliminate upper and right axes\n",
    "    ax.spines['right'].set_color('none')\n",
    "    ax.spines['top'].set_color('none')\n",
    "\n",
    "    # Show ticks in the left and lower axes only\n",
    "    ax.xaxis.set_ticks_position('bottom')\n",
    "    ax.yaxis.set_ticks_position('left')\n",
    "\n",
    "    ax.spines['bottom'].set_position('zero')\n",
    "    ax.spines['left'].set_position('zero')\n",
    "\n",
    "    ax.set_xlim([-1,7])\n",
    "    ax.set_ylim([-1,6])\n",
    "    ax.text(0.7,-0.05,'(Parallel) Windspeed',transform=ax.transAxes)\n",
    "    ax.text(-0.05,0.7,'Groundspeed',transform=ax.transAxes,rotation=90)\n",
    "\n",
    "    inputs = np.linspace(-4,7,100)\n",
    "    plt.plot(inputs,speed_sigmoid_func_kwarg(inputs,x_0a,x_0b,L,4.,y_0),color='orange')\n",
    "    plt.plot(x_0a*np.ones(100),np.linspace(-1,6,100),color='navy')\n",
    "    ax.text(x_0a-0.5,3.,'x_0a',horizontalalignment='center',transform=ax.transData,color='navy')\n",
    "\n",
    "    plt.plot(x_0b*np.ones(100),np.linspace(-1,6,100),color='navy')\n",
    "    ax.text(x_0b+0.5,3.,'x_0b',horizontalalignment='center',transform=ax.transData,color='navy')\n",
    "\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",
    "        \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",
    "        ax.scatter(x,y,marker='x',s=50,c='k')\n",
    "    \n",
    "    ax.set_aspect('equal')\n",
    "          \n",
    "#     ax.set_xticks([])\n",
    "#     ax.set_yticks([])\n",
    "    \n",
    "\n",
    "    \n",
    "    if toggle==1:\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",
    "    if toggle==0: \n",
    "        time = t*60.\n",
    "        mag = time*dispersing_speeds \n",
    "        x,y = mag*np.cos(track_heading_angles), \\\n",
    "            mag*np.sin(track_heading_angles)\n",
    "     \n",
    "    if toggle<2:\n",
    "        fly_dots = plt.scatter(x,y,alpha=0.5,c=dispersing_speeds,vmin=0.,vmax=4.,cmap='cool')\n",
    "        n = matplotlib.colors.Normalize(vmin = 0., vmax = 4.)\n",
    "        m = matplotlib.cm.ScalarMappable(norm=n, cmap='cool')\n",
    "        fly_dots.set_facecolor(m.to_rgba(dispersing_speeds))\n",
    "        fly_dots.set_clim(vmin=0., vmax=4.)\n",
    "                \n",
    "        axins = inset_axes(ax,\n",
    "            width=\"50%\",  # width = 50% of parent_bbox width\n",
    "            height=\"5%\",  # height : 5%\n",
    "            bbox_to_anchor=(-0.25,-0.55, 1, 0.5),\n",
    "            bbox_transform=ax.transAxes,\n",
    "            borderpad=0)\n",
    "        \n",
    "        fig.colorbar(fly_dots,cax=axins,pad=0.2,orientation='horizontal')         \n",
    "        \n",
    "        ax.text(0.5,-0.2,'Dispersal Speed (m/s)',horizontalalignment='center',\n",
    "            verticalalignment='center',fontsize=12,transform=ax.transAxes)\n",
    "\n",
    "    \n",
    "\n",
    "    elif toggle==2:\n",
    "#         sign = np.sign(dispersing_speeds)\n",
    "#         time = sign*t*60.\n",
    "        rad = 500.\n",
    "        n_bins = 2**6\n",
    "        bins = np.linspace(0,2*np.pi,n_bins)\n",
    "        \n",
    "        \n",
    "        n,bins = np.histogram(track_heading_angles%(2*np.pi),bins=bins)\n",
    "        \n",
    "        center_rs = 0*rad*np.ones_like(bins[:-1])\n",
    "        n_scale = rad/np.max(n)\n",
    "        outside_rs = center_rs + n_scale*n\n",
    "        inside_rs = center_rs - n_scale*n\n",
    "        \n",
    "#         plt.plot(center_rs*np.cos(bins[:-1]),center_rs*np.sin(bins[:-1]),color='orange')\n",
    "        plt.plot(outside_rs*np.cos(bins[:-1]),outside_rs*np.sin(bins[:-1]),color='orange')\n",
    "#         plt.plot(inside_rs*np.cos(bins[:-1]),inside_rs*np.sin(bins[:-1]),color='orange')\n",
    "        \n",
    "    \n",
    "        rad2 = 700\n",
    "        density_dots = plt.scatter(rad2*np.cos(bins[:-1]),rad2*np.sin(bins[:-1]),\n",
    "                                   c=n/(num_flies/n_bins),vmin=0,vmax=(np.max(n))/(num_flies/n_bins))\n",
    "\n",
    "        axins = inset_axes(ax,\n",
    "            width=\"50%\",  # width = 50% of parent_bbox width\n",
    "            height=\"5%\",  # height : 5%\n",
    "            loc='lower center',\n",
    "            bbox_to_anchor=(-0.25, -0.55, 1, .5),\n",
    "            bbox_transform=ax.transAxes,\n",
    "            borderpad=0)\n",
    "\n",
    "        \n",
    "        fig.colorbar(density_dots,cax=axins,pad=0.4,orientation='horizontal')  \n",
    "        \n",
    "        ax.text(0.5,-0.3,'Directional Density',horizontalalignment='center',\n",
    "         verticalalignment='center',fontsize=12,transform=ax.transAxes)\n",
    "\n",
    "\n",
    "        \n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "    \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 = num_flies/50.*np.ones(8)\n",
    "    for i in range(8):\n",
    "\n",
    "        row = sim_reorder[i]-1\n",
    "        ax =  fig.add_subplot(gs[row,5:9])    \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",
    "            patch_object = UpdatingVPatch(50,50)\n",
    "            ax.add_patch(patch_object.rectangle)\n",
    "            cdf_patches.append(patch_object)\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(num_flies/50.,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=15,\n",
    "            horizontalalignment='center',verticalalignment='center')\n",
    "        if sim_reorder[i]-1==7:\n",
    "            ax.set_ylabel('Flies Arrived',x=0.5,horizontalalignment='center',fontsize=12)\n",
    "            plt.tick_params(axis='y', which='major', labelsize=15)\n",
    "            ax.set_xlabel('Time (min)',x=0.5,horizontalalignment='center',fontsize=12)\n",
    "            plt.tick_params(axis='x', which='major', labelsize=15)\n",
    "\n",
    "        else:\n",
    "            ax.set_xticklabels('')\n",
    "        axes.append(ax)\n",
    "   \n",
    "    for i,ax in enumerate(axes):\n",
    "        ax.set_yticks((num_flies/100)*np.array([0.,1.,2.,3.,4.]))\n",
    "        ax.set_ylim([0,np.max(new_maxes)])\n",
    "        patch_object = cdf_patches[i]\n",
    "        try:\n",
    "            patch_object.rectangle.set_height(ax.get_ylim()[1])\n",
    "    \n",
    "        except(ValueError):\n",
    "            pass\n",
    "    \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 = 1.5/100.*num_flies\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",
    "    trap_patches = [plt.Circle(center, size) for center, size, i in zip(\n",
    "        trap_locs_2d, radius_scale*trap_counts/np.max(new_maxes),range(8))]\n",
    "\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('winter')\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",
    "        \n",
    "    axins = inset_axes(ax,\n",
    "                    width=\"50%\",  # width = 50% of parent_bbox width\n",
    "                    height=\"5%\",  # height : 5%\n",
    "                    loc='lower center',\n",
    "                    bbox_to_anchor=(0.0,-0.2, 1, 0.5),\n",
    "                    bbox_transform=ax.transAxes,\n",
    "  )\n",
    "\n",
    "    fig.colorbar(coll, cax=axins,pad=0.2,orientation='horizontal')\n",
    "\n",
    "    \n",
    "    ax.text(0.5,-0.3,'First Arrival Time (min)',horizontalalignment='center',\n",
    "        verticalalignment='center',fontsize=12,transform=ax.transAxes)\n",
    "    \n",
    "    \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",
    "x_0a =slider(-2,6,0.1,-0.4,'x_0a')\n",
    "x_0b =slider(-2,6,0.1,1.45,'x_0b')\n",
    "t=slider(0,30,1,5,'t')\n",
    "#radio = slider(0,30,1,5,'toggle')\n",
    "    \n",
    "sliders = windmag_slider, cone_angle_slider, x_0_slider,radio,x_0a,x_0b,t\n",
    "slider_names = [slider.description for slider in sliders]\n",
    "param_dict =  dict(zip(slider_names,sliders))\n",
    "out = widgets.interactive_output(f, param_dict)\n",
    "\n",
    "#Rename the slider.description for visual clarity\n",
    "slider_descriptions = ['Wind Speed','Plume Thickness','Plume Detection Range','x_0a','x_0b','Time']\n",
    "\n",
    "for slider,slider_description in zip([windmag_slider, cone_angle_slider, x_0_slider,x_0a,x_0b,t],slider_descriptions):\n",
    "    slider.description = slider_description\n",
    "    \n",
    "radio.description = ''\n",
    "    \n",
    "items = [Box([slider]) for slider in sliders]\n",
    "ui = Box(items, layout=Layout(\n",
    "    display='flex', flex_flow='column',\n",
    "    border='solid 2px', align_items='stretch',\n",
    "    width='50%'))\n",
    "\n",
    "\n",
    "display(ui,out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}