Gill-Chiel-eNeuro-2020-code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gill-Chiel-eNeuro-2020-code\n",
    "\n",
    "Analyses, statistics, and figures published in\n",
    "\n",
    "> Gill, J. P., & Chiel, H. J. (2020). Rapid Adaptation to Changing Mechanical Load by Ordered Recruitment of Identified Motor Neurons. eNeuro, 7(3). https://doi.org/10.1523/ENEURO.0016-20.2020\n",
    "\n",
    "Figures and statistics are exported as files. [Click here to open the output directory.](output/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of Contents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- [Preamble](#Preamble)\n",
    "  - [Import Packages](#Import-Packages)\n",
    "  - [Plot Settings](#Plot-Settings)\n",
    "  - [Data Parameters](#Data-Parameters)\n",
    "  - [Helper Functions](#Helper-Functions)\n",
    "  - [Plotting Functions](#Plotting-Functions)\n",
    "\n",
    "\n",
    "- [Crunch the Numbers](#Crunch-the-Numbers)\n",
    "  - [Download the Data](#Download-the-Data)\n",
    "  - [Import and Process the Data](#Import-and-Process-the-Data)\n",
    "  - [Sanity Checks](#🤪-Sanity-Checks)\n",
    "\n",
    "\n",
    "- [Figures](#Figures)\n",
    "  - [Figure 1](#[FIGURE-1])\n",
    "    - [Figure 1A](#🐌-Figure-1A)\n",
    "    - [Figures 1B & 1C](#🐌-Figures-1B-&-1C)\n",
    "  - [Figure 2](#[FIGURE-2])\n",
    "    - [Figure 2A](#🐌-Figure-2A)\n",
    "    - [Figure 2B](#🐌-Figure-2B)\n",
    "    - [Figure 2C](#🐌-Figure-2C)\n",
    "    - [Figure 2D](#🐌-Figure-2D)\n",
    "    - [Figure 2E](#🐌-Figure-2E)\n",
    "  - [Figure 3](#[FIGURE-3])\n",
    "    - [Figure 3A](#🐌-Figure-3A)\n",
    "    - [Figure 3B](#🐌-Figure-3B)\n",
    "    - [Figure 3C](#🐌-Figure-3C)\n",
    "    - [Figure 3D](#🐌-Figure-3D)\n",
    "    - [Figure 3E](#🐌-Figure-3E)\n",
    "  - [Figure 4](#[FIGURE-4])\n",
    "    - [Figure 4 Statistics](#🐌-Figure-4-Statistics)\n",
    "    - [Figure 4A](#🐌-Figure-4A)\n",
    "    - [Figure 4B](#🐌-Figure-4B)\n",
    "    - [Figure 4C](#🐌-Figure-4C)\n",
    "    - [Figure 4D](#🐌-Figure-4D)\n",
    "    - [Figure 4E](#🐌-Figure-4E)\n",
    "    - [Figure 4F](#🐌-Figure-4F)\n",
    "    - [Other frequencies (not plotted in manuscript)](#🐌-Other-frequencies-(not-plotted-in-manuscript))\n",
    "  - [Figure 5](#[FIGURE-5])\n",
    "    - [Figure 5A](#🐌-Figure-5A)\n",
    "    - [Figure 5B](#🐌-Figure-5B)\n",
    "    - [Figure 5C](#🐌-Figure-5C)\n",
    "    - [Figure 5D](#🐌-Figure-5D)\n",
    "    - [Figure 5E](#🐌-Figure-5E)\n",
    "    - [Figure 5F](#🐌-Figure-5F)\n",
    "\n",
    "\n",
    "- [Statistical Tables](#Statistical-Tables)\n",
    "  - [Table 2](#[TABLE-2])\n",
    "  - [Table 3](#[TABLE-3])\n",
    "  - [Table 4](#[TABLE-4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preamble"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "## Import Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import datetime\n",
    "from tqdm.notebook import tqdm\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import quantities as pq\n",
    "import elephant\n",
    "import pandas as pd\n",
    "import statsmodels.api as sm\n",
    "import neurotic\n",
    "from neurotic.datasets.data import _detect_spikes, _find_bursts\n",
    "from modules.utils import BehaviorsDataFrame, DownsampleNeoSignal\n",
    "from modules import r_stats\n",
    "from modules.plot_utils import add_scalebar, solve_figure_horizontal_dimensions, solve_figure_vertical_dimensions, plot_signals_with_scalebars\n",
    "\n",
    "pq.markup.config.use_unicode = True  # allow symbols like mu for micro in output\n",
    "pq.mN = pq.UnitQuantity('millinewton', pq.N/1e3, symbol = 'mN');  # define millinewton\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as mcolors\n",
    "import matplotlib.patches as patches\n",
    "import matplotlib.lines as mlines\n",
    "from matplotlib.markers import CARETLEFT, CARETRIGHT, CARETUP, CARETDOWN, CARETUPBASE\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "# don't warn about invalid comparisons to NaN\n",
    "np.seterr(invalid='ignore')\n",
    "\n",
    "# np.nanmax raises a warning if all values are NaN and returns NaN, which is the behavior we want\n",
    "warnings.filterwarnings('ignore', message='All-NaN slice encountered')\n",
    "\n",
    "# elephant.statistics.instantaneous_rate always complains about negative values\n",
    "warnings.filterwarnings('ignore', message='Instantaneous firing rate approximation contains '\n",
    "                                          'negative values, possibly caused due to machine '\n",
    "                                          'precision errors')\n",
    "\n",
    "# with matplotlib>=3.1 and seaborn<=0.9.0, deprecation warnings are raised\n",
    "# whenever tick marks are placed on the right axis but not the left\n",
    "from matplotlib.cbook import MatplotlibDeprecationWarning\n",
    "warnings.simplefilter('ignore', MatplotlibDeprecationWarning)\n",
    "\n",
    "# don't complain about opening too many figures\n",
    "plt.rcParams.update({'figure.max_open_warning': 0})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# make figures interactive and inline\n",
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "## Plot Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_dir = 'output'\n",
    "if not os.path.exists(export_dir):\n",
    "    os.mkdir(export_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# general plot settings\n",
    "sns.set(\n",
    "    style = 'ticks',\n",
    "    font_scale = 1,\n",
    "    font = ['Palatino Linotype', 'serif'],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "unit_colors = {\n",
    "    'B38':       '#EFBF46', # yellow\n",
    "    'I2':        '#DC5151', # red\n",
    "    'B8a/b':     '#DA8BC3', # pink\n",
    "    'B6/B9':     '#64B5CD', # light blue\n",
    "    'B3/B6/B9':  '#5A9BC5', # medium blue\n",
    "    'B3':        '#4F80BD', # dark blue\n",
    "    'B4/B5':     '#00A86B', # jade green\n",
    "    'Force':     '#1A1A1A', # very dark gray\n",
    "}\n",
    "force_colors = {\n",
    "    'shoulder':     unit_colors['B38'],\n",
    "    'dip':          unit_colors['I2'],\n",
    "    'initial rise': unit_colors['B8a/b'],\n",
    "    'rise':         unit_colors['B8a/b'],\n",
    "    'plateau':      unit_colors['B3/B6/B9'],\n",
    "    'drop':         '#808080', # medium gray\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# display the selected unit colors\n",
    "sns.palplot(unit_colors.values())\n",
    "with plt.rc_context({'font.weight': 'bold'}):\n",
    "    for i, (unit, color) in enumerate(unit_colors.items()):\n",
    "        plt.gca().annotate(f'{unit}\\n{color[1:]}', xy=(i, 0), ha='center', va='center', color='w')\n",
    "plt.axis('off')\n",
    "plt.subplots_adjust(0, 0, 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print hex codes for selected unit colors\n",
    "for unit, color in unit_colors.items():\n",
    "    print(f'{unit.ljust(10)} {mcolors.to_hex(color).upper()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# see simulated colorblindness for selected unit colors\n",
    "print(\n",
    "    'https://davidmathlogic.com/colorblind/#' +\n",
    "    '-'.join([mcolors.to_hex(c).upper().replace('#', '%23') for c in unit_colors.values()]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "## Data Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_file = 'Gill-Chiel-eNeuro-2020-metadata.yml'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "units = [\n",
    "    'B38',\n",
    "    'I2',\n",
    "    'B8a/b',\n",
    "    'B6/B9',\n",
    "    'B3',\n",
    "    'B4/B5',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "channel_units = [\n",
    "    'uV', # I2\n",
    "    'uV', # RN\n",
    "    'uV', # BN2\n",
    "    'uV', # BN3\n",
    "    'mN', # Force\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "channel_names_by_animal = {\n",
    "    'JG07': ['I2-L', 'RN-L', 'BN2-L', 'BN3-L',    'Force'],\n",
    "    'JG08': ['I2',   'RN',   'BN2',   'BN3',      'Force'],\n",
    "    'JG11': ['I2',   'RN',   'BN2',   'BN3-PROX', 'Force'],\n",
    "    'JG12': ['I2',   'RN',   'BN2',   'BN3-DIST', 'Force'],\n",
    "    'JG14': ['I2',   'RN',   'BN2',   'BN3-PROX', 'Force'],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch_types_by_food = {\n",
    "    'Regular nori': ['Swallow (regular 5-cm nori strip)'],\n",
    "    'Tape nori':    ['Swallow (tape nori)'],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feeding_bouts = {\n",
    "    # (animal, food, bout_index): (data_set_name, time_window)\n",
    "\n",
    "    ('JG07', 'Regular nori', 0): ('Animal 1 Unloaded 1', [1169, 1191]), # 4 swallows\n",
    "    ('JG07', 'Regular nori', 1): ('Animal 1 Unloaded 2', [1495, 1518]), # 4 swallows\n",
    "    ('JG07', 'Regular nori', 2): ('Animal 1 Unloaded 3', [1581, 1615]), # 5 swallows\n",
    "    ('JG08', 'Regular nori', 0): ('Animal 2 Unloaded 1', [ 256,  287]), # 4 swallows\n",
    "    ('JG08', 'Regular nori', 1): ('Animal 2 Unloaded 2', [ 454,  481]), # 4 swallows\n",
    "    ('JG11', 'Regular nori', 0): ('Animal 3 Unloaded 1', [1791, 1819]), # 5 swallows\n",
    "    ('JG11', 'Regular nori', 1): ('Animal 3 Unloaded 2', [ 551,  568]), # 3 swallows\n",
    "    ('JG12', 'Regular nori', 0): ('Animal 4 Unloaded 1', [ 147,  165]), # 3 swallows\n",
    "    ('JG12', 'Regular nori', 1): ('Animal 4 Unloaded 2', [ 228,  245]), # 3 swallows\n",
    "    ('JG12', 'Regular nori', 2): ('Animal 4 Unloaded 3', [ 277,  291]), # 3 swallows\n",
    "    ('JG14', 'Regular nori', 0): ('Animal 5 Unloaded 1', [1834, 1865]), # 4 swallows\n",
    "    ('JG14', 'Regular nori', 1): ('Animal 5 Unloaded 2', [1909, 1943]), # 5 swallows\n",
    "    ('JG14', 'Regular nori', 2): ('Animal 5 Unloaded 3', [2049, 2084]), # 5 swallows\n",
    "    \n",
    "    ('JG07', 'Tape nori',    0): ('Animal 1 Loaded 1', [2716, 2755]), # 5 swallows\n",
    "    ('JG08', 'Tape nori',    0): ('Animal 2 Loaded 1', [ 141,  210]), # 7 swallows, some bucket and head movement\n",
    "    ('JG08', 'Tape nori',    1): ('Animal 2 Loaded 2', [ 660,  701]), # 5 swallows, large bucket movement\n",
    "    ('JG08', 'Tape nori',    2): ('Animal 2 Loaded 3', [1448, 1477]), # 3 swallows, some bucket movement\n",
    "    ('JG11', 'Tape nori',    0): ('Animal 3 Loaded 1', [1226, 1280]), # 5 swallows, inward food movements had very low amplitude\n",
    "    ('JG12', 'Tape nori',    0): ('Animal 4 Loaded 1', [ 434,  465]), # 4 swallows\n",
    "    ('JG12', 'Tape nori',    1): ('Animal 4 Loaded 2', [2897, 2937]), # 5 swallows\n",
    "    ('JG14', 'Tape nori',    0): ('Animal 5 Loaded 1', [ 828,  870]), # 5 swallows\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for example figures only -- not used in majority of analysis because of long inter-swallow intervals\n",
    "exemplary_swallow = ('JG12', 'Tape nori', 102)\n",
    "feeding_bouts[exemplary_swallow] = ('Animal 4 Loaded 2 Superset', [2977.3, 2984.5]) # 1 swallow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# these unloaded swallows have unreliable inward movement measurement\n",
    "# and are excluded from some parts of the analysis\n",
    "unreliable_inward_movement = [\n",
    "    ('JG07', 'Regular nori', 0, 1), # inward movement unusually short (maybe strip tore?)\n",
    "    ('JG07', 'Regular nori', 0, 3), # finished strip mid-retraction\n",
    "    ('JG07', 'Regular nori', 1, 2), # inward movement unusually short (maybe strip tore?)\n",
    "    ('JG07', 'Regular nori', 1, 3), # finished strip mid-retraction\n",
    "    ('JG07', 'Regular nori', 2, 2), # inward movement unusually short (maybe strip tore?)\n",
    "    ('JG07', 'Regular nori', 2, 4), # finished strip mid-retraction\n",
    "    ('JG08', 'Regular nori', 0, 3), # finished strip mid-retraction\n",
    "    ('JG08', 'Regular nori', 1, 3), # finished strip mid-retraction\n",
    "    ('JG11', 'Regular nori', 0, 4), # finished strip mid-retraction\n",
    "    ('JG11', 'Regular nori', 1, 2), # finished strip mid-retraction\n",
    "    ('JG12', 'Regular nori', 0, 2), # finished strip mid-retraction\n",
    "    # JG12 Regular nori 0 1 is OK\n",
    "    ('JG12', 'Regular nori', 2, 2), # finished strip mid-retraction\n",
    "    ('JG14', 'Regular nori', 0, 3), # finished strip mid-retraction\n",
    "    ('JG14', 'Regular nori', 1, 0), # bite-swallow and inward movement unusually short (maybe strip tore?)\n",
    "    ('JG14', 'Regular nori', 1, 3), # inward movement unusually short (maybe strip tore?)\n",
    "    ('JG14', 'Regular nori', 1, 4), # finished strip mid-retraction\n",
    "    ('JG14', 'Regular nori', 2, 1), # not visible\n",
    "]\n",
    "\n",
    "# these behaviors at the start of each unloaded bout are also excluded\n",
    "# because they are actually bite-swallows instead of pure swallows\n",
    "bite_swallow_behaviors = [\n",
    "    ('JG07', 'Regular nori', 0, 0),\n",
    "    ('JG07', 'Regular nori', 1, 0),\n",
    "    ('JG07', 'Regular nori', 2, 0),\n",
    "    ('JG08', 'Regular nori', 0, 0),\n",
    "    ('JG08', 'Regular nori', 1, 0),\n",
    "    ('JG11', 'Regular nori', 0, 0),\n",
    "    ('JG11', 'Regular nori', 1, 0),\n",
    "    ('JG12', 'Regular nori', 0, 0),\n",
    "    ('JG12', 'Regular nori', 1, 0),\n",
    "    ('JG12', 'Regular nori', 2, 0),\n",
    "    ('JG14', 'Regular nori', 0, 0),\n",
    "    ('JG14', 'Regular nori', 1, 0),\n",
    "    ('JG14', 'Regular nori', 2, 0),\n",
    "]\n",
    "unreliable_inward_movement += bite_swallow_behaviors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "## Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def query_union(queries):\n",
    "    return '|'.join([f'({q})' for q in queries if q is not None])\n",
    "\n",
    "def label2query(label):\n",
    "    animal, food = label[:4], label[5:]\n",
    "    query = f'(Animal == \"{animal}\") & (Food == \"{food}\")'\n",
    "    return query\n",
    "\n",
    "def contains(series, string):\n",
    "#     return np.array([string in x for x in series])\n",
    "    return series.map(lambda x: string in x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def get_sig(blk, channel):\n",
    "    sig = next((sig for sig in blk.segments[0].analogsignals if sig.name == channel), None)\n",
    "    if sig is None:\n",
    "        raise Exception(f'Channel \"{channel}\" could not be found')\n",
    "    else:\n",
    "        return sig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sig_index(blk, channel):\n",
    "    index = next((i for i, sig in enumerate(blk.segments[0].analogsignals) if sig.name == channel), None)\n",
    "    if index is None:\n",
    "        raise Exception(f'Channel \"{channel}\" could not be found')\n",
    "    else:\n",
    "        return index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_filters(blk, metadata):\n",
    "    # nearly identical to neurotic's implementation except\n",
    "    # time_slice ensures proxies are loaded\n",
    "    \n",
    "    for sig_filter in metadata['filters']:\n",
    "        index = get_sig_index(blk, sig_filter['channel'])\n",
    "        high = sig_filter.get('highpass', None)\n",
    "        low  = sig_filter.get('lowpass',  None)\n",
    "        if high:\n",
    "            high *= pq.Hz\n",
    "        if low:\n",
    "            low  *= pq.Hz\n",
    "        blk.segments[0].analogsignals[index] = elephant.signal_processing.butter(  # may raise a FutureWarning\n",
    "            signal = blk.segments[0].analogsignals[index].time_slice(None, None),\n",
    "            highpass_freq = high,\n",
    "            lowpass_freq  = low,\n",
    "        )\n",
    "    \n",
    "    return blk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def is_good_burst(burst):\n",
    "    time, duration, n_spikes = burst\n",
    "    return duration >= 0.5*pq.s and n_spikes > 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# must first convert args from quantities to simple ndarrays\n",
    "# (use .rescale('s').magnitude)\n",
    "def normalize_time(fixed_times, t, extrapolate=True):\n",
    "    \n",
    "    if not isinstance(t, np.ndarray):\n",
    "        if type(t) is list:\n",
    "            t = np.array(t)\n",
    "        else:\n",
    "            t = np.array([t])\n",
    "    \n",
    "    assert not isinstance(fixed_times, pq.Quantity), f'fixed_times should not be a quantity (use .magnitude)'\n",
    "    assert not isinstance(t, pq.Quantity), f't should not be a quantity (use .magnitude)'\n",
    "    assert np.all(np.diff(fixed_times[~np.isnan(fixed_times)])>=0), f'fixed_times must be sorted: {fixed_times}'\n",
    "    \n",
    "    # create a copy in which NaNs are removed\n",
    "    # - this is done because searchsorted does not work well with NaNs\n",
    "    # - infinities are retained here\n",
    "    fixed_times_without_nans = fixed_times[np.where(~np.isnan(fixed_times))[0]]\n",
    "    \n",
    "    # find the indexes of the fixed values that are just after each value in t\n",
    "    indexes = np.searchsorted(fixed_times_without_nans, t)\n",
    "    \n",
    "    # adjust indexes to account for the NaNs that were removed\n",
    "    nan_indexes = np.where(np.isnan(fixed_times))[0]\n",
    "    for nan_index in nan_indexes:\n",
    "        indexes[indexes >= nan_index] += 1\n",
    "    \n",
    "    # increment/decrement any index equal to 0/N, where N=len(fixed_times)\n",
    "    # - this is needed for values in t that are less/greater than the min/max fixed\n",
    "    #   time, and for NaNs in t which get assigned an index of N by searchsorted\n",
    "    # - for values in t less/greater than the min/max fixed time, this\n",
    "    #   increment/decrement in index will prepare that value to be normalized\n",
    "    #   using extrapolation based on the first/last interval in fixed_times\n",
    "    indexes = np.clip(indexes, 1, len(fixed_times)-1)\n",
    "    \n",
    "    # compute the normalized values of t using linear interpolation between the\n",
    "    # bordering fixed times\n",
    "    # - normalization of values in t that are less than the min fixed time is\n",
    "    #   accomplished by extrapolation, as the fraction (after-t)/(after-before)\n",
    "    #   will be greater than 1 since the t value is in fact earlier than the\n",
    "    #   \"before\" fixed time\n",
    "    # - normalization of values in t that are greater than the max fixed time is\n",
    "    #   accomplished by extrapolation, as the fraction (after-t)/(after-before)\n",
    "    #   will be less than 0 since the t value is in fact later than the \"after\"\n",
    "    #   fixed time\n",
    "    before = fixed_times[indexes-1]\n",
    "    after  = fixed_times[indexes]\n",
    "    t_normalized = indexes - (after-t)/(after-before)\n",
    "    \n",
    "    if not extrapolate:\n",
    "        # extrapolation is already done, so here we undo it by setting to NaN\n",
    "        # the normalized value for any t that is less/greater than the min/max\n",
    "        # fixed time\n",
    "        with np.errstate(invalid='ignore'): # don't warn about invalid comparisons to NaN\n",
    "            t_normalized[t < np.nanmin(fixed_times)] = np.nan\n",
    "            t_normalized[t > np.nanmax(fixed_times)] = np.nan\n",
    "    \n",
    "    return t_normalized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# must first convert args from quantities to simple ndarrays\n",
    "# (use .rescale('s').magnitude)\n",
    "def unnormalize_time(fixed_times, t_normalized, extrapolate=True):\n",
    "    \n",
    "    if not isinstance(t_normalized, np.ndarray):\n",
    "        if type(t_normalized) is list:\n",
    "            t_normalized = np.array(t_normalized)\n",
    "        else:\n",
    "            t_normalized = np.array([t_normalized])\n",
    "    \n",
    "    assert not isinstance(fixed_times, pq.Quantity), f'fixed_times should not be a quantity (use .magnitude)'\n",
    "    assert not isinstance(t_normalized, pq.Quantity), f't_normalized should not be a quantity (use .magnitude)'\n",
    "    assert np.all(np.diff(fixed_times[~np.isnan(fixed_times)])>=0), f'fixed_times must be sorted: {fixed_times}'\n",
    "    \n",
    "    # get the index of the fixed time that comes after each t\n",
    "    indexes = np.ceil(t_normalized)\n",
    "    \n",
    "    # clip the \"after\" indexes so that the first or last interval\n",
    "    # in fixed_times will be used for extrapolation\n",
    "    indexes = np.clip(indexes, 1, len(fixed_times)-1)\n",
    "    \n",
    "    # get the fixed times before and after each t\n",
    "    before = np.array([fixed_times[int(i)-1] if not np.isnan(i) else np.nan for i in indexes])\n",
    "    after  = np.array([fixed_times[int(i)]   if not np.isnan(i) else np.nan for i in indexes])\n",
    "    \n",
    "    # compute the real values of t\n",
    "    t = after + (t_normalized-indexes)*(after-before)\n",
    "    \n",
    "    if not extrapolate:\n",
    "        # extrapolation is already done, so here we undo it by setting to NaN\n",
    "        # the value for any t that is less/greater than the min/max fixed time\n",
    "        with np.errstate(invalid='ignore'): # don't warn about invalid comparisons to NaN\n",
    "            t[t < np.nanmin(fixed_times)] = np.nan\n",
    "            t[t > np.nanmax(fixed_times)] = np.nan\n",
    "    \n",
    "    return t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# these arrays are used as interp_times for resample_sig_in_normalized_time\n",
    "# depending on the segmentation scheme\n",
    "# - linspace ensures samples will be taken at regular intervals in normalized time\n",
    "interp_resolution = 1000\n",
    "video_seg_interp_times = np.linspace(0, 9, interp_resolution)\n",
    "force_seg_interp_times = np.linspace(0, 9, interp_resolution)\n",
    "\n",
    "def resample_sig_in_normalized_time(fixed_times, sig, interp_times):\n",
    "    \n",
    "    # get normalized times and signal values\n",
    "    # - normalize_time will put into times_normalized a np.nan wherever a time was not normalizable,\n",
    "    #   i.e., wherever the time occurred adjacent to a missing fixed time (np.nan in fixed_times)\n",
    "    times = sig.times.rescale('s').magnitude\n",
    "    times_normalized = normalize_time(fixed_times.magnitude, times)\n",
    "    y = sig.magnitude.flatten()\n",
    "\n",
    "    # drop times that could not be normalized due to missing fixed times\n",
    "    # - interp1d will erroneously interpolate across the gaps created by this deletion,\n",
    "    #   but we will then replace those interpolated values with np.nan\n",
    "    where_normalizable = np.where(~np.isnan(times_normalized))[0]\n",
    "    times_normalized = times_normalized[where_normalizable]\n",
    "    y = y[where_normalizable]\n",
    "\n",
    "    # resample evenly in normalized time\n",
    "    # - with bounds_error=False and fill_value=np.nan, interp1d will put np.nan in places outside\n",
    "    #   the min and max of times_normalized\n",
    "    # - interp1d will interpolate across the regions we deleted, but we want np.nan there,\n",
    "    #   so we will insert them manually in the next step\n",
    "    interp_func = sp.interpolate.interp1d(times_normalized, y, kind='linear', bounds_error=False, fill_value=np.nan)\n",
    "    y_interp = interp_func(interp_times)\n",
    "\n",
    "    # replace erroneously interpolated values with np.nan\n",
    "    # - now the points in y_interp which would correspond to times that could not be normalized\n",
    "    #   have been set to np.nan\n",
    "    where_not_normalizable = np.where(np.isnan(unnormalize_time(fixed_times.magnitude, interp_times)))[0]\n",
    "    y_interp[where_not_normalizable] = np.nan\n",
    "\n",
    "    return y_interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_column_analysis(column):\n",
    "    unit = column.split('(')[-1].split(')')[0]\n",
    "    mean = df_all[column].mean()\n",
    "    std = df_all[column].std()\n",
    "    cv = std/abs(mean)\n",
    "    print(f'Overall mean +/- std: {mean:.2f} +/- {std:.2f} {unit}')\n",
    "    print(f'Overall coefficient of variation CV: {cv:.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def differences_test(x, y, x_label='GROUP 1', y_label='GROUP 2', measure_label='MEASURE', units='UNITS', alpha=0.05):\n",
    "\n",
    "    # descriptive statistics\n",
    "    x_mean = np.mean(x)\n",
    "    y_mean = np.mean(y)\n",
    "    x_std = np.std(x, ddof=1)\n",
    "    y_std = np.std(y, ddof=1)\n",
    "    x_n = len(x)\n",
    "    y_n = len(y)\n",
    "    assert x_n == y_n, 'expected x and y to be paired but they have unequal length'\n",
    "    print(f'{x_label}: (M = {x_mean:g}, SD = {x_std:g}, N = {x_n:g})')\n",
    "    print(f'{y_label}: (M = {y_mean:g}, SD = {y_std:g}, N = {y_n:g})')\n",
    "    print()\n",
    "\n",
    "    # Shapiro-Wilk test for normality of differences\n",
    "    # - equivalent R test: shapiro.test(x-y)\n",
    "    shapiro_result = r_stats.shapiro_test(x.values-y.values)\n",
    "    shapiro_W, shapiro_p = shapiro_result['W'], shapiro_result['p']\n",
    "    shapiro_signif = '*' if shapiro_p < alpha else '(n.s.)'\n",
    "    print(f'H0: Differences have normal distribution, W = {shapiro_W:g},\\tp = {shapiro_p:g} {shapiro_signif}')\n",
    "\n",
    "    if shapiro_p >= 0.05:\n",
    "        print('- Because the differences can be assumed to be normal, a paired t-test will be used')\n",
    "\n",
    "        # paired one-tailed T-test for an increase in means\n",
    "        # - equivalent R test : t.test(x, y, paired=TRUE, alternative=\"greater\")\n",
    "        ttest_result = r_stats.t_test(x.values, y.values, paired=True, alternative='greater')\n",
    "        ttest_t, ttest_p = ttest_result['t'], ttest_result['p']\n",
    "        ttest_signif = '*' if ttest_p < alpha and x_mean > y_mean else '(n.s.)'\n",
    "        print(f'H0: Difference in means is not positive,  t = {ttest_t:g},\\tp = {ttest_p:g} {ttest_signif}')\n",
    "\n",
    "        # Cohen's d for effect size\n",
    "        # - equivalent R function: library(effsize); cohen.d(x, y)\n",
    "        # - although the R function offers a paired=TRUE case, I'm not using it here\n",
    "        #   because I don't understand the paper cited in the documentation. It appears\n",
    "        #   to be a variation on a different effect size measure, so it's unclear that\n",
    "        #   it should be referred to as \"Cohen's d\". When paired=TRUE, the result is\n",
    "        #   generally a little smaller.\n",
    "        cohen_d = r_stats.cohen_d(x.values, y.values)['estimate']\n",
    "        print()\n",
    "        print(f'Effect size: Cohen\\'s d = {cohen_d:g}')\n",
    "        \n",
    "        if ttest_p < alpha and x_mean > y_mean:\n",
    "            print()\n",
    "            print(f'\"A paired-samples one-tailed t-test indicated that {measure_label} for the ' \\\n",
    "                  f'[{x_n}] animals were significantly [greater/longer] for ' \\\n",
    "                  f'{x_label} (M = {x_mean:.2f} {units}, SD = {x_std:.2f} {units}) than for ' \\\n",
    "                  f'{y_label} (M = {y_mean:.2f} {units}, SD = {y_std:.2f} {units}), ' \\\n",
    "                  f't = {ttest_t:.3f}, df = {x_n-1}, p = {ttest_p:.3f}, Cohen\\'s d = {cohen_d:.2f}.\"')\n",
    "        else:\n",
    "            print()\n",
    "            print(f'\"A paired-samples one-tailed t-test indicated no significant increase in {measure_label} in the [{x_n}] animals between ' \\\n",
    "                  f'{x_label} (M = {x_mean:.2f} {units}, SD = {x_std:.2f} {units}) and ' \\\n",
    "                  f'{y_label} (M = {y_mean:.2f} {units}, SD = {y_std:.2f} {units}), ' \\\n",
    "                  f't = {ttest_t:.3f}, df = {x_n-1}, p = {ttest_p:.3f}, Cohen\\'s d = {cohen_d:.2f}.\"')\n",
    "        \n",
    "        return ttest_signif\n",
    "\n",
    "    else:\n",
    "        print('- Because the differences cannot be assumed to be normal, a Wilcoxon signed rank test will be used')\n",
    "\n",
    "        # paired Wilcoxon signed rank test for increase in locations (medians?)\n",
    "        # - equivalent R test: wilcox.test(x, y, paired=TRUE, alternative=\"greater\")\n",
    "        wilcoxon_result = r_stats.wilcox_test(x.values, y.values, paired=True, alternative='greater')\n",
    "        wilcoxon_W, wilcoxon_p = wilcoxon_result['W'], wilcoxon_result['p']\n",
    "        wilcoxon_signif = '*' if wilcoxon_p < alpha else '(n.s.)'\n",
    "        print(f'H0: Difference in medians is not positive, W = {wilcoxon_W:g},\\tp = {wilcoxon_p:g} {wilcoxon_signif}')\n",
    "        \n",
    "        return wilcoxon_signif"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "## Plotting Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def boxplot_with_points(x, y, hue, data, ax=None, show_points=False, describe=True):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    boxcolor = '0.75' if hue is None else None\n",
    "    pointcolor = 'k' if hue is None else None\n",
    "    edgecolor = '0.25'\n",
    "    linewidth = 0 if hue is None else 1\n",
    "    size = 4\n",
    "    \n",
    "    data = data.dropna(subset=[y])\n",
    "    \n",
    "    sns.boxplot(x=x, y=y, hue=hue, data=data, ax=ax, color=boxcolor, whis=999) # whiskers span extrema\n",
    "    \n",
    "    if show_points:\n",
    "        sns.swarmplot(x=x, y=y, hue=hue, data=data, ax=ax, color=pointcolor, linewidth=linewidth, edgecolor=edgecolor, size=size, dodge=True)\n",
    "        \n",
    "        if hue is not None:\n",
    "            # avoid duplicate legend entries\n",
    "            handles, labels = ax.get_legend_handles_labels()\n",
    "            n = int(len(labels)/2)\n",
    "            ax.legend(handles[:n], labels[:n], title=hue)\n",
    "    \n",
    "    ax.set_xlabel(None)\n",
    "\n",
    "    if describe:\n",
    "        by = [x] if hue is None else [x, hue]\n",
    "        print(y)\n",
    "#         print(data.groupby(by)[y].describe())\n",
    "        print(data.groupby(by)[y].apply(lambda y: {\n",
    "            'N': f'{y.count()}',\n",
    "            'Median': y.median(),\n",
    "#             'Q1': y.quantile(0.25),\n",
    "#             'Q3': y.quantile(0.75),\n",
    "        }).unstack())\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "default_markers = ['.', '+', 'x', '1', '4', 'o', 'D']\n",
    "default_colors = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6']\n",
    "\n",
    "def scatter2d(ax, data_subsets, xlabel, xlim, ylabel, ylim, trend=False, trend_separately=False, tooltips=False, padding=0.05, colors=default_colors, markers=default_markers):\n",
    "    \n",
    "    for j, (label, query) in enumerate(data_subsets.items()):\n",
    "        if query is not None:\n",
    "            df = df_all.query(query)\n",
    "            ax.scatter(df[xlabel], df[ylabel],\n",
    "                       label=label, marker=markers[j], c=colors[j], clip_on=False)\n",
    "            \n",
    "    all_points = df_all.query(query_union(data_subsets.values()))[[xlabel, ylabel]].dropna()\n",
    "    \n",
    "    if len(all_points) > 0:\n",
    "\n",
    "        xrange = np.ptp(all_points.iloc[:, 0])\n",
    "        xmin = min(all_points.iloc[:, 0]) - xrange * padding\n",
    "        xmax = max(all_points.iloc[:, 0]) + xrange * padding\n",
    "        if xlim is None:\n",
    "            xlim = [xmin, xmax]\n",
    "        if xlim[0] is None:\n",
    "            xlim[0] = xmin\n",
    "        if xlim[1] is None:\n",
    "            xlim[1] = xmax\n",
    "\n",
    "        yrange = np.ptp(all_points.iloc[:, 1])\n",
    "        ymin = min(all_points.iloc[:, 1]) - yrange * padding\n",
    "        ymax = max(all_points.iloc[:, 1]) + yrange * padding\n",
    "        if ylim is None:\n",
    "            ylim = [ymin, ymax]\n",
    "        if ylim[0] is None:\n",
    "            ylim[0] = ymin\n",
    "        if ylim[1] is None:\n",
    "            ylim[1] = ymax\n",
    "    \n",
    "    if trend_separately:\n",
    "        for j, (label, query) in enumerate(data_subsets.items()):\n",
    "            if query is not None:\n",
    "                df = df_all.query(query)[[xlabel, ylabel]].dropna()\n",
    "                model = sm.OLS(df.iloc[:,1], sm.add_constant(df.iloc[:,0])).fit()\n",
    "                model_stats = 'R$^2$ = {:.2f}, p = {:.5f}, n = {}'.format(model.rsquared, model.pvalues[1], len(df))\n",
    "                print(f'{label}:', model_stats)\n",
    "                model_x = np.linspace(df.iloc[:,0].min()-np.ptp(df.iloc[:,0])*0.1, df.iloc[:,0].max()+np.ptp(df.iloc[:,0])*0.1, 100)\n",
    "                model_y = model.params[0] + model.params[1] * model_x\n",
    "                ax.plot(model_x, model_y, color=colors[j])#, label=model_stats)\n",
    "                \n",
    "    if trend:\n",
    "        model = sm.OLS(all_points.iloc[:,1], sm.add_constant(all_points.iloc[:,0])).fit()\n",
    "        model_stats = 'R$^2$ = {:.2f}, p = {:.5f}, n = {}'.format(model.rsquared, model.pvalues[1], len(all_points))\n",
    "        print('All points:', model_stats)\n",
    "        model_x = np.linspace(all_points.iloc[:,0].min()-np.ptp(all_points.iloc[:,0])*0.1, all_points.iloc[:,0].max()+np.ptp(all_points.iloc[:,0])*0.1, 100)\n",
    "        model_y = model.params[0] + model.params[1] * model_x\n",
    "        ax.plot(model_x, model_y, color='k')#, label=model_stats)\n",
    "    \n",
    "    if tooltips:\n",
    "        # create a transparent layer containing all points for detecting mouse events\n",
    "        sc = ax.scatter(all_points.iloc[:, 0], all_points.iloc[:, 1], label=None,\n",
    "                        alpha=0, # transparent points\n",
    "                        s=0.001, # small radius of tooltip activation\n",
    "                       )\n",
    "        \n",
    "        # initialize a hidden empty tooltip\n",
    "        annot = ax.annotate(\n",
    "            '', xy=(0, 0),                              # initialize empty at origin\n",
    "            xytext=(0, 15), textcoords='offset points', # position text above point\n",
    "            color='k', size=10, ha='center',            # small black centered text\n",
    "            bbox=dict(fc='w', lw=0, alpha=0.6),         # transparent white background\n",
    "            arrowprops=dict(shrink=0, headlength=7, headwidth=7, width=0, lw=0, color='k'), # small black arrow\n",
    "        )\n",
    "        annot.set_visible(False)\n",
    "        \n",
    "        # prepare tooltip contents\n",
    "        tooltip_text = [' '.join([str(x) for x in index]) for index in list(all_points.index)]\n",
    "        \n",
    "        # bind the animation of tooltips to a hover event\n",
    "        def hover(event):\n",
    "            if event.inaxes == ax:\n",
    "                cont, ind = sc.contains(event)\n",
    "                if cont:\n",
    "                    point_index = ind['ind'][0]\n",
    "                    pos = sc.get_offsets()[point_index]\n",
    "                    annot.xy = pos\n",
    "                    annot.set_text(tooltip_text[point_index])\n",
    "                    annot.set_visible(True)\n",
    "                    ax.figure.canvas.draw_idle()\n",
    "                else:\n",
    "                    if annot.get_visible():\n",
    "                        annot.set_visible(False)\n",
    "                        ax.figure.canvas.draw_idle()\n",
    "        ax.figure.canvas.mpl_connect('motion_notify_event', hover)\n",
    "        \n",
    "    ax.set_xlabel(xlabel)\n",
    "    ax.set_xlim(xlim)\n",
    "    ax.set_ylabel(ylabel)#.replace('rectified ',''))\n",
    "    ax.set_ylim(ylim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_unloaded_vs_loaded(df, y, ax, color='0.25', show_all=False, show_statistics=True, show_signif=True, bracket_width=1.0, alpha=0.05):\n",
    "    \n",
    "    df = df.reset_index()\n",
    "    df.loc[df['Food'] == 'Regular nori', 'Food'] = 'Unloaded'\n",
    "    df.loc[df['Food'] == 'Tape nori',    'Food'] = 'Loaded'\n",
    "    \n",
    "#     # plot means\n",
    "#     sns.pointplot(\n",
    "#         x='Animal', y=y, hue='Food',\n",
    "#         data=df,\n",
    "#         hue_order=['Unloaded', 'Loaded'],\n",
    "#         dodge=0.4,\n",
    "# #         palette=[color], # do not desaturate by food\n",
    "#         palette=['0.75'], # do not desaturate by food\n",
    "#         join=False,\n",
    "#         estimator=np.mean,\n",
    "#         ci='sd',\n",
    "#         scale=0.6,\n",
    "#         capsize=0.2,\n",
    "#         errwidth=1,\n",
    "#         ax=ax,\n",
    "#     )\n",
    "    \n",
    "#     # plot individual swallows points\n",
    "#     if show_all:\n",
    "#         sns.stripplot(\n",
    "#             x='Animal', y=y, hue='Food',\n",
    "#             data=df,\n",
    "#             hue_order=['Unloaded', 'Loaded'],\n",
    "#             jitter=False,\n",
    "#             dodge=True,\n",
    "#             linewidth=1,\n",
    "#             marker='_',\n",
    "#             color='k',\n",
    "#             ax=ax,\n",
    "#             clip_on=False,\n",
    "#         )\n",
    "\n",
    "    for food, offset in {'Unloaded': -0.2, 'Loaded': 0.2}.items():\n",
    "        for i, (animal, values) in enumerate(df[df['Food'] == food].groupby('Animal')[y]):\n",
    "            \n",
    "            # plot the mean\n",
    "            ax.scatter(\n",
    "                [i+offset], values.mean(),\n",
    "                color=color,\n",
    "                s=15,\n",
    "                zorder=2, clip_on=False,\n",
    "            )\n",
    "            \n",
    "            # plot the standard error of the mean\n",
    "            lower_sem = values.mean()-values.sem()\n",
    "            upper_sem = values.mean()+values.sem()\n",
    "            ax.plot(\n",
    "                [i+offset]*2, [lower_sem, upper_sem],\n",
    "                color=color, lw=1,\n",
    "                zorder=2, clip_on=False,\n",
    "            )\n",
    "            ax.plot(\n",
    "                [i+offset-0.1, i+offset+0.1], [lower_sem, lower_sem],\n",
    "                color=color, lw=1,\n",
    "                zorder=2, clip_on=False,\n",
    "            )\n",
    "            ax.plot(\n",
    "                [i+offset-0.1, i+offset+0.1], [upper_sem, upper_sem],\n",
    "                color=color, lw=1,\n",
    "                zorder=2, clip_on=False,\n",
    "            )\n",
    "            \n",
    "            if show_all:\n",
    "                # plot individual swallows points\n",
    "                ax.scatter(\n",
    "                    [i+offset]*values.size, values,\n",
    "\n",
    "                    facecolors='none',\n",
    "                    edgecolors='k',\n",
    "                    linewidths=1,\n",
    "                    s=12,\n",
    "\n",
    "#                     facecolors='k',\n",
    "#                     marker='_',\n",
    "#                     s=18,\n",
    "\n",
    "                    zorder=3,\n",
    "                )\n",
    "            ax.axvline(x=i+offset, c='0.9', lw=0.5, zorder=-2)\n",
    "            ax.annotate(\n",
    "                food[0],\n",
    "                xy=(i+offset, -0.02), xycoords=ax.get_xaxis_transform(),\n",
    "                ha='center', va='top', fontsize='x-small', color='0.5',\n",
    "            )\n",
    "            if food == 'Unloaded': # do it just once\n",
    "                ax.annotate(\n",
    "                    i+1,\n",
    "                    xy=(i, -0.02), xycoords=ax.get_xaxis_transform(),\n",
    "                    xytext=(0, -10), textcoords='offset points',\n",
    "                    ha='center', va='top', fontsize='x-small', color='0.5',\n",
    "                )\n",
    "\n",
    "    # plot lines connecting points\n",
    "    for i, (y_unloaded, y_loaded) in enumerate(zip(\n",
    "                df.query('Food == \"Unloaded\"').groupby('Animal')[y].mean(),\n",
    "                df.query('Food == \"Loaded\"').groupby('Animal')[y].mean())):\n",
    "        ax.plot(\n",
    "            [i-0.2, i+0.2],\n",
    "            [y_unloaded, y_loaded],\n",
    "            color='0.75',\n",
    "            zorder=-1,\n",
    "        )\n",
    "\n",
    "    ax.tick_params(bottom=False, labelbottom=False)\n",
    "#     ax.legend_.remove()\n",
    "    ax.set_xlabel(None)\n",
    "    sns.despine(ax=ax, bottom=True)\n",
    "\n",
    "    if show_statistics:\n",
    "        print(df.groupby(['Animal', 'Food'])[y].apply(lambda x: {'Mean': x.mean(), 'SEM': x.sem(), 'STD': x.std(), 'Count': x.count()}).unstack([1, 2])[['Unloaded', 'Loaded']])\n",
    "        print()\n",
    "        signif = differences_test(\n",
    "            x=df.query('Food == \"Loaded\"').groupby('Animal')[y].mean(),\n",
    "            y=df.query('Food == \"Unloaded\"').groupby('Animal')[y].mean(),\n",
    "            x_label='loaded swallows',\n",
    "            y_label='unloaded swallows',\n",
    "            measure_label='mean [' + ' '.join(y.split()[:-1]) + ']',\n",
    "            units=y.split()[-1].strip('()'),\n",
    "            alpha=alpha,\n",
    "        )\n",
    "        \n",
    "        if show_signif and signif == '*':\n",
    "            from matplotlib.patches import ArrowStyle\n",
    "            ax.annotate(\n",
    "                '*',\n",
    "                xy=(0.5, 1.03), xycoords='axes fraction',\n",
    "                xytext=(0, 0), textcoords='offset points',\n",
    "                arrowprops=dict(arrowstyle=ArrowStyle.BracketB(widthB=bracket_width, lengthB=0.2, angleB=None), color='0.5'),\n",
    "                ha='center', fontsize='x-large', color='0.5',\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Crunch the Numbers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# download each file that is not already stored locally\n",
    "metadata = neurotic.MetadataSelector(file=metadata_file)\n",
    "for (animal, food, bout_index), (data_set_name, time_window) in feeding_bouts.items():\n",
    "    metadata.select(data_set_name)\n",
    "    metadata.download_all_data_files()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import and Process the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start = datetime.datetime.now()\n",
    "\n",
    "# use Neo RawIO lazy loading to load much faster and using less memory\n",
    "# - with lazy=True, filtering parameters specified in metadata are ignored\n",
    "#     - note: filters are replaced below and applied manually anyway\n",
    "# - with lazy=True, loading via time_slice requires neo>=0.8.0\n",
    "lazy = True\n",
    "\n",
    "# load the metadata containing file paths\n",
    "metadata = neurotic.MetadataSelector(file=metadata_file)\n",
    "\n",
    "# filter epochs for each bout and perform calculations\n",
    "df_list = []\n",
    "last_data_set_name = None\n",
    "pbar = tqdm(total=len(feeding_bouts), unit='feeding bout')\n",
    "for (animal, food, bout_index), (data_set_name, time_window) in feeding_bouts.items():\n",
    "\n",
    "    channel_names = channel_names_by_animal[animal]\n",
    "    epoch_types = epoch_types_by_food[food]\n",
    "\n",
    "    ###\n",
    "    ### LOAD DATASET\n",
    "    ###\n",
    "\n",
    "    metadata.select(data_set_name)\n",
    "\n",
    "    if data_set_name is last_data_set_name:\n",
    "        # skip reloading the data if it's already in memory\n",
    "        pass\n",
    "    else:\n",
    "        blk = neurotic.load_dataset(metadata, lazy=lazy)\n",
    "        if lazy:\n",
    "            # manually perform filters\n",
    "            blk = apply_filters(blk, metadata)\n",
    "\n",
    "    last_data_set_name = data_set_name\n",
    "\n",
    "\n",
    "\n",
    "    ###\n",
    "    ### LOCATE BEHAVIOR EPOCHS AND SUBEPOCHS\n",
    "    ###\n",
    "\n",
    "    # construct a query for locating behaviors\n",
    "    behavior_query = f'(Type in {epoch_types}) & ({time_window[0]} <= Start) & (End <= {time_window[1]})'\n",
    "\n",
    "    # construct queries for locating epochs associated with each behavior\n",
    "    # - each query should match at most one epoch\n",
    "    # - dictionary keys are used as prefixes for the names of new columns\n",
    "    subepoch_queries = {}\n",
    "\n",
    "    subepoch_queries['B38 activity']            = (f'(Type == \"B38 activity\") & ' \\\n",
    "                                                   f'(@behavior_start-3 <= End) & (End <= @behavior_start+4)',\n",
    "                                                   'last') # use last if there are multiple matches\n",
    "                                                  # must end within a few seconds of behavior start (3 before or 4 after)\n",
    "\n",
    "    subepoch_queries['I2 protraction activity'] = f'(Type == \"I2 protraction activity\") & ' \\\n",
    "                                                  f'(@behavior_start-1 <= Start) & (End <= @behavior_end)'\n",
    "                                                  # must start no earlier than 1 second before behavior and end within it\n",
    "\n",
    "    subepoch_queries['B8 activity']             = f'(Type == \"B8 activity\") & ' \\\n",
    "                                                  f'(@behavior_start <= Start) & (End <= @behavior_end)'\n",
    "                                                  # must be fully contained within behavior\n",
    "\n",
    "    subepoch_queries['B3/6/9/10 activity']      = f'(Type == \"B3/6/9/10 activity\") & ' \\\n",
    "                                                  f'(@behavior_start <= Start) & (End <= @behavior_end)'\n",
    "                                                  # must be fully contained within behavior\n",
    "\n",
    "    subepoch_queries['B4/B5 activity']          = (f'(Type == \"B4/B5 activity\") & ' \\\n",
    "                                                   f'(@behavior_start <= Start) & (Start <= @behavior_end)',\n",
    "                                                   'first') # use first if there are multiple matches\n",
    "                                                  # must start within behavior\n",
    "\n",
    "    subepoch_queries['Force shoulder end']      = f'(Type == \"Force shoulder end\") & ' \\\n",
    "                                                  f'(@behavior_start-3 <= End) & (End <= @behavior_start+3)'\n",
    "                                                  # must end within 3 seconds of behavior start\n",
    "\n",
    "    subepoch_queries['Force rise start']        = f'(Type == \"Force rise start\") & ' \\\n",
    "                                                  f'(@behavior_start <= Start) & (Start <= @behavior_end)'\n",
    "                                                  # must start within behavior\n",
    "\n",
    "    subepoch_queries['Force plateau start']     = f'(Type == \"Force plateau start\") & ' \\\n",
    "                                                  f'(@behavior_start <= Start) & (Start <= @behavior_end)'\n",
    "                                                  # must start within behavior\n",
    "\n",
    "    subepoch_queries['Force plateau end']       = f'(Type == \"Force plateau end\") & ' \\\n",
    "                                                  f'(@behavior_end-2 <= Start) & (Start <= @behavior_end+2)'\n",
    "                                                  # must start within 2 seconds of behavior end\n",
    "\n",
    "    subepoch_queries['Force drop end']          = f'(Type == \"Force drop end\") & ' \\\n",
    "                                                  f'(@behavior_end-2 <= Start) & (Start <= @behavior_end+2)'\n",
    "                                                  # must start within 2 seconds of behavior end\n",
    "\n",
    "    subepoch_queries['Inward movement']         = f'(Type == \"Inward movement\") & ' \\\n",
    "                                                  f'(@behavior_start <= Start) & (Start <= @behavior_end)'\n",
    "                                                  # must start within behavior\n",
    "\n",
    "    # construct a table in which each row is a behavior and subepoch data\n",
    "    # is added as columns, e.g. df['B38 activity start (s)']\n",
    "    df = BehaviorsDataFrame(blk.segments[0].epochs, behavior_query, subepoch_queries)\n",
    "\n",
    "    # renumber behaviors assuming all behaviors are from a single contiguous sequence\n",
    "    df = df.sort_values('Start (s)').reset_index(drop=True).rename_axis('Behavior_index')\n",
    "\n",
    "\n",
    "\n",
    "    ###\n",
    "    ### START CALCULATIONS\n",
    "    ###\n",
    "\n",
    "    # some columns must have type 'object', which\n",
    "    # can be accomplished by initializing with None or np.nan\n",
    "    df['Video segmentation times (s)'] = None\n",
    "    df['Force segmentation times (s)'] = None\n",
    "    df['Force, video segmented interpolation (mN)'] = None\n",
    "    df['Force, force segmented interpolation (mN)'] = None\n",
    "    for unit in units:\n",
    "        df[f'{unit} spike train'] = None\n",
    "        df[f'{unit} firing rate (Hz)'] = None\n",
    "        df[f'{unit} firing rate, video segmented interpolation (Hz)'] = None\n",
    "        df[f'{unit} firing rate, force segmented interpolation (Hz)'] = None\n",
    "        df[f'{unit} all bursts'] = None\n",
    "\n",
    "        # while we're at it, initialize some other things that might otherwise never be given values\n",
    "        df[f'{unit} burst start (s)'] = np.nan\n",
    "        df[f'{unit} burst end (s)'] = np.nan\n",
    "        df[f'{unit} burst duration (s)'] = 0\n",
    "        df[f'{unit} burst spike count'] = 0\n",
    "        df[f'{unit} burst mean frequency (Hz)'] = 0\n",
    "        df[f'{unit} burst start (video seg normalized)'] = np.nan\n",
    "        df[f'{unit} burst end (video seg normalized)'] = np.nan\n",
    "        df[f'{unit} burst start (force seg normalized)'] = np.nan\n",
    "        df[f'{unit} burst end (force seg normalized)'] = np.nan\n",
    "    df['Inward movement start (video seg normalized)'] = np.nan\n",
    "    df['Inward movement end (video seg normalized)'] = np.nan\n",
    "    df['Inward movement start (force seg normalized)'] = np.nan\n",
    "    df['Inward movement end (force seg normalized)'] = np.nan\n",
    "\n",
    "\n",
    "\n",
    "    # get smoothed force for entire time window\n",
    "    channel = 'Force'\n",
    "    sig = get_sig(blk, channel)\n",
    "    if lazy:\n",
    "        sig = sig.time_slice(None, None)\n",
    "    sig = elephant.signal_processing.butter(sig, lowpass_freq = 5*pq.Hz)\n",
    "    sig = sig.time_slice(time_window[0]*pq.s, time_window[1]*pq.s)\n",
    "    sig = sig.rescale('mN')\n",
    "    force_smoothed_sig = sig\n",
    "\n",
    "\n",
    "\n",
    "    df['End to next start (s)'] = np.nan\n",
    "    df['Start to next start (s)'] = np.nan\n",
    "    previous_i = None\n",
    "\n",
    "    # iterate over all swallows\n",
    "    for j, i in enumerate(df.index):\n",
    "\n",
    "        behavior_start = df.loc[i, 'Start (s)']*pq.s\n",
    "        behavior_end = df.loc[i, 'End (s)']*pq.s\n",
    "        inward_movement_start = df.loc[i, 'Inward movement start (s)']*pq.s\n",
    "        inward_movement_end = df.loc[i, 'Inward movement end (s)']*pq.s\n",
    "\n",
    "        # calculate interbehavior intervals assuming all behaviors are from a single contiguous sequence\n",
    "        if previous_i is not None:\n",
    "            df.loc[previous_i, 'End to next start (s)']   = df.loc[i, 'Start (s)'] - df.loc[previous_i, 'End (s)']\n",
    "            df.loc[previous_i, 'Start to next start (s)'] = df.loc[i, 'Start (s)'] - df.loc[previous_i, 'Start (s)']\n",
    "            df.loc[previous_i, 'Inward movement start to next inward movement start (s)'] = \\\n",
    "                df.loc[i, 'Inward movement start (s)'] - df.loc[previous_i, 'Inward movement start (s)']\n",
    "            df.loc[previous_i, 'Inward movement end to next inward movement start (s)'] = \\\n",
    "                df.loc[i, 'Inward movement start (s)'] - df.loc[previous_i, 'Inward movement end (s)']\n",
    "        previous_i = i\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### VIDEO SEGMENTATION\n",
    "        ###\n",
    "\n",
    "        # inward movement start and end are required\n",
    "        video_is_segmented = np.all(np.isfinite(np.array([\n",
    "            inward_movement_start, inward_movement_end])))\n",
    "\n",
    "        if video_is_segmented:\n",
    "            inward_movement_duration = df.loc[i, 'Inward movement duration (s)']*pq.s\n",
    "\n",
    "            # get the list of fixed times for normalization\n",
    "            video_segmentation_times = df.at[i, 'Video segmentation times (s)'] = np.array([\n",
    "                inward_movement_start-inward_movement_duration*4,\n",
    "                inward_movement_start-inward_movement_duration*3,\n",
    "                inward_movement_start-inward_movement_duration*2,\n",
    "                inward_movement_start-inward_movement_duration*1,\n",
    "                inward_movement_start,\n",
    "                inward_movement_end,\n",
    "                inward_movement_end+inward_movement_duration*1,\n",
    "                inward_movement_end+inward_movement_duration*2,\n",
    "                inward_movement_end+inward_movement_duration*3,\n",
    "                inward_movement_end+inward_movement_duration*4,\n",
    "            ])*pq.s # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "        else: # video is not segmented\n",
    "            video_segmentation_times = df.at[i, 'Video segmentation times (s)'] = np.array([\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "            ])*pq.s # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### FORCE SEGMENTATION\n",
    "        ###\n",
    "\n",
    "        force_shoulder_end  = df.loc[i, 'Force shoulder end start (s)']*pq.s  # start of \"Force shoulder end\" epoch\n",
    "        force_rise_start    = df.loc[i, 'Force rise start start (s)']*pq.s    # start of \"Force rise start\" epoch\n",
    "        force_plateau_start = df.loc[i, 'Force plateau start start (s)']*pq.s # start of \"Force plateau start\" epoch\n",
    "        force_plateau_end   = df.loc[i, 'Force plateau end start (s)']*pq.s   # start of \"Force plateau end\" epoch\n",
    "        force_drop_end      = df.loc[i, 'Force drop end start (s)']*pq.s      # start of \"Force drop end\" epoch\n",
    "\n",
    "        # force rise start, plateau start and end, and drop end are required\n",
    "        force_is_segmented = np.all(np.isfinite(np.array([\n",
    "            force_rise_start, force_plateau_start, force_plateau_end, force_drop_end])))\n",
    "\n",
    "        if force_is_segmented:\n",
    "            # get some times for the previous and next swallow\n",
    "            epochs_force_shoulder_end  = next((ep for ep in blk.segments[0].epochs if ep.name == 'Force shoulder end'), None)\n",
    "            epochs_force_rise_start    = next((ep for ep in blk.segments[0].epochs if ep.name == 'Force rise start'), None)\n",
    "            epochs_force_plateau_start = next((ep for ep in blk.segments[0].epochs if ep.name == 'Force plateau start'), None)\n",
    "            epochs_force_plateau_end   = next((ep for ep in blk.segments[0].epochs if ep.name == 'Force plateau end'), None)\n",
    "            epochs_force_drop_end      = next((ep for ep in blk.segments[0].epochs if ep.name == 'Force drop end'), None)\n",
    "            assert epochs_force_shoulder_end  is not None, 'failed to find \"Force shoulder end\" epochs'\n",
    "            assert epochs_force_rise_start    is not None, 'failed to find \"Force rise start\" epochs'\n",
    "            assert epochs_force_plateau_start is not None, 'failed to find \"Force plateau start\" epochs'\n",
    "            assert epochs_force_plateau_end   is not None, 'failed to find \"Force plateau end\" epochs'\n",
    "            assert epochs_force_drop_end      is not None, 'failed to find \"Force drop end\" epochs'\n",
    "\n",
    "            try:\n",
    "                prev_force_plateau_start = df.loc[i, 'Previous force plateau start (s)'] = epochs_force_plateau_start.time_slice(None, force_rise_start)[-1]\n",
    "                assert force_rise_start-prev_force_plateau_start < 16*pq.s, f'for swallow {i}, previous force plateau start is too far away'\n",
    "            except IndexError:\n",
    "                prev_force_plateau_start = df.loc[i, 'Previous force plateau start (s)'] = np.nan\n",
    "\n",
    "            try:\n",
    "                prev_force_plateau_end = df.loc[i, 'Previous force plateau end (s)'] = epochs_force_plateau_end.time_slice(None, force_rise_start)[-1]\n",
    "                assert force_rise_start-prev_force_plateau_end < 12*pq.s, f'for swallow {i}, previous force plateau end is too far away'\n",
    "            except IndexError:\n",
    "                prev_force_plateau_end = df.loc[i, 'Previous force plateau end (s)'] = np.nan\n",
    "\n",
    "            try:\n",
    "                prev_force_drop_end = df.loc[i, 'Previous force drop end (s)'] = epochs_force_drop_end.time_slice(None, force_rise_start)[-1]\n",
    "                assert force_rise_start-prev_force_drop_end < 12*pq.s, f'for swallow {i}, previous force drop end is too far away'\n",
    "            except IndexError:\n",
    "                prev_force_drop_end = df.loc[i, 'Previous force drop end (s)'] = np.nan\n",
    "\n",
    "            try:\n",
    "                next_force_rise_start = df.loc[i, 'Next force rise start (s)'] = epochs_force_rise_start.time_slice(force_drop_end, None)[0]\n",
    "                assert next_force_rise_start-force_drop_end < 12*pq.s, f'for swallow {i}, next force rise start is too far away'\n",
    "            except IndexError:\n",
    "                next_force_rise_start = df.loc[i, 'Next force rise start (s)'] = np.nan\n",
    "\n",
    "            try:\n",
    "                next_force_shoulder_end = df.loc[i, 'Next force shoulder end (s)'] = epochs_force_shoulder_end.time_slice(force_drop_end, None)[0]\n",
    "                if next_force_shoulder_end > next_force_rise_start:\n",
    "                    # next swallow did not have a shoulder and we instead grabbed a later shoulder\n",
    "                    next_force_shoulder_end = df.loc[i, 'Next force shoulder end (s)'] = np.nan\n",
    "            except IndexError:\n",
    "                next_force_shoulder_end = df.loc[i, 'Next force shoulder end (s)'] = np.nan\n",
    "\n",
    "            # get the list of fixed times for normalization\n",
    "            force_segmentation_times = df.at[i, 'Force segmentation times (s)'] = np.array([\n",
    "                prev_force_plateau_start,\n",
    "                prev_force_plateau_end,\n",
    "                prev_force_drop_end,\n",
    "                force_shoulder_end,\n",
    "                force_rise_start,\n",
    "                force_plateau_start,\n",
    "                force_plateau_end,\n",
    "                force_drop_end,\n",
    "                next_force_shoulder_end,\n",
    "                next_force_rise_start,\n",
    "            ])*pq.s # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "        else: # force is not segmented\n",
    "            force_segmentation_times = df.at[i, 'Force segmentation times (s)'] = np.array([\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "                np.nan,\n",
    "            ])*pq.s # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### BEHAVIORAL MARKERS\n",
    "        ###\n",
    "\n",
    "        if video_is_segmented:\n",
    "            df.loc[i, 'Inward movement start (video seg normalized)'] = \\\n",
    "                normalize_time(video_segmentation_times.magnitude, float(inward_movement_start))\n",
    "            df.loc[i, 'Inward movement end (video seg normalized)'] = \\\n",
    "                normalize_time(video_segmentation_times.magnitude, float(inward_movement_end))\n",
    "\n",
    "        if force_is_segmented:\n",
    "            df.loc[i, 'Inward movement start (force seg normalized)'] = \\\n",
    "                normalize_time(force_segmentation_times.magnitude, float(inward_movement_start))\n",
    "            df.loc[i, 'Inward movement end (force seg normalized)'] = \\\n",
    "                normalize_time(force_segmentation_times.magnitude, float(inward_movement_end))\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### FORCE QUANTIFICATION\n",
    "        ###\n",
    "\n",
    "        if force_is_segmented:\n",
    "\n",
    "            # get smoothed force for whole behavior for remaining force calculations\n",
    "            sig = force_smoothed_sig\n",
    "            if np.isfinite(force_shoulder_end):\n",
    "                sig = sig.time_slice(force_shoulder_end - 0.01*pq.s, force_drop_end + 0.01*pq.s)\n",
    "            else:\n",
    "                sig = sig.time_slice(force_rise_start - 1*pq.s, force_drop_end + 0.01*pq.s)\n",
    "            sig = sig.rescale('mN')\n",
    "\n",
    "            # find force peak, baseline, and the increase\n",
    "            force_min_time = df.loc[i, 'Force minimum time (s)'] = elephant.spike_train_generation.peak_detection(sig, 999*pq.mN, sign='below')[0]\n",
    "            force_min = df.loc[i, 'Force minimum (mN)'] = sig[sig.time_index(force_min_time)][0]\n",
    "            force_peak_time = df.loc[i, 'Force peak time (s)'] = elephant.spike_train_generation.peak_detection(sig, 0*pq.mN)[0]\n",
    "            force_peak = df.loc[i, 'Force peak (mN)'] = sig[sig.time_index(force_peak_time)][0]\n",
    "            force_baseline = df.loc[i, 'Force baseline (mN)'] = sig[sig.time_index(force_rise_start)][0]\n",
    "            force_increase = df.loc[i, 'Force increase (mN)'] = force_peak-force_baseline\n",
    "\n",
    "            # find force plateau, drop, and shoulder values\n",
    "            force_plateau_start_value = df.loc[i, 'Force plateau start value (mN)'] = sig[sig.time_index(force_plateau_start)][0]\n",
    "            force_plateau_end_value = df.loc[i, 'Force plateau end value (mN)'] = sig[sig.time_index(force_plateau_end)][0]\n",
    "            force_drop_end_value = df.loc[i, 'Force drop end value (mN)'] = sig[sig.time_index(force_drop_end)][0]\n",
    "            if np.isfinite(force_shoulder_end):\n",
    "                force_shoulder_end_value = df.loc[i, 'Force shoulder end value (mN)'] = sig[sig.time_index(force_shoulder_end)][0]\n",
    "            else:\n",
    "                force_shoulder_end_value = np.nan\n",
    "\n",
    "            # find force rise and plateau durations\n",
    "            force_rise_duration = df.loc[i, 'Force rise duration (s)'] = force_plateau_start - force_rise_start\n",
    "            force_plateau_duration = df.loc[i, 'Force plateau duration (s)'] = force_plateau_end - force_plateau_start\n",
    "            force_rise_plateau_duration = df.loc[i, 'Force rise and plateau duration (s)'] = force_plateau_end - force_rise_start\n",
    "\n",
    "            # find average slope during rising phase\n",
    "            force_rise_increase = df.loc[i, 'Force rise increase (mN)'] = force_plateau_start_value - force_baseline\n",
    "            force_slope = df.loc[i, 'Force slope (mN/s)'] = (force_rise_increase/force_rise_duration).rescale('mN/s')\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### FORCE NORMALIZATION\n",
    "        ###\n",
    "\n",
    "        if force_is_segmented:\n",
    "            t_start = force_segmentation_times[0]-0.001*pq.s\n",
    "            t_stop = force_segmentation_times[-1]+0.001*pq.s\n",
    "\n",
    "            channel = 'Force'\n",
    "            sig = get_sig(blk, channel)\n",
    "            sig = sig.time_slice(t_start, t_stop)\n",
    "            sig = sig.rescale(channel_units[channel_names.index(channel)])\n",
    "\n",
    "            force_force_seg_interp = df.at[i, 'Force, force segmented interpolation (mN)'] = \\\n",
    "                resample_sig_in_normalized_time(force_segmentation_times, sig, force_seg_interp_times) # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### FIND SPIKE TRAINS\n",
    "        ###\n",
    "\n",
    "        if lazy:\n",
    "            if metadata['amplitude_discriminators'] is not None:\n",
    "                for discriminator in metadata['amplitude_discriminators']:\n",
    "                    sig = get_sig(blk, discriminator['channel'])\n",
    "                    if sig is not None:\n",
    "                        sig = sig.time_slice(max(sig.t_start, behavior_start - 5*pq.s), min(sig.t_stop, behavior_end + 5*pq.s))\n",
    "                        st = _detect_spikes(sig, discriminator, blk.segments[0].epochs)\n",
    "                        st_epoch_start = df.loc[i, discriminator['epoch']+' start (s)']*pq.s\n",
    "                        st_epoch_end = df.loc[i, discriminator['epoch']+' end (s)']*pq.s\n",
    "                        st = st.time_slice(st_epoch_start, st_epoch_end)\n",
    "                        df.at[i, f'{st.name} spike train'] = st # 'at', not 'loc', is important for inserting list into cell\n",
    "        else:\n",
    "            for spiketrain in blk.segments[0].spiketrains:\n",
    "                discriminator = next((d for d in metadata['amplitude_discriminators'] if d['name'] == spiketrain.name), None)\n",
    "                if discriminator is None:\n",
    "                    raise Exception(f'For data set \"{data_set_name}\", discriminator \"{spiketrain.name}\" could not be found')\n",
    "                st_epoch_start = df.loc[i, discriminator['epoch']+' start (s)']*pq.s\n",
    "                st_epoch_end = df.loc[i, discriminator['epoch']+' end (s)']*pq.s\n",
    "                if np.isfinite(st_epoch_start) and np.isfinite(st_epoch_end):\n",
    "                    st = spiketrain.time_slice(st_epoch_start, st_epoch_end)\n",
    "                else:\n",
    "                    # this unit's discriminator epoch was not located for this swallow\n",
    "                    st = None\n",
    "                df.at[i, f'{spiketrain.name} spike train'] = st # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### QUANTIFY SPIKE TRAINS AND BURSTS\n",
    "        ###\n",
    "\n",
    "        for k, unit in enumerate(units):\n",
    "            st = df.loc[i, f'{unit} spike train']\n",
    "            if st is not None:\n",
    "\n",
    "                # get the neural channel\n",
    "                channel = st.annotations['channels'][0]\n",
    "                sig = get_sig(blk, channel)\n",
    "\n",
    "                # create a continuous smoothed firing rate representation\n",
    "                # by convolving the spike train with a kernel\n",
    "                # - choice of t_start and t_stop here ensures firing rates are\n",
    "                #   recorded as zero far from the burst and can be resampled later\n",
    "                if video_is_segmented and force_is_segmented:\n",
    "                    t_start = min(video_segmentation_times[0], force_segmentation_times[0])-0.001*pq.s\n",
    "                    t_stop = max(video_segmentation_times[-1], force_segmentation_times[-1])+0.001*pq.s\n",
    "                elif video_is_segmented:\n",
    "                    t_start = video_segmentation_times[0]-0.001*pq.s\n",
    "                    t_stop = video_segmentation_times[-1]+0.001*pq.s\n",
    "                elif force_is_segmented:\n",
    "                    t_start = force_segmentation_times[0]-0.001*pq.s\n",
    "                    t_stop = force_segmentation_times[-1]+0.001*pq.s\n",
    "                else:\n",
    "                    # no segmentation available\n",
    "                    t_start = behavior_start-5*pq.s\n",
    "                    t_stop = behavior_end+5*pq.s\n",
    "                smoothing_kernel = elephant.kernels.GaussianKernel(0.2*pq.s) # 200 ms standard deviation\n",
    "                firing_rate = df.at[i, f'{unit} firing rate (Hz)'] = elephant.statistics.instantaneous_rate(\n",
    "                    spiketrain=st,\n",
    "                    t_start=t_start,\n",
    "                    t_stop=t_stop,\n",
    "                    sampling_period=sig.sampling_period,\n",
    "                    kernel=smoothing_kernel,\n",
    "                ) # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "                # normalization\n",
    "                if video_is_segmented:\n",
    "                    firing_rate_video_seg_interp = df.at[i, f'{unit} firing rate, video segmented interpolation (Hz)'] = \\\n",
    "                        resample_sig_in_normalized_time(video_segmentation_times, firing_rate, video_seg_interp_times) # 'at', not 'loc', is important for inserting list into cell\n",
    "                if force_is_segmented:\n",
    "                    firing_rate_force_seg_interp = df.at[i, f'{unit} firing rate, force segmented interpolation (Hz)'] = \\\n",
    "                        resample_sig_in_normalized_time(force_segmentation_times, firing_rate, force_seg_interp_times) # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "                if st.size > 0:\n",
    "\n",
    "                    # find every sequence of spikes that qualifies as a burst\n",
    "                    burst_thresholds = {d['spiketrain']: d['thresholds']*pq.Hz for d in metadata['burst_detectors']}\n",
    "                    bursts = df.at[i, f'{unit} all bursts'] = _find_bursts(st, burst_thresholds[unit][0], burst_thresholds[unit][1]) # 'at', not 'loc', is important for inserting list into cell\n",
    "\n",
    "                    first_burst_start = np.nan\n",
    "                    last_burst_end = np.nan\n",
    "                    if len(bursts) > 0:\n",
    "\n",
    "                        for burst in zip(bursts.times, bursts.durations, bursts.array_annotations['spikes']):\n",
    "                            if is_good_burst(burst):\n",
    "                                time, duration, n_spikes = burst\n",
    "                                first_burst_start = time\n",
    "                                break # quit after finding first good burst\n",
    "\n",
    "                        for burst in zip(reversed(bursts.times), reversed(bursts.durations), reversed(bursts.array_annotations['spikes'])):\n",
    "                            if is_good_burst(burst):\n",
    "                                time, duration, n_spikes = burst\n",
    "                                last_burst_end = time + duration\n",
    "                                break # quit after finding first (actually, last) good burst\n",
    "\n",
    "                    # merge the first and last good bursts and anything in between\n",
    "                    if np.isfinite(first_burst_start) and np.isfinite(last_burst_end):\n",
    "                        st_burst = st.time_slice(first_burst_start, last_burst_end)\n",
    "                        burst_start = df.loc[i, f'{unit} burst start (s)'] = first_burst_start\n",
    "                        burst_end = df.loc[i, f'{unit} burst end (s)'] = last_burst_end\n",
    "                        burst_duration = df.loc[i, f'{unit} burst duration (s)'] = (burst_end-burst_start)\n",
    "                        burst_spike_count = df.loc[i, f'{unit} burst spike count'] = st_burst.size\n",
    "                        if burst_spike_count > 1:\n",
    "                            burst_mean_freq = df.loc[i, f'{unit} burst mean frequency (Hz)'] = ((burst_spike_count-1)/burst_duration).rescale('Hz')\n",
    "                        if video_is_segmented:\n",
    "                            df.loc[i, f'{unit} burst start (video seg normalized)'] = normalize_time(video_segmentation_times.magnitude, float(first_burst_start))\n",
    "                            df.loc[i, f'{unit} burst end (video seg normalized)']   = normalize_time(video_segmentation_times.magnitude, float(last_burst_end))\n",
    "                        if force_is_segmented:\n",
    "                            df.loc[i, f'{unit} burst start (force seg normalized)'] = normalize_time(force_segmentation_times.magnitude, float(first_burst_start))\n",
    "                            df.loc[i, f'{unit} burst end (force seg normalized)']   = normalize_time(force_segmentation_times.magnitude, float(last_burst_end))\n",
    "\n",
    "        # B3/B6/B9\n",
    "        df['B3/B6/B9 burst start (s)'] = np.nan\n",
    "        df['B3/B6/B9 burst end (s)'] = np.nan\n",
    "        df['B3/B6/B9 burst duration (s)'] = 0\n",
    "        df['B3/B6/B9 burst spike count'] = 0\n",
    "        df['B3/B6/B9 burst mean frequency (Hz)'] = 0\n",
    "        b3b6b9_burst_start = df['B3/B6/B9 burst start (s)'] = df[['B6/B9 burst start (s)', 'B3 burst start (s)']].min(axis=1)\n",
    "        b3b6b9_burst_end = df['B3/B6/B9 burst end (s)']   = df[['B6/B9 burst end (s)',   'B3 burst end (s)']]  .max(axis=1)\n",
    "        b3b6b9_burst_duration = df['B3/B6/B9 burst duration (s)'] = b3b6b9_burst_end - b3b6b9_burst_start\n",
    "        b3b6b9_burst_spike_count = df['B3/B6/B9 burst spike count'] = df['B6/B9 burst spike count'] + df['B3 burst spike count']\n",
    "        b3b6b9_burst_mean_freq = df['B3/B6/B9 burst mean frequency (Hz)'] = (b3b6b9_burst_spike_count-1)/b3b6b9_burst_duration\n",
    "\n",
    "\n",
    "    ###\n",
    "    ### FINISH\n",
    "    ###\n",
    "\n",
    "    # index the table on 4 variables so that this dataframe can later be merged with others\n",
    "    df['Animal'] = animal\n",
    "    df['Food'] = food\n",
    "    df['Bout_index'] = bout_index\n",
    "    df = df.reset_index().set_index(['Animal', 'Food', 'Bout_index', 'Behavior_index'])\n",
    "\n",
    "    df_list += [df]\n",
    "\n",
    "    pbar.update()\n",
    "\n",
    "pbar.close()\n",
    "\n",
    "df_all = pd.concat(df_list, sort=False).sort_index()\n",
    "\n",
    "if exemplary_swallow in feeding_bouts:\n",
    "    # move exemplar to separate dataframe\n",
    "    df_exemplary_swallow = df_all.loc[exemplary_swallow].copy()\n",
    "    df_all = df_all.drop(exemplary_swallow)\n",
    "\n",
    "end = datetime.datetime.now()\n",
    "print('elapsed time:', end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🤪 Sanity Checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "skip_sanity_checks = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if not skip_sanity_checks:\n",
    "    start = datetime.datetime.now()\n",
    "\n",
    "    # reconstruct the original df_all, before exemplary_swallow was removed\n",
    "    df_list2 = [df_all]\n",
    "    if exemplary_swallow in feeding_bouts:\n",
    "        df_exemplary_swallow2 = df_exemplary_swallow.copy()\n",
    "        df_exemplary_swallow2['Animal'] = exemplary_swallow[0]\n",
    "        df_exemplary_swallow2['Food'] = exemplary_swallow[1]\n",
    "        df_exemplary_swallow2['Bout_index'] = exemplary_swallow[2]\n",
    "        df_exemplary_swallow2 = df_exemplary_swallow2.reset_index().set_index(['Animal', 'Food', 'Bout_index', 'Behavior_index'])\n",
    "        df_list2 += [df_exemplary_swallow2]\n",
    "    df_all2 = pd.concat(df_list2, sort=False).sort_index()\n",
    "\n",
    "    # use Neo RawIO lazy loading to load much faster and using less memory\n",
    "    # - with lazy=True, filtering parameters specified in metadata are ignored\n",
    "    #     - note: filters are replaced below and applied manually anyway\n",
    "    # - with lazy=True, loading via time_slice requires neo>=0.8.0\n",
    "    lazy = True\n",
    "\n",
    "    # load the metadata containing file paths\n",
    "    metadata = neurotic.MetadataSelector(file=metadata_file)\n",
    "\n",
    "    last_data_set_name = None\n",
    "    pbar = tqdm(total=len(feeding_bouts), unit='figure')\n",
    "    for (animal, food, bout_index), (data_set_name, time_window) in feeding_bouts.items():\n",
    "\n",
    "        channel_names = channel_names_by_animal[animal]\n",
    "        epoch_types = epoch_types_by_food[food]\n",
    "\n",
    "        df = df_all2.loc[animal, food, bout_index]\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### LOAD DATASET\n",
    "        ###\n",
    "\n",
    "        metadata.select(data_set_name)\n",
    "\n",
    "        if data_set_name is last_data_set_name:\n",
    "            # skip reloading the data if it's already in memory\n",
    "            pass\n",
    "        else:\n",
    "            blk = neurotic.load_dataset(metadata, lazy=lazy)\n",
    "            if lazy:\n",
    "                # manually perform filters\n",
    "                blk = apply_filters(blk, metadata)\n",
    "\n",
    "        last_data_set_name = data_set_name\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### START FIGURE\n",
    "        ###\n",
    "\n",
    "#             figsize = (9.5, 10) # dimensions for notebook\n",
    "#             figsize = (11, 8.5) # dimensions for printing\n",
    "        figsize = (16, 9) # dimensions for filling wide screens\n",
    "        fig, axes = plt.subplots(len(channel_names), 1, sharex=True, figsize=figsize)\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### PLOT SIGNALS\n",
    "        ###\n",
    "\n",
    "        # plot all channels for entire time window\n",
    "        for i, channel in enumerate(channel_names):\n",
    "            plt.sca(axes[i])\n",
    "            sig = get_sig(blk, channel)\n",
    "            sig = sig.time_slice(time_window[0]*pq.s, time_window[1]*pq.s)\n",
    "            sig = sig.rescale(channel_units[i])\n",
    "            plt.plot(sig.times, sig.magnitude, c='0.8', lw=1, zorder=-1)\n",
    "\n",
    "            if i == 0:\n",
    "                plt.title(f'({animal}, {food}, {bout_index}): {data_set_name}')\n",
    "\n",
    "            plt.ylabel(sig.name + ' (' + sig.units.dimensionality.string + ')')\n",
    "            axes[i].yaxis.set_label_coords(-0.06, 0.5)\n",
    "\n",
    "            if i < len(channel_names)-1:\n",
    "                # remove right, top, and bottom plot borders, and remove x-axis\n",
    "                sns.despine(ax=plt.gca(), bottom=True)\n",
    "                plt.gca().xaxis.set_visible(False)\n",
    "            else:\n",
    "                # remove right and top plot borders, and set x-label\n",
    "                sns.despine(ax=plt.gca())\n",
    "                plt.xlabel('Time (s)')\n",
    "\n",
    "        # plot smoothed force for entire time window\n",
    "        channel = 'Force'\n",
    "        sig = get_sig(blk, channel)\n",
    "        if lazy:\n",
    "            sig = sig.time_slice(None, None)\n",
    "        sig = elephant.signal_processing.butter(sig, lowpass_freq = 5*pq.Hz)\n",
    "        sig = sig.time_slice(time_window[0]*pq.s, time_window[1]*pq.s)\n",
    "        sig = sig.rescale('mN')\n",
    "        plt.plot(sig.times, sig.magnitude, c='k', lw=1, zorder=0)\n",
    "        force_smoothed_sig = sig\n",
    "\n",
    "\n",
    "\n",
    "        # iterate over all swallows\n",
    "        for j, i in enumerate(df.index):\n",
    "\n",
    "            behavior_start = df.loc[i, 'Start (s)']*pq.s\n",
    "            behavior_end   = df.loc[i, 'End (s)']*pq.s\n",
    "            \n",
    "            ###\n",
    "            ### MOVEMENTS\n",
    "            ###\n",
    "            \n",
    "            # plot inward food movement\n",
    "            inward_movement_start = df.loc[i, 'Inward movement start (s)']*pq.s\n",
    "            inward_movement_end   = df.loc[i, 'Inward movement end (s)']*pq.s\n",
    "            if np.isfinite(inward_movement_start):\n",
    "                channel = 'Force'\n",
    "                ax = axes[channel_names.index(channel)]\n",
    "                ax.axvspan(\n",
    "                    inward_movement_start, inward_movement_end,\n",
    "                    0.99, 1,\n",
    "                    facecolor='k', edgecolor=None, lw=0)\n",
    "            \n",
    "            \n",
    "\n",
    "            ###\n",
    "            ### FORCE SEGMENTATION\n",
    "            ###\n",
    "\n",
    "            prev_force_drop_end = df.loc[i, 'Previous force drop end (s)']*pq.s   # start of previous \"Force drop end\" epoch\n",
    "            force_shoulder_end  = df.loc[i, 'Force shoulder end start (s)']*pq.s  # start of \"Force shoulder end\" epoch\n",
    "            force_rise_start    = df.loc[i, 'Force rise start start (s)']*pq.s    # start of \"Force rise start\" epoch\n",
    "            force_plateau_start = df.loc[i, 'Force plateau start start (s)']*pq.s # start of \"Force plateau start\" epoch\n",
    "            force_plateau_end   = df.loc[i, 'Force plateau end start (s)']*pq.s   # start of \"Force plateau end\" epoch\n",
    "            force_drop_end      = df.loc[i, 'Force drop end start (s)']*pq.s      # start of \"Force drop end\" epoch\n",
    "\n",
    "            # force rise start, plateau start and end, and drop end are required\n",
    "            force_is_segmented = np.all(np.isfinite(np.array([\n",
    "                force_rise_start, force_plateau_start, force_plateau_end, force_drop_end])))\n",
    "\n",
    "            force_segmentation_times = df.at[i, 'Force segmentation times (s)']\n",
    "\n",
    "            if force_is_segmented:\n",
    "\n",
    "                plt.sca(axes[channel_names.index('Force')])\n",
    "                \n",
    "                force_min_time = df.loc[i, 'Force minimum time (s)']\n",
    "                force_min = df.loc[i, 'Force minimum (mN)']\n",
    "                force_peak_time = df.loc[i, 'Force peak time (s)']\n",
    "                force_peak = df.loc[i, 'Force peak (mN)']\n",
    "                force_baseline = df.loc[i, 'Force baseline (mN)']\n",
    "\n",
    "                force_plateau_start_value = df.loc[i, 'Force plateau start value (mN)']\n",
    "                force_plateau_end_value = df.loc[i, 'Force plateau end value (mN)']\n",
    "                force_drop_end_value = df.loc[i, 'Force drop end value (mN)']\n",
    "                force_shoulder_end_value = df.loc[i, 'Force shoulder end value (mN)']\n",
    "\n",
    "                # get smoothed force for whole behavior for remaining force calculations\n",
    "                sig = force_smoothed_sig\n",
    "                sig = sig.time_slice(prev_force_drop_end - 0.01*pq.s, force_drop_end + 0.01*pq.s)\n",
    "                sig = sig.rescale('mN')\n",
    "\n",
    "                # plot force shoulder in color\n",
    "                if np.isfinite(force_shoulder_end):\n",
    "                    sig2 = sig.time_slice(prev_force_drop_end, force_shoulder_end)\n",
    "                    plt.plot(sig2.times, sig2.magnitude, c=force_colors['shoulder'], lw=2, zorder=1)\n",
    "                    \n",
    "                # plot force rise in color\n",
    "                sig2 = sig.time_slice(force_rise_start, force_plateau_start)\n",
    "                plt.plot(sig2.times, sig2.magnitude, c=force_colors['rise'], lw=2, zorder=1)\n",
    "\n",
    "                # plot force plateau in color\n",
    "                sig2 = sig.time_slice(force_plateau_start, force_plateau_end)\n",
    "                plt.plot(sig2.times, sig2.magnitude, c=force_colors['plateau'], lw=2, zorder=1)\n",
    "\n",
    "                # plot force peak, baseline, and plateau values\n",
    "                plt.plot([force_peak_time],     [force_peak],                marker=CARETDOWN,  markersize=5, color='k')\n",
    "#                 plt.plot([force_min_time],      [force_min],                 marker=CARETUP,    markersize=5, color='k')\n",
    "                plt.plot([force_rise_start],    [force_baseline],            marker=CARETUP,    markersize=5, color='k')\n",
    "                plt.plot([force_plateau_start], [force_plateau_start_value], marker=CARETRIGHT, markersize=5, color='k')\n",
    "                plt.plot([force_plateau_end],   [force_plateau_end_value],   marker=CARETLEFT,  markersize=5, color='k')\n",
    "\n",
    "                # plot segmentation boundaries\n",
    "                for (t, y, c) in [\n",
    "                        (force_shoulder_end,  force_shoulder_end_value,  force_colors['shoulder']),\n",
    "                        (force_rise_start,    force_baseline,            force_colors['rise']),\n",
    "                        (force_plateau_start, force_plateau_start_value, force_colors['plateau']),\n",
    "                        (force_plateau_end,   force_plateau_end_value,   force_colors['plateau']),\n",
    "                        (force_drop_end,      force_drop_end_value,      force_colors['drop'])]:\n",
    "                    if np.isfinite(y):\n",
    "                        axes[-1].add_artist(patches.ConnectionPatch(\n",
    "                            xyA=(t, y), xyB=(t, 0),\n",
    "                            coordsA='data', coordsB=axes[-1].get_xaxis_transform(),\n",
    "                            axesA=axes[-1], axesB=axes[-1],\n",
    "                            color='0.8', lw=1, ls=':', zorder=-2))\n",
    "\n",
    "\n",
    "\n",
    "            ###\n",
    "            ### SPIKES AND BURSTS\n",
    "            ###\n",
    "\n",
    "            for k, unit in enumerate(units):\n",
    "                st = df.loc[i, f'{unit} spike train']\n",
    "                if st is not None and st.size > 0:\n",
    "\n",
    "                    # get the signal for the behavior with 10 seconds cushion for spikes outside the behavior duration\n",
    "                    channel = st.annotations['channels'][0]\n",
    "                    sig = get_sig(blk, channel)\n",
    "                    sig = sig.time_slice(max(sig.t_start, behavior_start - 10*pq.s), min(sig.t_stop, behavior_end + 10*pq.s))\n",
    "                    sig = sig.rescale(channel_units[channel_names.index(channel)])\n",
    "\n",
    "                    # plot spikes\n",
    "                    plt.sca(axes[channel_names.index(channel)])\n",
    "                    marker = ['.', 'x'][j%2] # alternate markers between behaviors\n",
    "                    spike_amplitudes = np.array([sig[sig.time_index(t)] for t in st]) * pq.Quantity(sig.units)\n",
    "                    plt.scatter(st.times.rescale('s'), spike_amplitudes, marker=marker, c=unit_colors[unit])\n",
    "\n",
    "                    # plot burst windows\n",
    "                    discriminator = next((d for d in metadata['amplitude_discriminators'] if d['name'] == st.name), None)\n",
    "                    if discriminator is None:\n",
    "                        raise Exception(f'For data set \"{data_set_name}\", discriminator \"{st.name}\" could not be found')\n",
    "                    bottom = pq.Quantity(discriminator['amplitude'][0], discriminator['units']).rescale(sig.units)\n",
    "                    top = pq.Quantity(discriminator['amplitude'][1], discriminator['units']).rescale(sig.units)\n",
    "                    height = top-bottom\n",
    "                    bursts = df.at[i, f'{unit} all bursts']\n",
    "                    for burst in zip(bursts.times, bursts.durations, bursts.array_annotations['spikes']):\n",
    "                        time, duration, n_spikes = burst\n",
    "                        left = time\n",
    "                        right = time + duration\n",
    "                        width = right-left\n",
    "                        linestyle = '-' if is_good_burst(burst) else '--'\n",
    "                        rect = patches.Rectangle((left, bottom), width, height, ls=linestyle, edgecolor=unit_colors[unit], fill=False)\n",
    "                        plt.gca().add_patch(rect)\n",
    "\n",
    "                    # plot markers for edges of bursts\n",
    "                    burst_start = df.loc[i, f'{unit} burst start (s)']\n",
    "                    burst_end = df.loc[i, f'{unit} burst end (s)']\n",
    "                    if top > 0:\n",
    "                        plt.plot([burst_start], [top],    marker=CARETDOWN, markersize=5, color='k')\n",
    "                        plt.plot([burst_end],   [top],    marker=CARETDOWN, markersize=5, color='k')\n",
    "                    else:\n",
    "                        plt.plot([burst_start], [bottom], marker=CARETUP,   markersize=5, color='k')\n",
    "                        plt.plot([burst_end],   [bottom], marker=CARETUP,   markersize=5, color='k')\n",
    "\n",
    "\n",
    "\n",
    "        ###\n",
    "        ### FINISH FIGURE\n",
    "        ###\n",
    "\n",
    "        # optimize plot margins\n",
    "        plt.subplots_adjust(\n",
    "            left   = 0.1,\n",
    "            right  = 0.99,\n",
    "            top    = 0.96,\n",
    "            bottom = 0.06,\n",
    "            hspace = 0.15,\n",
    "        )\n",
    "\n",
    "        # export figure\n",
    "        export_dir2 = os.path.join(export_dir, 'sanity-checks')\n",
    "        if not os.path.exists(export_dir2):\n",
    "            os.mkdir(export_dir2)\n",
    "        plt.gcf().savefig(os.path.join(export_dir2, f'{animal} {food} {bout_index}.png'), dpi=300)\n",
    "        \n",
    "        pbar.update()\n",
    "        \n",
    "    pbar.close()\n",
    "\n",
    "    if exemplary_swallow in feeding_bouts:\n",
    "        # rename output file for exemplar\n",
    "        animal, food, bout_index = exemplary_swallow\n",
    "        old_path = os.path.join(export_dir2, f'{animal} {food} {bout_index}.png')\n",
    "        new_path = os.path.join(export_dir2, 'exemplary_swallow.png')\n",
    "        if os.path.exists(old_path):\n",
    "            if os.path.exists(new_path):\n",
    "                os.remove(new_path)\n",
    "            os.rename(old_path, new_path)\n",
    "\n",
    "    del df_all2, df_list2\n",
    "\n",
    "    end = datetime.datetime.now()\n",
    "    print('elapsed time:', end-start)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if not skip_sanity_checks:\n",
    "    start = datetime.datetime.now()\n",
    "\n",
    "    if exemplary_swallow in feeding_bouts:\n",
    "        pbar = tqdm(total=len(feeding_bouts)-1, unit='figure')\n",
    "    else:\n",
    "        pbar = tqdm(total=len(feeding_bouts), unit='figure')\n",
    "    for (animal, food, bout_index), (data_set_name, time_window) in feeding_bouts.items():\n",
    "\n",
    "        channel_names = channel_names_by_animal[animal]\n",
    "\n",
    "        if (animal, food, bout_index) == exemplary_swallow:\n",
    "            # skip the lone swallow\n",
    "            continue\n",
    "        else:\n",
    "            df = df_all.loc[animal, food, bout_index]\n",
    "\n",
    "        # load the data\n",
    "        metadata = neurotic.MetadataSelector(metadata_file)\n",
    "        metadata.select(data_set_name)\n",
    "        blk = neurotic.load_dataset(metadata, lazy=True)\n",
    "\n",
    "        # figsize = (9.5, 10) # dimensions for notebook\n",
    "        # figsize = (11, 8.5) # dimensions for printing\n",
    "        # figsize = (16, 9) # dimensions for filling wide screens\n",
    "        figsize = (20, 9)\n",
    "        fig, axes = plt.subplots(len(units)+1, 3, sharex='col', figsize=figsize)\n",
    "\n",
    "        for k, unit in enumerate(units):\n",
    "            # get the subplot axes handles\n",
    "            ax_real_time, ax_force_seg, ax_video_seg = axes[k]\n",
    "\n",
    "            # set y-axis label\n",
    "            ax_real_time.set_ylabel(f'{unit} (Hz)')\n",
    "            ax_real_time.yaxis.set_label_coords(-0.06, 0.5)\n",
    "\n",
    "            # remove right, top, and bottom plot borders, and remove x-axis\n",
    "            sns.despine(ax=ax_real_time, bottom=True)\n",
    "            sns.despine(ax=ax_force_seg, bottom=True)\n",
    "            sns.despine(ax=ax_video_seg, bottom=True)\n",
    "            ax_real_time.xaxis.set_visible(False)\n",
    "            ax_force_seg.xaxis.set_visible(False)\n",
    "            ax_video_seg.xaxis.set_visible(False)\n",
    "\n",
    "            # set time plot ranges\n",
    "            ax_real_time.set_xlim([time_window[0]-5, time_window[1]+5])\n",
    "            ax_force_seg.set_xlim([force_seg_interp_times.min(), force_seg_interp_times.max()])\n",
    "#             ax_video_seg.set_xlim([video_seg_interp_times.min(), video_seg_interp_times.max()])\n",
    "            ax_video_seg.set_xlim([2, 6.5])\n",
    "\n",
    "            # elevate the Axes for units and remove background colors so that\n",
    "            # each vertical ConnectionPatch drawn later is visible behind it\n",
    "            ax_real_time.set_zorder(1)\n",
    "            ax_force_seg.set_zorder(1)\n",
    "            ax_video_seg.set_zorder(1)\n",
    "            ax_real_time.set_facecolor('none')\n",
    "            ax_force_seg.set_facecolor('none')\n",
    "            ax_video_seg.set_facecolor('none')\n",
    "\n",
    "        # remove right and top plot borders from bottom panels, and set x-label\n",
    "        ax_real_time, ax_force_seg, ax_video_seg = axes[-1]\n",
    "        sns.despine(ax=ax_real_time)\n",
    "        sns.despine(ax=ax_force_seg)\n",
    "        sns.despine(ax=ax_video_seg)\n",
    "        ax_real_time.set_xlabel('Time (s)')\n",
    "        ax_force_seg.set_xlabel('Time (normalized using force segmentation)')\n",
    "        ax_video_seg.set_xlabel('Time (normalized using video segmentation)')\n",
    "\n",
    "        # set time plot ranges\n",
    "        ax_real_time.set_xlim([time_window[0]-5, time_window[1]+5])\n",
    "        ax_force_seg.set_xlim([force_seg_interp_times.min(), force_seg_interp_times.max()])\n",
    "#         ax_video_seg.set_xlim([video_seg_interp_times.min(), video_seg_interp_times.max()])\n",
    "        ax_video_seg.set_xlim([2, 6.5])\n",
    "\n",
    "        # plot force in real time\n",
    "        ax_real_time, ax_force_seg, ax_video_seg = axes[-1]\n",
    "        channel = 'Force'\n",
    "        sig = get_sig(blk, channel)\n",
    "        sig = sig.time_slice(time_window[0]*pq.s, time_window[1]*pq.s)\n",
    "        sig = sig.rescale(channel_units[channel_names.index(channel)])\n",
    "        ax_real_time.plot(sig.times, sig.magnitude, c='0.8', lw=1)\n",
    "        ax_real_time.set_ylabel(sig.name + ' (' + sig.units.dimensionality.string + ')')\n",
    "        ax_real_time.yaxis.set_label_coords(-0.06, 0.5)\n",
    "\n",
    "        all_force_seg_normalized_times_series = {}\n",
    "        all_video_seg_normalized_times_series = {}\n",
    "        for unit in units:\n",
    "            all_force_seg_normalized_times_series[unit] = np.zeros((0, force_seg_interp_times.size))\n",
    "            all_video_seg_normalized_times_series[unit] = np.zeros((0, video_seg_interp_times.size))\n",
    "        all_force_seg_normalized_times_series['Force'] = np.zeros((0, force_seg_interp_times.size))\n",
    "        all_video_seg_normalized_times_series['Force'] = np.zeros((0, video_seg_interp_times.size))\n",
    "        \n",
    "        for j, i in enumerate(df.index):\n",
    "            \n",
    "            \n",
    "            # plot inward food movement\n",
    "            inward_movement_start = df.loc[i, 'Inward movement start (s)']\n",
    "            inward_movement_end   = df.loc[i, 'Inward movement end (s)']\n",
    "            if np.isfinite(inward_movement_start):\n",
    "                ax_real_time, ax_force_seg, ax_video_seg = axes[-1]\n",
    "                ax_real_time.axvspan(\n",
    "                    inward_movement_start, inward_movement_end,\n",
    "                    0.98, 1,\n",
    "                    facecolor='k', edgecolor=None, lw=0)\n",
    "            \n",
    "            \n",
    "            for k, unit in enumerate(units):\n",
    "                ax_real_time, ax_force_seg, ax_video_seg = axes[k]\n",
    "\n",
    "\n",
    "                # raster plot\n",
    "                st = df.loc[i, f'{unit} spike train']\n",
    "                if st is not None:\n",
    "                    ax_real_time.eventplot(positions=st, lineoffsets=-1, colors=unit_colors[unit])\n",
    "\n",
    "\n",
    "                # plot the firing rates in real time\n",
    "                firing_rate = df.loc[i, f'{unit} firing rate (Hz)']\n",
    "                if firing_rate is not None:\n",
    "                    ax_real_time.plot(firing_rate.times.rescale('s'), firing_rate, c=unit_colors[unit])\n",
    "\n",
    "\n",
    "                # plot firing rates in normalized time\n",
    "                firing_rate_force_seg_interp = df.loc[i, f'{unit} firing rate, force segmented interpolation (Hz)']\n",
    "                if firing_rate_force_seg_interp is not None:\n",
    "                    all_force_seg_normalized_times_series[unit] = np.concatenate([all_force_seg_normalized_times_series[unit], firing_rate_force_seg_interp[np.newaxis, :]])\n",
    "                    ax_force_seg.plot(force_seg_interp_times, firing_rate_force_seg_interp, c=unit_colors[unit])\n",
    "                firing_rate_video_seg_interp = df.loc[i, f'{unit} firing rate, video segmented interpolation (Hz)']\n",
    "                if firing_rate_video_seg_interp is not None:\n",
    "                    all_video_seg_normalized_times_series[unit] = np.concatenate([all_video_seg_normalized_times_series[unit], firing_rate_video_seg_interp[np.newaxis, :]])\n",
    "                    ax_video_seg.plot(video_seg_interp_times, firing_rate_video_seg_interp, c=unit_colors[unit])\n",
    "\n",
    "\n",
    "            # plot force in normalized time\n",
    "            ax_real_time, ax_force_seg, ax_video_seg = axes[-1]\n",
    "            force_force_seg_interp = df.at[i, 'Force, force segmented interpolation (mN)']\n",
    "            if force_force_seg_interp is not None:\n",
    "                all_force_seg_normalized_times_series['Force'] = np.concatenate([all_force_seg_normalized_times_series['Force'], force_force_seg_interp[np.newaxis, :]])\n",
    "                ax_force_seg.plot(force_seg_interp_times, force_force_seg_interp, c='0.8', lw=1)\n",
    "            force_video_seg_interp = df.at[i, 'Force, video segmented interpolation (mN)']\n",
    "            if force_video_seg_interp is not None:\n",
    "                all_video_seg_normalized_times_series['Force'] = np.concatenate([all_video_seg_normalized_times_series['Force'], force_video_seg_interp[np.newaxis, :]])\n",
    "                ax_video_seg.plot(video_seg_interp_times, force_video_seg_interp, c='0.8', lw=1)\n",
    "\n",
    "\n",
    "            # plot force phase boundaries in real time\n",
    "            force_segmentation_times = df.at[i, 'Force segmentation times (s)'].rescale('s').magnitude\n",
    "            for m, t in enumerate(force_segmentation_times[3:8]): # 3 = end of shoulder, 8 = end of next shoulder\n",
    "                if np.isfinite(t):\n",
    "                    if m == 1: # 1 = start of rise\n",
    "                        color = force_colors['rise']\n",
    "                    else:\n",
    "                        color = '0.75'\n",
    "                    axes[-1][0].add_artist(patches.ConnectionPatch(\n",
    "                        xyA=(t, 0), xyB=(t, 1),\n",
    "                        coordsA=axes[-1][0].get_xaxis_transform(), coordsB=axes[0][0].get_xaxis_transform(),\n",
    "                        axesA=axes[-1][0], axesB=axes[0][0],\n",
    "                        color=color, lw=1, ls=':'))\n",
    "\n",
    "\n",
    "        # plot force phase boundaries in normalized time\n",
    "        for m in range(len(force_segmentation_times)):\n",
    "            if m == 4: # 4 = start of rise\n",
    "                color = force_colors['rise']\n",
    "            else:\n",
    "                color = '0.75'\n",
    "            axes[-1][1].add_artist(patches.ConnectionPatch(\n",
    "                xyA=(m, 0), xyB=(m, 1),\n",
    "                coordsA=axes[-1][1].get_xaxis_transform(), coordsB=axes[0][1].get_xaxis_transform(),\n",
    "                axesA=axes[-1][1], axesB=axes[0][1],\n",
    "                color=color, lw=1, ls=':'))\n",
    "        \n",
    "        # plot video phase boundaries in normalized time\n",
    "        video_segmentation_times = df.at[i, 'Video segmentation times (s)'].rescale('s').magnitude # grab last swallow's as an example\n",
    "#         for m in range(len(video_segmentation_times)):\n",
    "        for m in [2, 3, 4, 5, 6]:\n",
    "            if m == 4: # 4 = start of inward movement\n",
    "                color = force_colors['rise']\n",
    "            else:\n",
    "                color = '0.75'\n",
    "            axes[-1][2].add_artist(patches.ConnectionPatch(\n",
    "                xyA=(m, 0), xyB=(m, 1),\n",
    "                coordsA=axes[-1][2].get_xaxis_transform(), coordsB=axes[0][2].get_xaxis_transform(),\n",
    "                axesA=axes[-1][2], axesB=axes[0][2],\n",
    "                color=color, lw=1, ls=':'))\n",
    "\n",
    "\n",
    "        # plot firing rate distributions\n",
    "        for k, unit in enumerate(units):\n",
    "            ax_real_time, ax_force_seg, ax_video_seg = axes[k]\n",
    "\n",
    "            firing_rate_median = np.nanmedian(all_force_seg_normalized_times_series[unit], axis=0)\n",
    "            firing_rate_q1 = np.nanquantile(all_force_seg_normalized_times_series[unit], q=0.25, axis=0)\n",
    "            firing_rate_q3 = np.nanquantile(all_force_seg_normalized_times_series[unit], q=0.75, axis=0)\n",
    "            ax_force_seg.plot(force_seg_interp_times, firing_rate_median, c='k', lw=2, zorder=3)\n",
    "            ax_force_seg.plot(force_seg_interp_times, firing_rate_q1, c='k', lw=2, ls='--')\n",
    "            ax_force_seg.plot(force_seg_interp_times, firing_rate_q3, c='k', lw=2, ls='--')\n",
    "            \n",
    "            firing_rate_median = np.nanmedian(all_video_seg_normalized_times_series[unit], axis=0)\n",
    "            firing_rate_q1 = np.nanquantile(all_video_seg_normalized_times_series[unit], q=0.25, axis=0)\n",
    "            firing_rate_q3 = np.nanquantile(all_video_seg_normalized_times_series[unit], q=0.75, axis=0)\n",
    "            ax_video_seg.plot(video_seg_interp_times, firing_rate_median, c='k', lw=2, zorder=3)\n",
    "            ax_video_seg.plot(video_seg_interp_times, firing_rate_q1, c='k', lw=2, ls='--')\n",
    "            ax_video_seg.plot(video_seg_interp_times, firing_rate_q3, c='k', lw=2, ls='--')\n",
    "\n",
    "\n",
    "        # plot force distribution\n",
    "        ax_real_time, ax_force_seg, ax_video_seg = axes[-1]\n",
    "\n",
    "        force_median = np.nanmedian(all_force_seg_normalized_times_series['Force'], axis=0)\n",
    "        force_q1 = np.nanquantile(all_force_seg_normalized_times_series['Force'], q=0.25, axis=0)\n",
    "        force_q3 = np.nanquantile(all_force_seg_normalized_times_series['Force'], q=0.75, axis=0)\n",
    "        ax_force_seg.plot(force_seg_interp_times, force_median, c='k', lw=2, zorder=3)\n",
    "        ax_force_seg.plot(force_seg_interp_times, force_q1, c='k', lw=2, ls='--')\n",
    "        ax_force_seg.plot(force_seg_interp_times, force_q3, c='k', lw=2, ls='--')\n",
    "        \n",
    "        force_median = np.nanmedian(all_video_seg_normalized_times_series['Force'], axis=0)\n",
    "        force_q1 = np.nanquantile(all_video_seg_normalized_times_series['Force'], q=0.25, axis=0)\n",
    "        force_q3 = np.nanquantile(all_video_seg_normalized_times_series['Force'], q=0.75, axis=0)\n",
    "        ax_video_seg.plot(video_seg_interp_times, force_median, c='k', lw=2, zorder=3)\n",
    "        ax_video_seg.plot(video_seg_interp_times, force_q1, c='k', lw=2, ls='--')\n",
    "        ax_video_seg.plot(video_seg_interp_times, force_q3, c='k', lw=2, ls='--')\n",
    "\n",
    "\n",
    "        plt.suptitle(f'({animal}, {food}, {bout_index}): {data_set_name}')\n",
    "        plt.tight_layout(rect=(0, 0, 1, 0.97))\n",
    "\n",
    "        # export figure\n",
    "        export_dir3 = os.path.join(export_dir, 'sanity-checks-firing-rates')\n",
    "        if not os.path.exists(export_dir3):\n",
    "            os.mkdir(export_dir3)\n",
    "        plt.gcf().savefig(os.path.join(export_dir3, f'{animal} {food} {bout_index}.png'), dpi=300)\n",
    "        \n",
    "        pbar.update()\n",
    "    \n",
    "    pbar.close()\n",
    "\n",
    "    end = datetime.datetime.now()\n",
    "    print('elapsed time:', end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [FIGURE 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 1A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = datetime.datetime.now()\n",
    "\n",
    "df = df_exemplary_swallow\n",
    "(data_set_name, time_window) = feeding_bouts[exemplary_swallow]\n",
    "animal, _, _ = exemplary_swallow\n",
    "\n",
    "t_start, t_stop = time_window*pq.s\n",
    "plots = [\n",
    "    {'channel': 'I2',       'units': 'uV', 'ylim': [ -35,  35], 'scalebar': 50}, #, 'decimation_factor': 100},\n",
    "    {'channel': 'RN',       'units': 'uV', 'ylim': [ -25,  25], 'scalebar': 25}, #, 'decimation_factor': 100},\n",
    "    {'channel': 'BN2',      'units': 'uV', 'ylim': [ -45,  45], 'scalebar': 50}, #, 'decimation_factor': 100},\n",
    "    {'channel': 'BN3-DIST', 'units': 'uV', 'ylim': [ -70,  70], 'scalebar': 50, 'ylabel': 'BN3'}, #, 'decimation_factor': 100},\n",
    "    {'channel': 'Force',    'units': 'mN', 'ylim': [-100, 400], 'scalebar': 200}, #, 'decimation_factor': 100},\n",
    "]\n",
    "plot_names = [p['channel'] for p in plots]\n",
    "plot_units = [p['units'] for p in plots]\n",
    "\n",
    "kwargs = dict(\n",
    "    figsize = (5, 6.5),\n",
    "    linewidth = 0.5,\n",
    "    x_scalebar = 1*pq.s,\n",
    ")\n",
    "\n",
    "# load the metadata\n",
    "metadata = neurotic.MetadataSelector(metadata_file)\n",
    "metadata.select(data_set_name)\n",
    "\n",
    "# download data that is not already stored locally\n",
    "metadata.download_all_data_files()\n",
    "\n",
    "# load the data\n",
    "blk = neurotic.load_dataset(metadata, lazy=True)\n",
    "\n",
    "# manually perform filters\n",
    "blk = apply_filters(blk, metadata)\n",
    "\n",
    "# plot the signals\n",
    "fig, axes = plot_signals_with_scalebars(blk, t_start, t_stop, plots, **kwargs)\n",
    "\n",
    "unit_burst_boxes = {\n",
    "    'B38':   [-12, 12],\n",
    "    'I2':    [-30, 25],\n",
    "    'B8a/b': [-20, 12],\n",
    "    'B6/B9': [-20, 15],\n",
    "    'B3':    [-45, 35],\n",
    "    'B4/B5': [-60, 55],\n",
    "}\n",
    "\n",
    "for j, i in enumerate(df.index):\n",
    "    for k, unit in enumerate(units):\n",
    "        st = df.loc[i, f'{unit} spike train']\n",
    "        if st is not None and st.size > 0:\n",
    "\n",
    "            # get the neural channel\n",
    "            channel = st.annotations['channels'][0]\n",
    "            sig = get_sig(blk, channel)\n",
    "\n",
    "            # get the signal for the entire bout\n",
    "            sig = sig.time_slice(t_start, t_stop)\n",
    "            sig = sig.rescale(plot_units[plot_names.index(channel)])\n",
    "\n",
    "            # plot spikes\n",
    "            ax = axes[plot_names.index(channel)]\n",
    "            spike_amplitudes = np.array([sig[sig.time_index(t)] for t in st]) * pq.Quantity(sig.units)\n",
    "            ax.scatter(st.times.rescale('s'), spike_amplitudes, marker='.', s=20, c=unit_colors[unit], zorder=3)\n",
    "\n",
    "            # plot burst windows\n",
    "            left = df.loc[i, f'{unit} burst start (s)']\n",
    "            right = df.loc[i, f'{unit} burst end (s)']\n",
    "            if np.isfinite(left) and np.isfinite(right):\n",
    "                width = right-left\n",
    "                bottom, top = unit_burst_boxes[unit]\n",
    "                height = top-bottom\n",
    "                rect = patches.Rectangle((left, bottom), width, height, linewidth=2, ls='-', edgecolor=unit_colors[unit], fill=False, zorder=3, clip_on=False)\n",
    "                ax.add_patch(rect)\n",
    "\n",
    "# plot a zero baseline for force\n",
    "force_zero = -11.5 # mN, average before animal began swallowing, not zero because of DC offset\n",
    "axes[-1].axhline(force_zero, color='0.75', lw=1, ls='--', zorder=-1)\n",
    "\n",
    "# plot force phase boundaries\n",
    "times = df.loc[0, 'Force segmentation times (s)'][2:2+6]\n",
    "for t in times:\n",
    "    axes[-1].axvline(x=t, ymin=0.15, lw=1, ls=':', c='gray', zorder=-1)\n",
    "\n",
    "# add arabic numerals for force phase boundaries\n",
    "axes[-1].annotate('1', xy=(times[0], -0.05), xycoords=('data', 'axes fraction'), ha='center', va='bottom')\n",
    "axes[-1].annotate('2', xy=(times[1], -0.05), xycoords=('data', 'axes fraction'), ha='center', va='bottom')\n",
    "axes[-1].annotate('3', xy=(times[2], -0.05), xycoords=('data', 'axes fraction'), ha='center', va='bottom')\n",
    "axes[-1].annotate('4', xy=(times[3], -0.05), xycoords=('data', 'axes fraction'), ha='center', va='bottom')\n",
    "axes[-1].annotate('5', xy=(times[4], -0.05), xycoords=('data', 'axes fraction'), ha='center', va='bottom')\n",
    "axes[-1].annotate('1', xy=(times[5], -0.05), xycoords=('data', 'axes fraction'), ha='center', va='bottom')\n",
    "\n",
    "# add roman numerals for force phases\n",
    "axes[-1].annotate('I',   xy=(times[0:2].mean(), 1), xycoords=('data', 'axes fraction'), ha='center', va='top')\n",
    "axes[-1].annotate('II',  xy=(times[1:3].mean(), 1), xycoords=('data', 'axes fraction'), ha='center', va='top')\n",
    "axes[-1].annotate('III', xy=(times[2:4].mean(), 1), xycoords=('data', 'axes fraction'), ha='center', va='top')\n",
    "axes[-1].annotate('IV',  xy=(times[3:5].mean(), 1), xycoords=('data', 'axes fraction'), ha='center', va='top')\n",
    "axes[-1].annotate('V',   xy=(times[4:6].mean(), 1), xycoords=('data', 'axes fraction'), ha='center', va='top')\n",
    "\n",
    "# add unit name labels\n",
    "axes[2].annotate('B38',   xy=(2978.15, 0.70), xycoords=('data', 'axes fraction'), ha='center', c=unit_colors['B38'])\n",
    "axes[0].annotate('I2',    xy=(2979.40, 0.75), xycoords=('data', 'axes fraction'), ha='left',   c=unit_colors['I2'])\n",
    "axes[1].annotate('B8a/b', xy=(2981.10, 0.80), xycoords=('data', 'axes fraction'), ha='center', c=unit_colors['B8a/b'])\n",
    "axes[2].annotate('B6/B9', xy=(2980.10, 0.73), xycoords=('data', 'axes fraction'), ha='center', c=unit_colors['B6/B9'])\n",
    "axes[2].annotate('B3',    xy=(2981.85, 0.79), xycoords=('data', 'axes fraction'), ha='left',   c=unit_colors['B3'])\n",
    "axes[3].annotate('B4/B5', xy=(2979.10, 0.80), xycoords=('data', 'axes fraction'), ha='right',  c=unit_colors['B4/B5'])\n",
    "\n",
    "# add protraction box\n",
    "left, right = df.loc[0, ['I2 burst start (s)', 'I2 burst end (s)']]\n",
    "bottom, top = (1, 1.15)\n",
    "width = right-left\n",
    "height = top-bottom\n",
    "rect = patches.Rectangle((left, bottom), width, height, linewidth=2, ls='-', edgecolor='k', fill=False, zorder=3, clip_on=False, transform=axes[0].get_xaxis_transform())\n",
    "axes[0].add_patch(rect)\n",
    "axes[0].annotate('Prot.', xy=((left+right)/2, 1.065), xycoords=('data', 'axes fraction'), ha='center', va='center', c='k', fontsize='small')\n",
    "\n",
    "# add retraction box\n",
    "left, right = df.loc[0, ['I2 burst end (s)', 'End (s)']] # behavior ends with end of B43 burst\n",
    "bottom, top = (1, 1.15)\n",
    "width = right-left\n",
    "height = top-bottom\n",
    "rect = patches.Rectangle((left, bottom), width, height, linewidth=2, ls='-', edgecolor='k', facecolor='k', fill=True, zorder=3, clip_on=False, transform=axes[0].get_xaxis_transform())\n",
    "axes[0].add_patch(rect)\n",
    "axes[0].annotate('Retraction', xy=((left+right)/2, 1.065), xycoords=('data', 'axes fraction'), ha='center', va='center', c='w', fontsize='small')\n",
    "\n",
    "# add inward movement box\n",
    "left, right = df.loc[0, ['Inward movement start (s)', 'Inward movement end (s)']]\n",
    "bottom, top = (1.20, 1.35)\n",
    "width = right-left\n",
    "height = top-bottom\n",
    "rect = patches.Rectangle((left, bottom), width, height, linewidth=2, ls='-', edgecolor='k', facecolor='0.75', fill=True, zorder=3, clip_on=False, transform=axes[0].get_xaxis_transform())\n",
    "axes[0].add_patch(rect)\n",
    "axes[0].annotate('Inward', xy=((left+right)/2, 1.260), xycoords=('data', 'axes fraction'), ha='center', va='center', c='k', fontsize='small')\n",
    "\n",
    "# add markers for video frame times\n",
    "video_times = {\n",
    "    'B': 2980.0,\n",
    "    'C': 2982.5,\n",
    "}\n",
    "for label, video_time in video_times.items():\n",
    "    axes[0].plot([video_time], [1.39], marker=CARETDOWN, markersize=8, color='k', transform=axes[0].get_xaxis_transform(), clip_on=False)\n",
    "    axes[0].annotate(label, xy=(video_time, 1.46), xycoords=('data', 'axes fraction'), ha='center', va='bottom')\n",
    "\n",
    "fig.tight_layout(h_pad=0, w_pad=0, pad=0)\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-1A.png'), dpi=600)\n",
    "\n",
    "end = datetime.datetime.now()\n",
    "print('render time:', end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figures 1B & 1C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Video frames (see code for Figure 1A for times of video frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## [FIGURE 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# shared plot settings\n",
    "inches_per_second = 0.05\n",
    "subplot_height_in_inches = 0.6\n",
    "top_margin_in_inches = 0.12\n",
    "bottom_margin_in_inches = 0.39\n",
    "left_margin_in_inches = 0.72\n",
    "right_margin_in_inches = 0.8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 🐌 Figure 2A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = datetime.datetime.now()\n",
    "\n",
    "data_set_name = 'Animal 4 Unloaded 2 Superset'\n",
    "animal = 'JG12'\n",
    "t_start, t_stop = [223.4, 261.4] * pq.s\n",
    "plots = [\n",
    "    {'channel': 'I2',       'units': 'uV', 'ylim': [ -60,  60], 'scalebar': 50}, #, 'decimation_factor': 400},\n",
    "    {'channel': 'RN',       'units': 'uV', 'ylim': [ -25,  25], 'scalebar': 25}, #, 'decimation_factor': 400},\n",
    "    {'channel': 'BN2',      'units': 'uV', 'ylim': [ -45,  45], 'scalebar': 50}, #, 'decimation_factor': 400},\n",
    "    {'channel': 'BN3-DIST', 'units': 'uV', 'ylim': [ -60,  60], 'scalebar': 50, 'ylabel': 'BN3'}, #, 'decimation_factor': 400},\n",
    "]\n",
    "\n",
    "fig_height_in_inches, bottom_fraction, top_fraction = solve_figure_vertical_dimensions(\n",
    "    len(plots), subplot_height_in_inches, bottom_margin_in_inches, top_margin_in_inches, 0)\n",
    "fig_width_in_inches, left_fraction, right_fraction = solve_figure_horizontal_dimensions(\n",
    "    1, inches_per_second*(t_stop-t_start).magnitude, left_margin_in_inches, right_margin_in_inches, 0)\n",
    "\n",
    "kwargs = dict(\n",
    "    formats = ['png'],\n",
    "    dpi = 600,\n",
    "    figsize = (fig_width_in_inches, fig_height_in_inches),\n",
    "    linewidth = 0.5,\n",
    "    x_scalebar = 10*pq.s,\n",
    "    layout_settings = dict(\n",
    "        left   = left_fraction,\n",
    "        right  = right_fraction,\n",
    "        bottom = bottom_fraction,\n",
    "        top    = top_fraction,\n",
    "        hspace = 0,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# load the metadata\n",
    "metadata = neurotic.MetadataSelector(metadata_file)\n",
    "metadata.select(data_set_name)\n",
    "\n",
    "# download data that is not already stored locally\n",
    "metadata.download_all_data_files()\n",
    "\n",
    "# load the data\n",
    "blk = neurotic.load_dataset(metadata, lazy=True)\n",
    "\n",
    "# manually perform filters\n",
    "blk = apply_filters(blk, metadata)\n",
    "\n",
    "# plot the signals\n",
    "fig, axes = plot_signals_with_scalebars(blk, t_start, t_stop, plots, **kwargs)\n",
    "\n",
    "# plot the inward movements\n",
    "y_inward = 1.03\n",
    "ep = next((ep for ep in blk.segments[0].epochs if ep.name=='Inward movement'), None)\n",
    "assert ep is not None\n",
    "ep = ep.time_slice(t_start, t_stop)\n",
    "for t, dur in zip(ep.times, ep.durations):\n",
    "    left, right = t, t + dur\n",
    "    bottom, top = y_inward-0.025, y_inward+0.025\n",
    "    width = right-left\n",
    "    height = top-bottom\n",
    "    rect = patches.Rectangle((left, bottom), width, height, linewidth=0, facecolor='k', fill=True, clip_on=False, transform=axes[0].get_xaxis_transform())\n",
    "    axes[0].add_patch(rect)\n",
    "\n",
    "# add bite/swallow markers\n",
    "behavior_markers = [\n",
    "    (226.4, 'B'),\n",
    "    (246.3, 'S'),\n",
    "    (250.6, 'S'),\n",
    "    (256.2, 'S'),\n",
    "#     (263.1, 'B'),\n",
    "]\n",
    "for t, label in behavior_markers:\n",
    "    axes[0].annotate(\n",
    "        label,\n",
    "        xy=(t, y_inward), xycoords=axes[0].get_xaxis_transform(),\n",
    "        ha='center', va='center', fontsize='small',\n",
    "    )\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-2A.png'), dpi=600)\n",
    "\n",
    "end = datetime.datetime.now()\n",
    "print('render time:', end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### 🐌 Figure 2B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start = datetime.datetime.now()\n",
    "\n",
    "data_set_name = 'Animal 4 Loaded 2 Superset'\n",
    "animal = 'JG12'\n",
    "t_start, t_stop = [2875.3, 3039.3] * pq.s\n",
    "plots = [\n",
    "    {'channel': 'I2',       'units': 'uV', 'ylim': [ -60,  60], 'scalebar': 50}, #, 'decimation_factor': 400},\n",
    "    {'channel': 'RN',       'units': 'uV', 'ylim': [ -25,  25], 'scalebar': 25}, #, 'decimation_factor': 400},\n",
    "    {'channel': 'BN2',      'units': 'uV', 'ylim': [ -45,  45], 'scalebar': 50}, #, 'decimation_factor': 400},\n",
    "    {'channel': 'BN3-DIST', 'units': 'uV', 'ylim': [ -60,  60], 'scalebar': 50, 'ylabel': 'BN3'}, #, 'decimation_factor': 400},\n",
    "    {'channel': 'Force',    'units': 'mN', 'ylim': [ -50, 450], 'scalebar': 300}, #, 'decimation_factor': 100},\n",
    "]\n",
    "\n",
    "fig_height_in_inches, bottom_fraction, top_fraction = solve_figure_vertical_dimensions(\n",
    "    len(plots), subplot_height_in_inches, bottom_margin_in_inches, top_margin_in_inches, 0)\n",
    "fig_width_in_inches, left_fraction, right_fraction = solve_figure_horizontal_dimensions(\n",
    "    1, inches_per_second*(t_stop-t_start).magnitude, left_margin_in_inches, right_margin_in_inches, 0)\n",
    "\n",
    "kwargs = dict(\n",
    "    formats = ['png'],\n",
    "    dpi = 600,\n",
    "    figsize = (fig_width_in_inches, fig_height_in_inches),\n",
    "    linewidth = 0.5,\n",
    "    x_scalebar = 10*pq.s,\n",
    "    layout_settings = dict(\n",
    "        left   = left_fraction,\n",
    "        right  = right_fraction,\n",
    "        bottom = bottom_fraction,\n",
    "        top    = top_fraction,\n",
    "        hspace = 0,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# load the metadata\n",
    "metadata = neurotic.MetadataSelector(metadata_file)\n",
    "metadata.select(data_set_name)\n",
    "\n",
    "# download data that is not already stored locally\n",
    "metadata.download_all_data_files()\n",
    "\n",
    "# load the data\n",
    "blk = neurotic.load_dataset(metadata, lazy=True)\n",
    "\n",
    "# manually perform filters\n",
    "blk = apply_filters(blk, metadata)\n",
    "\n",
    "# plot the signals\n",
    "fig, axes = plot_signals_with_scalebars(blk, t_start, t_stop, plots, **kwargs)\n",
    "\n",
    "# plot the inward movements\n",
    "y_inward = 1.03\n",
    "ep = next((ep for ep in blk.segments[0].epochs if ep.name=='Inward movement'), None)\n",
    "assert ep is not None\n",
    "ep = ep.time_slice(t_start, t_stop)\n",
    "for t, dur in zip(ep.times, ep.durations):\n",
    "    left, right = t, t + dur\n",
    "    bottom, top = y_inward-0.025, y_inward+0.025\n",
    "    width = right-left\n",
    "    height = top-bottom\n",
    "    rect = patches.Rectangle((left, bottom), width, height, linewidth=0, facecolor='k', fill=True, clip_on=False, transform=axes[0].get_xaxis_transform())\n",
    "    axes[0].add_patch(rect)\n",
    "\n",
    "# add bite/swallow markers\n",
    "behavior_markers = [\n",
    "    (2877.5, 'B'),\n",
    "]\n",
    "for t, label in behavior_markers:\n",
    "    axes[0].annotate(\n",
    "        label,\n",
    "        xy=(t, y_inward), xycoords=axes[0].get_xaxis_transform(),\n",
    "        ha='center', va='center', fontsize='small',\n",
    "    )\n",
    "\n",
    "# plot a zero baseline for force\n",
    "ax = axes[-1]\n",
    "force_zero = -11.5 # mN, average before animal began swallowing, not zero because of DC offset\n",
    "ax.axhline(force_zero, color='0.5', lw=1, ls='--', zorder=-1)\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-2B.png'), dpi=600)\n",
    "\n",
    "end = datetime.datetime.now()\n",
    "print('render time:', end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 2C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'Inward movement start to next inward movement start (s)'\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 3))\n",
    "plot_unloaded_vs_loaded(df_all.drop(unreliable_inward_movement), y, ax, bracket_width=2.9)\n",
    "ax.set_ylim([0, 10])\n",
    "ax.set_ylabel('Start of one inward\\nmovement to the next (s)')\n",
    "plt.subplots_adjust(left=0.32, right=0.99, top=0.91, bottom=0.10)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-2C.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 5))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(unreliable_inward_movement).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "y = 'Inward movement duration (s)'\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 3))\n",
    "plot_unloaded_vs_loaded(df_all.drop(unreliable_inward_movement), y, ax, bracket_width=2.9)\n",
    "ax.set_ylim([0, 10])\n",
    "ax.set_ylabel('Duration of\\ninward movement (s)')\n",
    "plt.subplots_adjust(left=0.32, right=0.99, top=0.91, bottom=0.10)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-2D.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 5))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(unreliable_inward_movement).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 2E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'Inward movement end to next inward movement start (s)'\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 3))\n",
    "plot_unloaded_vs_loaded(df_all.drop(unreliable_inward_movement), y, ax, bracket_width=2.9)\n",
    "ax.set_ylim([0, 10])\n",
    "ax.set_ylabel('Time between\\ninward movements (s)')\n",
    "plt.subplots_adjust(left=0.32, right=0.99, top=0.91, bottom=0.10)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-2E.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 5))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(unreliable_inward_movement).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [FIGURE 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "video_seg_phase_labels = [\n",
    "    '',\n",
    "    '',\n",
    "    '',\n",
    "    '',\n",
    "    'Inward movement',\n",
    "    '',\n",
    "    '',\n",
    "    '',\n",
    "    '',\n",
    "]\n",
    "\n",
    "def get_median_video_seg_phase_boundaries(df, show_summary=False):\n",
    "    \n",
    "    # the dataframe column 'Video segmentation times (s)' is an array of these times:\n",
    "    # - 0: inward_movement_start-inward_movement_duration*4,\n",
    "    # - 1: inward_movement_start-inward_movement_duration*3,\n",
    "    # - 2: inward_movement_start-inward_movement_duration*2,\n",
    "    # - 3: inward_movement_start-inward_movement_duration*1,\n",
    "    # - 4: inward_movement_start,\n",
    "    # - 5: inward_movement_end,\n",
    "    # - 6: inward_movement_end+inward_movement_duration*1,\n",
    "    # - 7: inward_movement_end+inward_movement_duration*2,\n",
    "    # - 8: inward_movement_end+inward_movement_duration*3,\n",
    "    # - 9: inward_movement_end+inward_movement_duration*4,\n",
    "\n",
    "    # get all normalization times from inward_movement_start to inward_movement_end\n",
    "    t = np.stack(df['Video segmentation times (s)']).magnitude[:, 4:5+1]\n",
    "\n",
    "    # get all phase durations\n",
    "    all_video_seg_phase_durations = np.diff(t).T\n",
    "\n",
    "    # find median video segmentation phase durations\n",
    "    # - 0: inward movement\n",
    "    median_video_seg_phase_durations = np.nanmedian(all_video_seg_phase_durations, axis=1)\n",
    "\n",
    "    # copy phase durations\n",
    "    # - 0: inward movement (duplicate)\n",
    "    # - 1: inward movement (duplicate)\n",
    "    # - 2: inward movement (duplicate)\n",
    "    # - 3: inward movement (duplicate)\n",
    "    # - 4: inward movement\n",
    "    # - 5: inward movement (duplicate)\n",
    "    # - 6: inward movement (duplicate)\n",
    "    # - 7: inward movement (duplicate)\n",
    "    # - 8: inward movement (duplicate)\n",
    "    median_video_seg_phase_durations = np.concatenate([\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "        median_video_seg_phase_durations,\n",
    "    ])\n",
    "\n",
    "    # convert median durations into median boundary timings\n",
    "    # - 0: inward_movement_start-inward_movement_duration*4 <-- reference point, zero by definition\n",
    "    # - 1: inward_movement_start-inward_movement_duration*3\n",
    "    # - 2: inward_movement_start-inward_movement_duration*2\n",
    "    # - 3: inward_movement_start-inward_movement_duration*1\n",
    "    # - 4: inward_movement_start\n",
    "    # - 5: inward_movement_end\n",
    "    # - 6: inward_movement_end+inward_movement_duration*1\n",
    "    # - 7: inward_movement_end+inward_movement_duration*2\n",
    "    # - 8: inward_movement_end+inward_movement_duration*3\n",
    "    # - 9: inward_movement_end+inward_movement_duration*4\n",
    "    median_video_seg_phase_boundaries = np.concatenate([[0], median_video_seg_phase_durations]).cumsum()\n",
    "    \n",
    "    if show_summary:\n",
    "        # plot the distributions of phase durations\n",
    "        plt.figure(figsize=(4,3))\n",
    "        plt.boxplot(\n",
    "            [a[np.isfinite(a)] for a in list(all_video_seg_phase_durations)],\n",
    "            labels=video_seg_phase_labels[4:5],\n",
    "            showmeans=True,\n",
    "        )\n",
    "        plt.ylim([0, None])\n",
    "        plt.ylabel('Phase duration (s)')\n",
    "        plt.title('Video segmentation')\n",
    "        plt.tight_layout()\n",
    "\n",
    "        # print summaries of the phase durations\n",
    "        for t, l in zip(all_video_seg_phase_durations, video_seg_phase_labels[4:5]):\n",
    "            l = l.replace('\\n', ' ')\n",
    "            print(f'{l}:\\tmedian {np.nanmedian(t):g}, mean {np.nanmean(t):g} (n={t[np.isfinite(t)].size})')\n",
    "\n",
    "    return median_video_seg_phase_boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# force_seg_phase_labels = [\n",
    "#     'IV\\nPrevious force\\nmaintenance',\n",
    "#     'V\\nPrevious major\\nforce drop',\n",
    "#     'I\\nPartial force\\nmaintenance',\n",
    "#     'II\\nForce dip',\n",
    "#     'III\\nForce rise',\n",
    "#     'IV\\nForce\\nmaintenance',\n",
    "#     'V\\nMajor\\nforce drop',\n",
    "#     'I\\nNext partial\\nforce maintenance',\n",
    "#     'II\\nNext\\nforce dip',\n",
    "# ]\n",
    "force_seg_phase_labels = [\n",
    "    '',\n",
    "    'V',\n",
    "    'I',\n",
    "    'II',\n",
    "    'III',\n",
    "    'IV',\n",
    "    'V',\n",
    "    '',\n",
    "    '',\n",
    "]\n",
    "\n",
    "def get_median_force_seg_phase_boundaries(df, show_summary=False):\n",
    "    \n",
    "    # the dataframe column 'Force segmentation times (s)' is an array of these times:\n",
    "    # - 0: prev_force_plateau_start\n",
    "    # - 1: prev_force_plateau_end\n",
    "    # - 2: prev_force_drop_end\n",
    "    # - 3: force_shoulder_end\n",
    "    # - 4: force_rise_start\n",
    "    # - 5: force_plateau_start\n",
    "    # - 6: force_plateau_end\n",
    "    # - 7: force_drop_end\n",
    "    # - 8: next_force_shoulder_end\n",
    "    # - 9: next_force_rise_start\n",
    "\n",
    "    # get all normalization times from previous_force_drop_end to current force_drop_end\n",
    "    t = np.stack(df['Force segmentation times (s)']).magnitude[:, 2:7+1]\n",
    "\n",
    "    # get all phase durations\n",
    "    all_force_seg_phase_durations = np.diff(t).T\n",
    "\n",
    "    # find median force segmentation phase durations\n",
    "    # - 0: partial force maintenance\n",
    "    # - 1: force dip\n",
    "    # - 2: force rise\n",
    "    # - 3: force maintenance\n",
    "    # - 4: major force drop\n",
    "    median_force_seg_phase_durations = np.nanmedian(all_force_seg_phase_durations, axis=1)\n",
    "\n",
    "    # copy phase durations to represent \"previous\" and \"next\" swallows\n",
    "    # - 0: previous force maintenance\n",
    "    # - 1: previous major force drop\n",
    "    # - 2: partial force maintenance\n",
    "    # - 3: force dip\n",
    "    # - 4: force rise\n",
    "    # - 5: force maintenance\n",
    "    # - 6: major force drop\n",
    "    # - 7: next partial force maintenance\n",
    "    # - 8: next force dip\n",
    "    median_force_seg_phase_durations = np.concatenate([\n",
    "        median_force_seg_phase_durations[-2:],\n",
    "        median_force_seg_phase_durations,\n",
    "        median_force_seg_phase_durations[:2],\n",
    "    ])\n",
    "\n",
    "    # convert median durations into median boundary timings\n",
    "    # - 0: prev_force_plateau_start <-- reference point, zero by definition\n",
    "    # - 1: prev_force_plateau_end\n",
    "    # - 2: prev_force_drop_end\n",
    "    # - 3: force_shoulder_end\n",
    "    # - 4: force_rise_start\n",
    "    # - 5: force_plateau_start\n",
    "    # - 6: force_plateau_end\n",
    "    # - 7: force_drop_end\n",
    "    # - 8: next_force_shoulder_end\n",
    "    # - 9: next_force_rise_start\n",
    "    median_force_seg_phase_boundaries = np.concatenate([[0], median_force_seg_phase_durations]).cumsum()\n",
    "    \n",
    "    if show_summary:\n",
    "        # plot the distributions of phase durations\n",
    "        plt.figure(figsize=(4,3))\n",
    "        plt.boxplot(\n",
    "            [a[np.isfinite(a)] for a in list(all_force_seg_phase_durations)],\n",
    "            labels=force_seg_phase_labels[2:7],\n",
    "            showmeans=True,\n",
    "        )\n",
    "        plt.ylim([0, None])\n",
    "        plt.ylabel('Phase duration (s)')\n",
    "        plt.title('Force segmentation')\n",
    "        plt.tight_layout()\n",
    "\n",
    "        # print summaries of the phase durations\n",
    "        for t, l in zip(all_force_seg_phase_durations, force_seg_phase_labels[2:7]):\n",
    "            l = l.replace('\\n', ' ')\n",
    "            print(f'{l}:\\tmedian {np.nanmedian(t):g}, mean {np.nanmean(t):g} (n={t[np.isfinite(t)].size})')\n",
    "\n",
    "    return median_force_seg_phase_boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "###\n",
    "### LOADED, FORCE SEGEMENTATION\n",
    "###\n",
    "\n",
    "median_loaded_force_seg_phase_boundaries = get_median_force_seg_phase_boundaries(df_all.query('Food == \"Tape nori\"'), show_summary=True)\n",
    "plt.title('Force segmentation of loaded swallows')\n",
    "\n",
    "###\n",
    "### UNLOADED, VIDEO SEGEMENTATION\n",
    "###\n",
    "\n",
    "median_unloaded_video_seg_phase_boundaries = get_median_video_seg_phase_boundaries(df_all.query('Food == \"Regular nori\"').drop(unreliable_inward_movement), show_summary=True)\n",
    "plt.title('Video segmentation of unloaded swallows')\n",
    "\n",
    "###\n",
    "### LOADED, VIDEO SEGEMENTATION\n",
    "###\n",
    "\n",
    "median_loaded_video_seg_phase_boundaries = get_median_video_seg_phase_boundaries(df_all.query('Food == \"Tape nori\"'), show_summary=True)\n",
    "plt.title('Video segmentation of loaded swallows')\n",
    "\n",
    "###\n",
    "### SHARED FIGURE SETTINGS\n",
    "###\n",
    "\n",
    "fig_width = 7 # inches\n",
    "fig_left_margin = 0.10 # axes fraction\n",
    "fig_right_margin = 0.90 # axes fraction\n",
    "t_width = 9.0 # sec\n",
    "time_before_force_rise = 4.1 # sec\n",
    "xlim_loaded_force_seg = [\n",
    "    (median_loaded_force_seg_phase_boundaries[4]-time_before_force_rise),\n",
    "    (median_loaded_force_seg_phase_boundaries[4]-time_before_force_rise) + t_width,\n",
    "]\n",
    "xlim_unloaded_video_seg = [\n",
    "    (median_unloaded_video_seg_phase_boundaries[4]-time_before_force_rise),\n",
    "    (median_unloaded_video_seg_phase_boundaries[4]-time_before_force_rise) + t_width,\n",
    "]\n",
    "xlim_loaded_video_seg = [\n",
    "    (median_loaded_video_seg_phase_boundaries[4]-time_before_force_rise),\n",
    "    (median_loaded_video_seg_phase_boundaries[4]-time_before_force_rise) + t_width,\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unloaded_b3_burst_start_times = df_all.query('Food == \"Regular nori\"').drop(unreliable_inward_movement)['B3 burst start (s)'].values\n",
    "print(f'{np.sum(np.isfinite(unloaded_b3_burst_start_times))} out of {unloaded_b3_burst_start_times.size} regular nori swallows have B3 bursts')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shoulder_end_times = np.stack(df_all.query('Food == \"Tape nori\"')['Force segmentation times (s)'].values).magnitude[:,3]\n",
    "print(f'{np.sum(np.isnan(shoulder_end_times))} out of {shoulder_end_times.size} tape nori swallows have missing shoulders')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 3A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Biomechanics schematic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 3B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(fig_width, 2.5))\n",
    "ax = plt.gca()\n",
    "\n",
    "# specify the data subset to use\n",
    "df = df_all.query('Food == \"Regular nori\"').drop(unreliable_inward_movement)\n",
    "\n",
    "# print number of swallows per animal\n",
    "print(df.groupby('Animal')['B6/B9 burst end (video seg normalized)'].count())\n",
    "print()\n",
    "\n",
    "# use a time plot range appropriate for the unnormalized burst timing\n",
    "xlim = xlim_unloaded_video_seg\n",
    "\n",
    "# specify boundaries where vertical lines will be drawn, and the labels between them\n",
    "phase_boundaries = median_unloaded_video_seg_phase_boundaries[4:6]\n",
    "phase_labels = video_seg_phase_labels[4:6]\n",
    "\n",
    "# specify times used for unnormalization\n",
    "unnormalization_fixed_times = median_unloaded_video_seg_phase_boundaries\n",
    "\n",
    "###\n",
    "### UNIT BOXES\n",
    "###\n",
    "\n",
    "for i, unit in enumerate(units):\n",
    "    t0_data    = df[f'{unit} burst start (video seg normalized)'].dropna()\n",
    "    t0_data[:] = unnormalize_time(unnormalization_fixed_times, t0_data.values)\n",
    "    t0_median  = t0_data.median()\n",
    "    t0_q1      = t0_data.quantile(0.25)\n",
    "    t0_q3      = t0_data.quantile(0.75)\n",
    "    \n",
    "    t1_data    = df[f'{unit} burst end (video seg normalized)'].dropna()\n",
    "    t1_data[:] = unnormalize_time(unnormalization_fixed_times, t1_data.values)\n",
    "    t1_median  = t1_data.median()\n",
    "    t1_q1      = t1_data.quantile(0.25)\n",
    "    t1_q3      = t1_data.quantile(0.75)\n",
    "    \n",
    "    # print summaries of the unit timing\n",
    "    print(f'{unit}:\\t[{t0_median:.2f} (n={t0_data.size}), {t1_median:.2f} (n={t1_data.size})], duration: {t1_median-t0_median:.2f}')\n",
    "    \n",
    "    if unit == 'B3':\n",
    "        # skip plotting B3 since it rarely burst for unloaded swallows\n",
    "        continue\n",
    "    \n",
    "    # plot boxes using medians\n",
    "    height = 0.8\n",
    "    lw = 1\n",
    "    rect = patches.Rectangle((t0_median, i-height/2), t1_median-t0_median, height, facecolor=unit_colors[unit], edgecolor='k', lw=lw, clip_on=False)\n",
    "    ax.add_patch(rect)\n",
    "    \n",
    "    # plot quartiles (25% and 75% quantiles)\n",
    "    ax.add_line(mlines.Line2D([t0_q1, t0_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t0_q3, t0_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t1_q1, t1_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t1_q3, t1_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t0_q1, t0_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t1_q1, t1_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "\n",
    "###\n",
    "### INWARD SEAWEED MOVEMENT\n",
    "###\n",
    "\n",
    "t0_data    = df['Inward movement start (video seg normalized)'].dropna()\n",
    "t0_data[:] = unnormalize_time(unnormalization_fixed_times, t0_data.values)\n",
    "t0_median  = t0_data.median()\n",
    "t0_q1      = t0_data.quantile(0.25)\n",
    "t0_q3      = t0_data.quantile(0.75)\n",
    "\n",
    "t1_data    = df['Inward movement end (video seg normalized)'].dropna()\n",
    "t1_data[:] = unnormalize_time(unnormalization_fixed_times, t1_data.values)\n",
    "t1_median  = t1_data.median()\n",
    "t1_q1      = t1_data.quantile(0.25)\n",
    "t1_q3      = t1_data.quantile(0.75)\n",
    "\n",
    "# print summaries of the unit timing\n",
    "print(f'Inward:\\t[{t0_median:.2f} (n={t0_data.size}), {t1_median:.2f} (n={t1_data.size})], duration: {t1_median-t0_median:.2f}')\n",
    "\n",
    "# plot box using medians\n",
    "i = len(units)\n",
    "height = 0.8\n",
    "lw = 1\n",
    "rect = patches.Rectangle((t0_median, i-height/2), t1_median-t0_median, height, facecolor='0.75', edgecolor='k', lw=lw, clip_on=False)\n",
    "ax.add_patch(rect)\n",
    "\n",
    "# # plot quartiles (25% and 75% quantiles)\n",
    "# ax.add_line(mlines.Line2D([t0_q1, t0_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t0_q3, t0_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t1_q1, t1_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t1_q3, t1_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t0_q1, t0_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t1_q1, t1_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "\n",
    "###\n",
    "### LABELS AND ANNOTATIONS\n",
    "###\n",
    "\n",
    "# plot vertical lines\n",
    "for t in phase_boundaries:\n",
    "    ax.axvline(x=t, ls=':', lw=1, c='gray', zorder=-1)\n",
    "\n",
    "# add phase labels on bottom edge\n",
    "ax.tick_params(bottom=False) # disable tick marks\n",
    "ax.set_xticks(phase_boundaries[:-1]+np.diff(phase_boundaries)/2)\n",
    "ax.set_xticklabels(phase_labels, fontsize='small')\n",
    "\n",
    "# add box labels on left edge\n",
    "ax.tick_params(left=False) # disable tick marks\n",
    "ax.set_yticks(range(len(units)+1))\n",
    "ax.set_yticklabels(units + ['Inward'], fontsize='medium')\n",
    "\n",
    "# add time scale bar\n",
    "add_scalebar(ax,\n",
    "    sizex=1, labelx='1 s',\n",
    "    loc='lower right', bbox_to_anchor=(1, 0),\n",
    "    borderpad=0.5, sep=5, barwidth=2,\n",
    ")\n",
    "\n",
    "###\n",
    "### FINISH\n",
    "###\n",
    "\n",
    "# set plot ranges\n",
    "ax.set_xlim(xlim)\n",
    "ax.set_ylim(6.75, -0.75)\n",
    "\n",
    "# remove box around figure\n",
    "sns.despine(bottom=True, left=True)\n",
    "\n",
    "plt.subplots_adjust(\n",
    "    left   = fig_left_margin,\n",
    "    right  = fig_right_margin,\n",
    "    bottom = 0.12,\n",
    "    top    = 1.00,\n",
    ")\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-3B.png'), dpi=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 3C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(fig_width, 2.5))\n",
    "ax = plt.gca()\n",
    "\n",
    "# specify the data subset to use\n",
    "df = df_all.query('Food == \"Tape nori\"')\n",
    "\n",
    "# print number of swallows per animal\n",
    "print(df.groupby('Animal')['B6/B9 burst end (force seg normalized)'].count())\n",
    "print()\n",
    "\n",
    "# use a time plot range appropriate for the unnormalized burst timing\n",
    "xlim = xlim_loaded_force_seg\n",
    "\n",
    "# specify boundaries where vertical lines will be drawn, and the labels between them\n",
    "phase_boundaries = median_loaded_force_seg_phase_boundaries[1:8]\n",
    "phase_labels = force_seg_phase_labels[1:8]\n",
    "\n",
    "# specify times used for unnormalization\n",
    "unnormalization_fixed_times = median_loaded_force_seg_phase_boundaries\n",
    "\n",
    "###\n",
    "### UNIT BOXES\n",
    "###\n",
    "\n",
    "for i, unit in enumerate(units):\n",
    "    t0_data    = df[f'{unit} burst start (force seg normalized)'].dropna()\n",
    "    t0_data[:] = unnormalize_time(unnormalization_fixed_times, t0_data.values)\n",
    "    t0_median  = t0_data.median()\n",
    "    t0_q1      = t0_data.quantile(0.25)\n",
    "    t0_q3      = t0_data.quantile(0.75)\n",
    "    \n",
    "    t1_data    = df[f'{unit} burst end (force seg normalized)'].dropna()\n",
    "    t1_data[:] = unnormalize_time(unnormalization_fixed_times, t1_data.values)\n",
    "    t1_median  = t1_data.median()\n",
    "    t1_q1      = t1_data.quantile(0.25)\n",
    "    t1_q3      = t1_data.quantile(0.75)\n",
    "    \n",
    "    # print summaries of the unit timing\n",
    "    print(f'{unit}:\\t[{t0_median:.2f} (n={t0_data.size}), {t1_median:.2f} (n={t1_data.size})], duration: {t1_median-t0_median:.2f}')\n",
    "    \n",
    "    # plot boxes using medians\n",
    "    height = 0.8\n",
    "    lw = 1\n",
    "    rect = patches.Rectangle((t0_median, i-height/2), t1_median-t0_median, height, facecolor=unit_colors[unit], edgecolor='k', lw=lw, clip_on=False)\n",
    "    ax.add_patch(rect)\n",
    "    \n",
    "    # plot quartiles (25% and 75% quantiles)\n",
    "    ax.add_line(mlines.Line2D([t0_q1, t0_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t0_q3, t0_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t1_q1, t1_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t1_q3, t1_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t0_q1, t0_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "    ax.add_line(mlines.Line2D([t1_q1, t1_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "\n",
    "###\n",
    "### INWARD SEAWEED MOVEMENT\n",
    "###\n",
    "\n",
    "t0_data    = df['Inward movement start (force seg normalized)'].dropna()\n",
    "t0_data[:] = unnormalize_time(unnormalization_fixed_times, t0_data.values)\n",
    "t0_median  = t0_data.median()\n",
    "t0_q1      = t0_data.quantile(0.25)\n",
    "t0_q3      = t0_data.quantile(0.75)\n",
    "\n",
    "t1_data    = df['Inward movement end (force seg normalized)'].dropna()\n",
    "t1_data[:] = unnormalize_time(unnormalization_fixed_times, t1_data.values)\n",
    "t1_median  = t1_data.median()\n",
    "t1_q1      = t1_data.quantile(0.25)\n",
    "t1_q3      = t1_data.quantile(0.75)\n",
    "\n",
    "# print summaries of the unit timing\n",
    "print(f'Inward:\\t[{t0_median:.2f} (n={t0_data.size}), {t1_median:.2f} (n={t1_data.size})], duration: {t1_median-t0_median:.2f}')\n",
    "\n",
    "# plot box using medians\n",
    "i = len(units)\n",
    "height = 0.8\n",
    "lw = 1\n",
    "rect = patches.Rectangle((t0_median, i-height/2), t1_median-t0_median, height, facecolor='0.75', edgecolor='k', lw=lw, clip_on=False)\n",
    "ax.add_patch(rect)\n",
    "\n",
    "# plot quartiles (25% and 75% quantiles)\n",
    "ax.add_line(mlines.Line2D([t0_q1, t0_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t0_q3, t0_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t1_q1, t1_q1], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t1_q3, t1_q3], [i-0.2, i+0.2], color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t0_q1, t0_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t1_q1, t1_q3], [i, i], color='k', lw=lw, clip_on=False))\n",
    "\n",
    "###\n",
    "### LABELS AND ANNOTATIONS\n",
    "###\n",
    "\n",
    "# plot vertical lines\n",
    "for t in phase_boundaries:\n",
    "    ax.axvline(x=t, ls=':', lw=1, c='gray', zorder=-1)\n",
    "\n",
    "# add phase labels on bottom edge\n",
    "ax.tick_params(bottom=False) # disable tick marks\n",
    "ax.set_xticks(phase_boundaries[:-1]+np.diff(phase_boundaries)/2)\n",
    "ax.set_xticklabels(phase_labels, fontsize='small')\n",
    "\n",
    "# add box labels on left edge\n",
    "ax.tick_params(left=False) # disable tick marks\n",
    "ax.set_yticks(range(len(units)+1))\n",
    "ax.set_yticklabels(units + ['Inward'], fontsize='medium')\n",
    "\n",
    "# add time scale bar\n",
    "add_scalebar(ax,\n",
    "    sizex=1, labelx='1 s',\n",
    "    loc='lower right', bbox_to_anchor=(1, 0),\n",
    "    borderpad=0.5, sep=5, barwidth=2,\n",
    ")\n",
    "\n",
    "###\n",
    "### FINISH\n",
    "###\n",
    "\n",
    "# set plot ranges\n",
    "ax.set_xlim(xlim)\n",
    "ax.set_ylim(6.75, -0.75)\n",
    "\n",
    "# remove box around figure\n",
    "sns.despine(bottom=True, left=True)\n",
    "\n",
    "plt.subplots_adjust(\n",
    "    left   = fig_left_margin,\n",
    "    right  = fig_right_margin,\n",
    "    bottom = 0.12,\n",
    "    top    = 1.00,\n",
    ")\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-3C.png'), dpi=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrows = len(units)\n",
    "subplot_height_in_inches = 412/600\n",
    "top_margin_in_inches = 300/600\n",
    "bottom_margin_in_inches = 180/600\n",
    "hspace = 0.2\n",
    "\n",
    "fig_height_in_inches, bottom_fraction, top_fraction = solve_figure_vertical_dimensions(\n",
    "    nrows, subplot_height_in_inches, bottom_margin_in_inches, top_margin_in_inches, hspace)\n",
    "\n",
    "fig, axes = plt.subplots(nrows, 1, sharex='col', figsize=(fig_width, fig_height_in_inches))\n",
    "\n",
    "# specify the data subset to use\n",
    "df = df_all.query('Food == \"Regular nori\"').drop(unreliable_inward_movement)\n",
    "\n",
    "# print number of swallows per animal\n",
    "print(df.groupby('Animal')['B6/B9 firing rate, video segmented interpolation (Hz)'].count())\n",
    "print()\n",
    "\n",
    "# use a time plot range appropriate for the unnormalized burst timing\n",
    "xlim = xlim_unloaded_video_seg\n",
    "\n",
    "# specify boundaries where vertical lines will be drawn, and the labels between them\n",
    "phase_boundaries = median_unloaded_video_seg_phase_boundaries[4:6]\n",
    "phase_labels = video_seg_phase_labels[4:6]\n",
    "\n",
    "# specify times used for unnormalization\n",
    "unnormalization_fixed_times = median_unloaded_video_seg_phase_boundaries\n",
    "\n",
    "# use these time values for plotting interpolated functions in unnormalized time\n",
    "interp_times_unnormalized = unnormalize_time(unnormalization_fixed_times, video_seg_interp_times)\n",
    "\n",
    "# specify dataframe columns to plot\n",
    "df_columns = {unit: f'{unit} firing rate, video segmented interpolation (Hz)' for unit in units}\n",
    "# df_columns['Force'] = 'Force, video segmented interpolation (mN)'\n",
    "\n",
    "###\n",
    "### UNIT FREQUENCIES AND FORCE\n",
    "###\n",
    "\n",
    "for i, (label, column) in enumerate(df_columns.items()):\n",
    "    ax = axes[i]\n",
    "\n",
    "    # find the median and quartiles\n",
    "    data   = np.stack(df[column])\n",
    "    median = np.nanmedian(data, axis=0)\n",
    "    q1     = np.nanquantile(data, q=0.25, axis=0)\n",
    "    q3     = np.nanquantile(data, q=0.75, axis=0)\n",
    "\n",
    "    # plot the median and quartiles (median last so it's on top)\n",
    "    ax.plot(interp_times_unnormalized, q1,     c=unit_colors[label], lw=1, ls='--', zorder=2)\n",
    "    ax.plot(interp_times_unnormalized, q3,     c=unit_colors[label], lw=1, ls='--', zorder=2)\n",
    "    ax.plot(interp_times_unnormalized, median, c=unit_colors[label], lw=2, zorder=2)\n",
    "\n",
    "###\n",
    "### INWARD SEAWEED MOVEMENT\n",
    "###\n",
    "\n",
    "t0_data    = df['Inward movement start (video seg normalized)'].dropna()\n",
    "t0_data[:] = unnormalize_time(unnormalization_fixed_times, t0_data.values)\n",
    "t0_median  = t0_data.median()\n",
    "t0_q1      = t0_data.quantile(0.25)\n",
    "t0_q3      = t0_data.quantile(0.75)\n",
    "\n",
    "t1_data    = df['Inward movement end (video seg normalized)'].dropna()\n",
    "t1_data[:] = unnormalize_time(unnormalization_fixed_times, t1_data.values)\n",
    "t1_median  = t1_data.median()\n",
    "t1_q1      = t1_data.quantile(0.25)\n",
    "t1_q3      = t1_data.quantile(0.75)\n",
    "\n",
    "ax = axes[0]\n",
    "y_inward = 1.3 # axes fractions above first panel\n",
    "\n",
    "# plot label\n",
    "ax.annotate(\n",
    "    'Inward',\n",
    "    xy=(0, y_inward), xycoords='axes fraction',\n",
    "    xytext=(-10, 0), textcoords='offset points',\n",
    "    ha='right', va='center',\n",
    ")\n",
    "\n",
    "# plot box using medians\n",
    "height = 0.2\n",
    "lw = 1\n",
    "rect = patches.Rectangle((t0_median, y_inward-height/2), t1_median-t0_median, height, transform=ax.get_xaxis_transform(), facecolor='0.75', edgecolor='k', lw=lw, clip_on=False)\n",
    "ax.add_patch(rect)\n",
    "\n",
    "# # plot quartiles (25% and 75% quantiles)\n",
    "# whisker_height = 0.1\n",
    "# ax.add_line(mlines.Line2D([t0_q1, t0_q1], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t0_q3, t0_q3], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t1_q1, t1_q1], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t1_q3, t1_q3], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t0_q1, t0_q3], [y_inward, y_inward], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "# ax.add_line(mlines.Line2D([t1_q1, t1_q3], [y_inward, y_inward], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "\n",
    "###\n",
    "### LABELS AND ANNOTATIONS\n",
    "###\n",
    "\n",
    "ax = axes[-1]\n",
    "\n",
    "# plot vertical lines\n",
    "for t in phase_boundaries:\n",
    "    ymax = nrows + (nrows-1)*hspace + (y_inward-1)*2\n",
    "    ax.axvline(x=t, ymax=ymax, ls=':', lw=1, c='gray', zorder=-1, clip_on=False) # axvline's dotted line style is better than ConnectionPath's\n",
    "for i in range(len(df_columns)-1):\n",
    "    axes[i].set_zorder(1) # elevate axes so vertical lines are behind it\n",
    "    axes[i].set_facecolor('none') # remove background so vertical line is visible\n",
    "\n",
    "# plot horizontal grid lines\n",
    "for i in range(len(df_columns)):\n",
    "    axes[i].grid(axis='y', clip_on=False)\n",
    "\n",
    "# add phase labels on bottom edge\n",
    "for i in range(len(df_columns)):\n",
    "    axes[i].tick_params(bottom=False) # disable tick marks\n",
    "ax.set_xticks(phase_boundaries[:-1]+np.diff(phase_boundaries)/2)\n",
    "ax.set_xticklabels(phase_labels, fontsize='small')\n",
    "\n",
    "# add trace labels on left edge\n",
    "for i, label in enumerate(df_columns):\n",
    "    axes[i].set_ylabel(label, rotation='horizontal', ha='right', va='center', labelpad=10, fontsize='medium')\n",
    "\n",
    "# add yticks on right edge\n",
    "yticks = {\n",
    "    'B38':   [0,  10,  20],\n",
    "    'I2':    [0,  10,  20],\n",
    "    'B8a/b': [0,  20,  40],\n",
    "    'B6/B9': [0,  25,  50],\n",
    "    'B3':    [0,   5,  10],\n",
    "    'B4/B5': [0,  10,  20],\n",
    "    'Force': [0, 150, 300],\n",
    "}\n",
    "for i, (label, column) in enumerate(df_columns.items()):\n",
    "    axes[i].tick_params(right=True, labelright=True) # enable tick marks and labels\n",
    "    axes[i].set_ylim([min(yticks[label]), max(yticks[label])])\n",
    "    axes[i].set_yticks(yticks[label])\n",
    "    yticklabels = [f'{y:g}' for y in axes[i].get_yticks()]\n",
    "    yunits = column.split()[-1].strip('()')\n",
    "    yticklabels[-1] += f' {yunits}' # append units to max label\n",
    "    axes[i].set_yticklabels(yticklabels)\n",
    "\n",
    "# add time scale bar\n",
    "add_scalebar(ax,\n",
    "    sizex=1, labelx='1 s',\n",
    "    loc='upper right', bbox_to_anchor=(1, 1),\n",
    "    borderpad=0.5, sep=5, barwidth=2,\n",
    ")\n",
    "\n",
    "###\n",
    "### FINISH\n",
    "###\n",
    "\n",
    "# set time range\n",
    "ax.set_xlim(xlim)\n",
    "\n",
    "# remove left and top plot borders\n",
    "sns.despine(fig=fig, left=True, right=False)\n",
    "\n",
    "plt.subplots_adjust(\n",
    "    left   = fig_left_margin,\n",
    "    right  = fig_right_margin,\n",
    "    bottom = bottom_fraction,\n",
    "    top    = top_fraction,\n",
    "    hspace = hspace,\n",
    ")\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-3D.png'), dpi=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 3E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "nrows = len(units)+1\n",
    "subplot_height_in_inches = 412/600\n",
    "top_margin_in_inches = 300/600\n",
    "bottom_margin_in_inches = 180/600\n",
    "hspace = 0.2\n",
    "\n",
    "fig_height_in_inches, bottom_fraction, top_fraction = solve_figure_vertical_dimensions(\n",
    "    nrows, subplot_height_in_inches, bottom_margin_in_inches, top_margin_in_inches, hspace)\n",
    "\n",
    "fig, axes = plt.subplots(nrows, 1, sharex='col', figsize=(fig_width, fig_height_in_inches))\n",
    "\n",
    "# specify the data subset to use\n",
    "df = df_all.query('Food == \"Tape nori\"')\n",
    "\n",
    "# print number of swallows per animal\n",
    "print(df.groupby('Animal')['B6/B9 firing rate, force segmented interpolation (Hz)'].count())\n",
    "print()\n",
    "\n",
    "# use a time plot range appropriate for the unnormalized burst timing\n",
    "xlim = xlim_loaded_force_seg\n",
    "\n",
    "# specify boundaries where vertical lines will be drawn, and the labels between them\n",
    "phase_boundaries = median_loaded_force_seg_phase_boundaries[1:8]\n",
    "phase_labels = force_seg_phase_labels[1:8]\n",
    "\n",
    "# specify times used for unnormalization\n",
    "unnormalization_fixed_times = median_loaded_force_seg_phase_boundaries\n",
    "\n",
    "# use these time values for plotting interpolated functions in unnormalized time\n",
    "interp_times_unnormalized = unnormalize_time(unnormalization_fixed_times, force_seg_interp_times)\n",
    "\n",
    "# specify dataframe columns to plot\n",
    "df_columns = {unit: f'{unit} firing rate, force segmented interpolation (Hz)' for unit in units}\n",
    "df_columns['Force'] = 'Force, force segmented interpolation (mN)'\n",
    "\n",
    "###\n",
    "### UNIT FREQUENCIES AND FORCE\n",
    "###\n",
    "\n",
    "for i, (label, column) in enumerate(df_columns.items()):\n",
    "    ax = axes[i]\n",
    "\n",
    "    # find the median and quartiles\n",
    "    data   = np.stack(df[column])\n",
    "    median = np.nanmedian(data, axis=0)\n",
    "    q1     = np.nanquantile(data, q=0.25, axis=0)\n",
    "    q3     = np.nanquantile(data, q=0.75, axis=0)\n",
    "\n",
    "    # plot the median and quartiles (median last so it's on top)\n",
    "    ax.plot(interp_times_unnormalized, q1,     c=unit_colors[label], lw=1, ls='--', zorder=2)\n",
    "    ax.plot(interp_times_unnormalized, q3,     c=unit_colors[label], lw=1, ls='--', zorder=2)\n",
    "    ax.plot(interp_times_unnormalized, median, c=unit_colors[label], lw=2, zorder=2)\n",
    "\n",
    "###\n",
    "### INWARD SEAWEED MOVEMENT\n",
    "###\n",
    "\n",
    "t0_data    = df['Inward movement start (force seg normalized)'].dropna()\n",
    "t0_data[:] = unnormalize_time(unnormalization_fixed_times, t0_data.values)\n",
    "t0_median  = t0_data.median()\n",
    "t0_q1      = t0_data.quantile(0.25)\n",
    "t0_q3      = t0_data.quantile(0.75)\n",
    "\n",
    "t1_data    = df['Inward movement end (force seg normalized)'].dropna()\n",
    "t1_data[:] = unnormalize_time(unnormalization_fixed_times, t1_data.values)\n",
    "t1_median  = t1_data.median()\n",
    "t1_q1      = t1_data.quantile(0.25)\n",
    "t1_q3      = t1_data.quantile(0.75)\n",
    "\n",
    "ax = axes[0]\n",
    "y_inward = 1.3 # axes fractions above first panel\n",
    "\n",
    "# plot label\n",
    "ax.annotate(\n",
    "    'Inward',\n",
    "    xy=(0, y_inward), xycoords='axes fraction',\n",
    "    xytext=(-10, 0), textcoords='offset points',\n",
    "    ha='right', va='center',\n",
    ")\n",
    "\n",
    "# plot box using medians\n",
    "height = 0.2\n",
    "lw = 1\n",
    "rect = patches.Rectangle((t0_median, y_inward-height/2), t1_median-t0_median, height, transform=ax.get_xaxis_transform(), facecolor='0.75', edgecolor='k', lw=lw, clip_on=False)\n",
    "ax.add_patch(rect)\n",
    "\n",
    "# plot quartiles (25% and 75% quantiles)\n",
    "whisker_height = 0.1\n",
    "ax.add_line(mlines.Line2D([t0_q1, t0_q1], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t0_q3, t0_q3], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t1_q1, t1_q1], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t1_q3, t1_q3], [y_inward-whisker_height/2, y_inward+whisker_height/2], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t0_q1, t0_q3], [y_inward, y_inward], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "ax.add_line(mlines.Line2D([t1_q1, t1_q3], [y_inward, y_inward], transform=ax.get_xaxis_transform(), color='k', lw=lw, clip_on=False))\n",
    "\n",
    "###\n",
    "### LABELS AND ANNOTATIONS\n",
    "###\n",
    "\n",
    "ax = axes[-1]\n",
    "\n",
    "# plot vertical lines\n",
    "for t in phase_boundaries:\n",
    "    ymax = nrows + (nrows-1)*hspace + (y_inward-1)*2\n",
    "    ax.axvline(x=t, ymax=ymax, ls=':', lw=1, c='gray', zorder=-1, clip_on=False) # axvline's dotted line style is better than ConnectionPath's\n",
    "for i in range(len(df_columns)-1):\n",
    "    axes[i].set_zorder(1) # elevate axes so vertical lines are behind it\n",
    "    axes[i].set_facecolor('none') # remove background so vertical line is visible\n",
    "\n",
    "# plot horizontal grid lines\n",
    "for i in range(len(df_columns)):\n",
    "    axes[i].grid(axis='y', clip_on=False)\n",
    "\n",
    "# add phase labels on bottom edge\n",
    "for i in range(len(df_columns)):\n",
    "    axes[i].tick_params(bottom=False) # disable tick marks\n",
    "ax.set_xticks(phase_boundaries[:-1]+np.diff(phase_boundaries)/2)\n",
    "ax.set_xticklabels(phase_labels, fontsize='small')\n",
    "\n",
    "# add trace labels on left edge\n",
    "for i, label in enumerate(df_columns):\n",
    "    axes[i].set_ylabel(label, rotation='horizontal', ha='right', va='center', labelpad=10, fontsize='medium')\n",
    "\n",
    "# add yticks on right edge\n",
    "yticks = {\n",
    "    'B38':   [0,  10,  20],\n",
    "    'I2':    [0,  10,  20],\n",
    "    'B8a/b': [0,  20,  40],\n",
    "    'B6/B9': [0,  25,  50],\n",
    "    'B3':    [0,   5,  10],\n",
    "    'B4/B5': [0,  10,  20],\n",
    "    'Force': [0, 150, 300],\n",
    "}\n",
    "for i, (label, column) in enumerate(df_columns.items()):\n",
    "    axes[i].tick_params(right=True, labelright=True) # enable tick marks and labels\n",
    "    axes[i].set_ylim([min(yticks[label]), max(yticks[label])])\n",
    "    axes[i].set_yticks(yticks[label])\n",
    "    yticklabels = [f'{y:g}' for y in axes[i].get_yticks()]\n",
    "    yunits = column.split()[-1].strip('()')\n",
    "    yticklabels[-1] += f' {yunits}' # append units to max label\n",
    "    axes[i].set_yticklabels(yticklabels)\n",
    "\n",
    "# add time scale bar\n",
    "add_scalebar(ax,\n",
    "    sizex=1, labelx='1 s',\n",
    "    loc='upper right', bbox_to_anchor=(1, 1),\n",
    "    borderpad=0.5, sep=5, barwidth=2,\n",
    ")\n",
    "\n",
    "###\n",
    "### FINISH\n",
    "###\n",
    "\n",
    "# set time range\n",
    "ax.set_xlim(xlim)\n",
    "\n",
    "# remove left and top plot borders\n",
    "sns.despine(fig=fig, left=True, right=False)\n",
    "\n",
    "plt.subplots_adjust(\n",
    "    left   = fig_left_margin,\n",
    "    right  = fig_right_margin,\n",
    "    bottom = bottom_fraction,\n",
    "    top    = top_fraction,\n",
    "    hspace = hspace,\n",
    ")\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-3E.png'), dpi=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [FIGURE 4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 4 Statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Burst duration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Protraction-phase Hotelling's T-squared"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['B38 burst duration (s)', 'I2 burst duration (s)']\n",
    "\n",
    "df = df_all.drop(bite_swallow_behaviors)\n",
    "x = df.query('Food == \"Regular nori\"')[columns].groupby('Animal').mean().values\n",
    "y = df.query('Food == \"Tape nori\"')[columns].groupby('Animal').mean().values\n",
    "result = r_stats.T2_test(x-y)\n",
    "print(f\"T^2 = {result['T2']:.3f}, F = {result['F']:.3f}, df_num = {result['df_num']}, df_den = {result['df_den']}, p = {result['p']:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"A paired-samples Hotelling’s T-squared test indicated no significant difference in durations of mean protraction-phase motor activity (mean B38 and mean I2 burst durations) in the five animals between loaded and unloaded swallows (T^2 = XXX, F = XXX, df_num = 2, df_den = 3, p = XXX).\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Retraction-phase Hotelling's T-squared"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['B8a/b burst duration (s)', 'B3/B6/B9 burst duration (s)']\n",
    "\n",
    "df = df_all.drop(bite_swallow_behaviors)\n",
    "x = df.query('Food == \"Regular nori\"')[columns].groupby('Animal').mean().values\n",
    "y = df.query('Food == \"Tape nori\"')[columns].groupby('Animal').mean().values\n",
    "result = r_stats.T2_test(x-y)\n",
    "print(f\"T^2 = {result['T2']:.3f}, F = {result['F']:.3f}, df_num = {result['df_num']}, df_den = {result['df_den']}, p = {result['p']:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"In contrast, a paired-samples Hotelling’s T-squared test indicated that durations of mean retraction-phase motor activity (mean B8a/b and mean B3/B6/B9 burst durations) in the five animals were significantly different between loaded and unloaded swallows (T^2 = XXX, F = XXX, df_num = 2, df_den = 3, p = XXX).\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### All t-tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df_all.drop(bite_swallow_behaviors)\n",
    "\n",
    "columns = [\n",
    "#     'B38 burst duration (s)',\n",
    "#     'I2 burst duration (s)',\n",
    "    'B8a/b burst duration (s)',\n",
    "    'B3/B6/B9 burst duration (s)',\n",
    "    'B4/B5 burst duration (s)',\n",
    "]\n",
    "\n",
    "print('NOTE: This table assumes by using t-tests that all Shapiro-Wilk tests conducted below were not significant!')\n",
    "print()\n",
    "print('\\t\\tUnloaded\\t\\t\\tLoaded\\t\\t\\t\\tPaired t-test\\tEffect size')\n",
    "print('Unit\\t\\tMean\\tSEM\\tSTD\\tCount\\tMean\\tSEM\\tSTD\\tCount\\tt\\tp\\td')\n",
    "for column in columns:\n",
    "    x = df.query('Food == \"Tape nori\"')[column].groupby('Animal').mean()\n",
    "    y = df.query('Food == \"Regular nori\"')[column].groupby('Animal').mean()\n",
    "\n",
    "    ttest_result = r_stats.t_test(x.values, y.values, paired=True, alternative='greater')\n",
    "    ttest_t, ttest_p = ttest_result['t'], ttest_result['p']\n",
    "    \n",
    "    cohen_d = r_stats.cohen_d(x.values, y.values)['estimate']\n",
    "\n",
    "    print(f'{column.split()[0].ljust(8)}\\t{y.mean():.2f}\\t{y.sem():.2f}\\t{y.std():.2f}\\t{y.count()}\\t{x.mean():.2f}\\t{x.sem():.2f}\\t{x.std():.2f}\\t{x.count()}\\t{ttest_t:.3f}\\t{ttest_p:.3f}\\t{cohen_d:.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"Post hoc paired-samples one-tailed t-tests indicated that both mean B8a/b burst duration (t = XXX, df = 4, p = XXX, Cohen's d = XXX) and mean B3/B6/B9 burst duration (t = XXX, df = 4, p = XXX, Cohen's d = XXX) were significantly longer for loaded swallows.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"Separately, a paired-samples one-tailed t-test indicated that the mean burst duration of multi-action neurons B4/B5 was also significantly longer for loaded swallows (t = XXX, df = 4, p = XXX, Cohen's d = XXX).\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Burst mean frequency"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Protraction-phase Hotelling's T-squared"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['B38 burst mean frequency (Hz)', 'I2 burst mean frequency (Hz)']\n",
    "\n",
    "df = df_all.drop(bite_swallow_behaviors)\n",
    "x = df.query('Food == \"Regular nori\"')[columns].groupby('Animal').mean().values\n",
    "y = df.query('Food == \"Tape nori\"')[columns].groupby('Animal').mean().values\n",
    "result = r_stats.T2_test(x-y)\n",
    "print(f\"T^2 = {result['T2']:.3f}, F = {result['F']:.3f}, df_num = {result['df_num']}, df_den = {result['df_den']}, p = {result['p']:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"When a similar analysis was applied to the mean firing rates of the protraction-phase motor activity (mean B38 and mean I2 burst firing frequencies), no significant difference was found between loaded and unloaded swallows (T^2 = XXX, F = XXX, df_num = 2, df_den = 3, p = XXX).\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Retraction-phase motor units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['B8a/b burst mean frequency (Hz)', 'B3/B6/B9 burst mean frequency (Hz)']\n",
    "\n",
    "df = df_all.drop(bite_swallow_behaviors)\n",
    "x = df.query('Food == \"Regular nori\"')[columns].groupby('Animal').mean().values\n",
    "y = df.query('Food == \"Tape nori\"')[columns].groupby('Animal').mean().values\n",
    "result = r_stats.T2_test(x-y)\n",
    "print(f\"T^2 = {result['T2']:.3f}, F = {result['F']:.3f}, df_num = {result['df_num']}, df_den = {result['df_den']}, p = {result['p']:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"In contrast, the mean firing rates of the retraction-phase motor activity (mean B8a/b and mean B3/B6/B9 burst firing frequencies) were significantly different between loaded and unloaded swallows (T^2 = XXX, F = XXX, df_num = 2, df_den = 3, p = XXX).\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Post hoc t-tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df_all.drop(bite_swallow_behaviors)\n",
    "\n",
    "columns = [\n",
    "#     'B38 burst mean frequency (Hz)',\n",
    "#     'I2 burst mean frequency (Hz)',\n",
    "    'B8a/b burst mean frequency (Hz)',\n",
    "    'B3/B6/B9 burst mean frequency (Hz)',\n",
    "#     'B4/B5 burst mean frequency (Hz)',\n",
    "]\n",
    "\n",
    "print('NOTE: This table assumes by using t-tests that all Shapiro-Wilk tests conducted below were not significant!')\n",
    "print()\n",
    "print('\\t\\tUnloaded\\t\\t\\tLoaded\\t\\t\\t\\tPaired t-test\\tEffect size')\n",
    "print('Unit\\t\\tMean\\tSEM\\tSTD\\tCount\\tMean\\tSEM\\tSTD\\tCount\\tt\\tp\\td')\n",
    "for column in columns:\n",
    "    x = df.query('Food == \"Tape nori\"')[column].groupby('Animal').mean()\n",
    "    y = df.query('Food == \"Regular nori\"')[column].groupby('Animal').mean()\n",
    "\n",
    "    ttest_result = r_stats.t_test(x.values, y.values, paired=True, alternative='greater')\n",
    "    ttest_t, ttest_p = ttest_result['t'], ttest_result['p']\n",
    "    \n",
    "    cohen_d = r_stats.cohen_d(x.values, y.values)['estimate']\n",
    "\n",
    "    print(f'{column.split()[0].ljust(8)}\\t{y.mean():.2f}\\t{y.sem():.2f}\\t{y.std():.2f}\\t{y.count()}\\t{x.mean():.2f}\\t{x.sem():.2f}\\t{x.std():.2f}\\t{x.count()}\\t{ttest_t:.3f}\\t{ttest_p:.3f}\\t{cohen_d:.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"Post hoc paired-samples one-tailed t-tests indicated that mean B3/B6/B9 firing rate (t = XXX, df = 4, p = XXX, Cohen's d = XXX) but not B8a/b firing rate (t = XXX, df = 4, p = XXX, Cohen's d = XXX) was significantly greater for loaded swallows.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### B4/B5 Wilcoxon signed-rank test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "column = 'B4/B5 burst mean frequency (Hz)'\n",
    "\n",
    "df = df_all.drop(bite_swallow_behaviors)\n",
    "df = df.reset_index()\n",
    "x = df.query('Food == \"Tape nori\"').groupby('Animal')[column].mean().values\n",
    "y = df.query('Food == \"Regular nori\"').groupby('Animal')[column].mean().values\n",
    "\n",
    "result = r_stats.shapiro_test(x-y)\n",
    "print(f\"Shapiro-Wilk: W = {result['W']:g}, p = {result['p']:.4f}\")\n",
    "\n",
    "result = r_stats.wilcox_test(x, y, paired=True, alternative='greater')\n",
    "print(f\"Wilcoxon: W = {result['W']:g}, p = {result['p']:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"Finally, a separate paired-samples one-tailed Wilcoxon signed-rank test (conducted because the normality assumption of the t-test was not satisfied) indicated that the mean firing rate of multi-action neurons B4/B5 was not significantly greater for loaded swallows (W = XXX, p = XXX).\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 4A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'B38 burst duration (s)'\n",
    "color = unit_colors['B38']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 5])\n",
    "ax.set_ylabel(y)\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-4A.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 4B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'I2 burst duration (s)'\n",
    "color = unit_colors['I2']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 5])\n",
    "ax.set_ylabel(y)\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-4B.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 4C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "y = 'B8a/b burst duration (s)'\n",
    "color = unit_colors['B8a/b']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 5])\n",
    "ax.set_ylabel(y)\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-4C.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 4D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'B3/B6/B9 burst duration (s)'\n",
    "color = unit_colors['B3/B6/B9']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 5])\n",
    "ax.set_ylabel(y)\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-4D.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 4E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'B4/B5 burst duration (s)'\n",
    "color = unit_colors['B4/B5']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 5])\n",
    "ax.set_ylabel(y)\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-4E.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 4F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'B3/B6/B9 burst mean frequency (Hz)'\n",
    "color = unit_colors['B3/B6/B9']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 50])\n",
    "ax.set_ylabel('B3/B6/B9 firing rate (Hz)')\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "fig.savefig(os.path.join(export_dir, 'figure-4F.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Other frequencies (not plotted in manuscript)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'B38 burst mean frequency (Hz)'\n",
    "color = unit_colors['B38']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 50])\n",
    "ax.set_ylabel('B38 firing rate (Hz)')\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "# fig.savefig(os.path.join(export_dir, 'figure-4F.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'I2 burst mean frequency (Hz)'\n",
    "color = unit_colors['I2']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 50])\n",
    "ax.set_ylabel('I2 firing rate (Hz)')\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "# fig.savefig(os.path.join(export_dir, 'figure-4F.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "y = 'B8a/b burst mean frequency (Hz)'\n",
    "color = unit_colors['B8a/b']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 50])\n",
    "ax.set_ylabel('B8a/b firing rate (Hz)')\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "# fig.savefig(os.path.join(export_dir, 'figure-4F.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 'B4/B5 burst mean frequency (Hz)'\n",
    "color = unit_colors['B4/B5']\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(2.25, 2.5))\n",
    "plot_unloaded_vs_loaded(df_all.drop(bite_swallow_behaviors), y, ax, color, bracket_width=3.2)\n",
    "ax.set_ylim([0, 50])\n",
    "ax.set_ylabel('B4/B5 firing rate (Hz)')\n",
    "ax.yaxis.set_label_coords(-0.2, 0.5)\n",
    "plt.subplots_adjust(left=0.25, right=0.99, top=0.89, bottom=0.11)\n",
    "# fig.savefig(os.path.join(export_dir, 'figure-4F.png'), dpi=300)\n",
    "\n",
    "# # plot by animal\n",
    "# fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "# boxplot_with_points('Animal', y, 'Food', df_all.drop(bite_swallow_behaviors).reset_index(), ax, show_points=True)\n",
    "# ax.set_ylim([0, None])\n",
    "# plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [FIGURE 5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 5A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(3, 3))\n",
    "\n",
    "data_subsets = [\n",
    "    'JG07 Tape nori',\n",
    "    'JG08 Tape nori',\n",
    "    'JG11 Tape nori',\n",
    "    'JG12 Tape nori',\n",
    "    'JG14 Tape nori',\n",
    "]\n",
    "data_subsets = {label:label2query(label) for label in data_subsets}\n",
    "\n",
    "xlabel, xlim = 'B3/B6/B9 burst duration (s)', [0, 5]\n",
    "ylabel, ylabel_alt, ylim = 'Force plateau duration (s)', 'Force maintenance duration (s)', [0, 5]\n",
    "\n",
    "scatter2d(ax, data_subsets, xlabel, xlim, ylabel, ylim, trend=True, tooltips=False, colors=['k']*5, markers=['.']*5)\n",
    "plt.plot([0, 999], [0, 999], ls=':', c='gray', zorder=0)\n",
    "plt.ylabel(ylabel_alt)\n",
    "sns.despine(ax=ax, trim=True)\n",
    "plt.tight_layout()\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-5A.png'), dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 5B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(3, 3))\n",
    "\n",
    "data_subsets = [\n",
    "    'JG07 Tape nori',\n",
    "]\n",
    "data_subsets = {label:label2query(label) for label in data_subsets}\n",
    "\n",
    "xlabel, xlim = 'B3/B6/B9 burst duration (s)', [0, 5]\n",
    "ylabel, ylabel_alt, ylim = 'Force plateau duration (s)', 'Force maintenance duration (s)', [0, 5]\n",
    "\n",
    "scatter2d(ax, data_subsets, xlabel, xlim, ylabel, ylim, trend_separately=True, tooltips=False, colors=['k'])\n",
    "plt.plot([0, 999], [0, 999], ls=':', c='gray', zorder=0)\n",
    "plt.ylabel(ylabel_alt)\n",
    "sns.despine(ax=ax, trim=True)\n",
    "plt.tight_layout()\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-5B.png'), dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 5C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(3, 3))\n",
    "\n",
    "data_subsets = [\n",
    "    'JG08 Tape nori',\n",
    "]\n",
    "data_subsets = {label:label2query(label) for label in data_subsets}\n",
    "\n",
    "xlabel, xlim = 'B3/B6/B9 burst duration (s)', [0, 5]\n",
    "ylabel, ylabel_alt, ylim = 'Force plateau duration (s)', 'Force maintenance duration (s)', [0, 5]\n",
    "\n",
    "scatter2d(ax, data_subsets, xlabel, xlim, ylabel, ylim, trend_separately=True, tooltips=False, colors=['k'])\n",
    "plt.plot([0, 999], [0, 999], ls=':', c='gray', zorder=0)\n",
    "plt.ylabel(ylabel_alt)\n",
    "sns.despine(ax=ax, trim=True)\n",
    "plt.tight_layout()\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-5C.png'), dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 5D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(3, 3))\n",
    "\n",
    "data_subsets = [\n",
    "    'JG11 Tape nori',\n",
    "]\n",
    "data_subsets = {label:label2query(label) for label in data_subsets}\n",
    "\n",
    "xlabel, xlim = 'B3/B6/B9 burst duration (s)', [0, 5]\n",
    "ylabel, ylabel_alt, ylim = 'Force plateau duration (s)', 'Force maintenance duration (s)', [0, 5]\n",
    "\n",
    "scatter2d(ax, data_subsets, xlabel, xlim, ylabel, ylim, trend_separately=True, tooltips=False, colors=['k'])\n",
    "plt.plot([0, 999], [0, 999], ls=':', c='gray', zorder=0)\n",
    "plt.ylabel(ylabel_alt)\n",
    "sns.despine(ax=ax, trim=True)\n",
    "plt.tight_layout()\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-5D.png'), dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 5E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(3, 3))\n",
    "\n",
    "data_subsets = [\n",
    "    'JG12 Tape nori',\n",
    "]\n",
    "data_subsets = {label:label2query(label) for label in data_subsets}\n",
    "\n",
    "xlabel, xlim = 'B3/B6/B9 burst duration (s)', [0, 5]\n",
    "ylabel, ylabel_alt, ylim = 'Force plateau duration (s)', 'Force maintenance duration (s)', [0, 5]\n",
    "\n",
    "scatter2d(ax, data_subsets, xlabel, xlim, ylabel, ylim, trend_separately=True, tooltips=False, colors=['k'])\n",
    "plt.plot([0, 999], [0, 999], ls=':', c='gray', zorder=0)\n",
    "plt.ylabel(ylabel_alt)\n",
    "sns.despine(ax=ax, trim=True)\n",
    "plt.tight_layout()\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-5E.png'), dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐌 Figure 5F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(3, 3))\n",
    "\n",
    "data_subsets = [\n",
    "    'JG14 Tape nori',\n",
    "]\n",
    "data_subsets = {label:label2query(label) for label in data_subsets}\n",
    "\n",
    "xlabel, xlim = 'B3/B6/B9 burst duration (s)', [0, 5]\n",
    "ylabel, ylabel_alt, ylim = 'Force plateau duration (s)', 'Force maintenance duration (s)', [0, 5]\n",
    "\n",
    "scatter2d(ax, data_subsets, xlabel, xlim, ylabel, ylim, trend_separately=True, tooltips=False, colors=['k'])\n",
    "plt.plot([0, 999], [0, 999], ls=':', c='gray', zorder=0)\n",
    "plt.ylabel(ylabel_alt)\n",
    "sns.despine(ax=ax, trim=True)\n",
    "plt.tight_layout()\n",
    "\n",
    "fig.savefig(os.path.join(export_dir, 'figure-5F.png'), dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Statistical Tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stats_table(figure_map, df, filename=None):\n",
    "\n",
    "    df = df.reset_index()\n",
    "    df.loc[df['Food'] == 'Regular nori', 'Food'] = 'Unloaded'\n",
    "    df.loc[df['Food'] == 'Tape nori',    'Food'] = 'Loaded'\n",
    "    df.loc[df['Animal'] == 'JG07', 'Animal'] = 1\n",
    "    df.loc[df['Animal'] == 'JG08', 'Animal'] = 2\n",
    "    df.loc[df['Animal'] == 'JG11', 'Animal'] = 3\n",
    "    df.loc[df['Animal'] == 'JG12', 'Animal'] = 4\n",
    "    df.loc[df['Animal'] == 'JG14', 'Animal'] = 5\n",
    "    \n",
    "\n",
    "    if filename is None:\n",
    "        file = sys.stdout\n",
    "    else:\n",
    "        file = open(filename, 'w', encoding='utf8')\n",
    "        \n",
    "    for fig, columns in figure_map.items():\n",
    "\n",
    "        x = df.query('Food == \"Loaded\"').groupby('Animal')[columns].mean()\n",
    "        y = df.query('Food == \"Unloaded\"').groupby('Animal')[columns].mean()\n",
    "\n",
    "        print('-----------------------------------------', file=file)\n",
    "        print(f'FIGURE {fig}: {columns}', file=file)\n",
    "        print(file=file)\n",
    "\n",
    "        if isinstance(columns, list):\n",
    "            # OMNIBUS TEST\n",
    "\n",
    "            hotelling_result = r_stats.T2_test(x.values-y.values)\n",
    "            print(f\"Hotelling's T-squared, T^2 = {hotelling_result['T2']:.3f}, F({hotelling_result['df_num']},{hotelling_result['df_den']}) = {hotelling_result['F']:.3f}, p = {hotelling_result['p']:.3f} {'(n.s.)' if hotelling_result['p'] > 0.05 else '(sig.)'}\", file=file)\n",
    "\n",
    "        else:\n",
    "            # POST-HOC TEST\n",
    "\n",
    "#             differences_test(x, y)\n",
    "\n",
    "            print(df.groupby(['Animal', 'Food'])[columns].apply(lambda x: {\n",
    "                'Mean ± SEM (N)': f'{x.mean():.2f} ± {x.sem():.2f} ({x.count()})'.ljust(18),\n",
    "            }).unstack([1, 2])[['Unloaded', 'Loaded']], file=file)\n",
    "            print(file=file)\n",
    "\n",
    "            print(f\"Difference of means, Mean ± SEM (N): {(x-y).mean():.2f} ± {(x-y).sem():.2f} ({(x-y).count()})\", file=file)\n",
    "            percent_change = ((x-y)/y).replace([-np.inf, np.inf], np.nan).dropna()\n",
    "            print(f\"Proportional change, Mean ± SEM (N): {100*percent_change.mean():.0f}% ± {100*percent_change.sem():.0f}% ({percent_change.count()})\", file=file)\n",
    "            print(file=file)\n",
    "\n",
    "            shapiro_result = r_stats.shapiro_test(x.values-y.values)\n",
    "            if shapiro_result['p'] > 0.05:\n",
    "                print(f\"Shapiro-Wilk, W = {shapiro_result['W']:.2f}, p = {shapiro_result['p']:.2f} (n.s.)\", file=file)\n",
    "            else:\n",
    "                print(f\"Shapiro-Wilk, W = {shapiro_result['W']:.2f}, p = {shapiro_result['p']:.3f} (sig.)\", file=file)\n",
    "\n",
    "            if shapiro_result['p'] > 0.05:\n",
    "                ttest_result = r_stats.t_test(x.values, y.values, paired=True, alternative='greater')\n",
    "                print(f\"Paired t-test, t({ttest_result['df']:g}) = {ttest_result['t']:.3f}, p = {ttest_result['p']:.3f} {'(n.s.)' if ttest_result['p'] > 0.05 else '(sig.)'}\", file=file)\n",
    "            else:\n",
    "                wilcoxon_result = r_stats.wilcox_test(x.values, y.values, paired=True, alternative='greater')\n",
    "                print(f\"Paired Wilcoxon signed rank, W = {wilcoxon_result['W']:g}, p = {wilcoxon_result['p']:.2f} {'(n.s.)' if wilcoxon_result['p'] > 0.05 else '(sig.)'}\", file=file)\n",
    "\n",
    "            cohen_d = r_stats.cohen_d(x.values, y.values)['estimate']\n",
    "            print(f\"Cohen's d = {cohen_d:.2f}\", file=file)\n",
    "\n",
    "        print(file=file)\n",
    "    \n",
    "    if filename is not None:\n",
    "        file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [TABLE 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure_map = {\n",
    "    '2C': 'Inward movement start to next inward movement start (s)',\n",
    "    '2D': 'Inward movement duration (s)',\n",
    "    '2E': 'Inward movement end to next inward movement start (s)',\n",
    "}\n",
    "\n",
    "filename = os.path.join(export_dir, 'table-2.txt')\n",
    "stats_table(figure_map, df_all.drop(unreliable_inward_movement), filename)\n",
    "\n",
    "if filename is not None:\n",
    "    with open(filename, 'r', encoding='utf8') as f:\n",
    "        print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [TABLE 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "figure_map = {\n",
    "    '4A+4B': ['B38 burst duration (s)', 'I2 burst duration (s)'],\n",
    "    '4A':    'B38 burst duration (s)',\n",
    "    '4B':    'I2 burst duration (s)',\n",
    "    '4C+4D': ['B8a/b burst duration (s)', 'B3/B6/B9 burst duration (s)'],\n",
    "    '4C':    'B8a/b burst duration (s)',\n",
    "    '4D':    'B3/B6/B9 burst duration (s)',\n",
    "    '4E':    'B4/B5 burst duration (s)',\n",
    "}\n",
    "\n",
    "filename = os.path.join(export_dir, 'table-3.txt')\n",
    "stats_table(figure_map, df_all.drop(bite_swallow_behaviors), filename)\n",
    "\n",
    "if filename is not None:\n",
    "    with open(filename, 'r', encoding='utf8') as f:\n",
    "        print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [TABLE 4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "figure_map = {\n",
    "    'X1+X2': ['B38 burst mean frequency (Hz)', 'I2 burst mean frequency (Hz)'],\n",
    "    'X1':    'B38 burst mean frequency (Hz)',\n",
    "    'X2':    'I2 burst mean frequency (Hz)',\n",
    "    'X3+4F': ['B8a/b burst mean frequency (Hz)', 'B3/B6/B9 burst mean frequency (Hz)'],\n",
    "    'X3':    'B8a/b burst mean frequency (Hz)',\n",
    "    '4F':    'B3/B6/B9 burst mean frequency (Hz)',\n",
    "    'X4':    'B4/B5 burst mean frequency (Hz)',\n",
    "}\n",
    "\n",
    "filename = os.path.join(export_dir, 'table-4.txt')\n",
    "stats_table(figure_map, df_all.drop(bite_swallow_behaviors), filename)\n",
    "\n",
    "if filename is not None:\n",
    "    with open(filename, 'r', encoding='utf8') as f:\n",
    "        print(f.read())"
   ]
  },
  {
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}