5-label_acoustics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "41254766-7844-485b-ae7b-f5953d62a1da",
   "metadata": {},
   "source": [
    "# Label acoustic data\n",
    "\n",
    "Optional notebook within the chronic ephys processing pipeline\n",
    "- 1-preprocess_acoustics\n",
    "- 2-curate_acoustics\n",
    "- 3-sort_spikes\n",
    "- 4-curate_spikes\n",
    "- **5-label_acoustics**\n",
    "\n",
    "*Currently contains functionality to label social context and syllables*\n",
    "\n",
    "Use the environment **birdsong** to run this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04e33001-b4ad-4fb3-93f5-8fe7dc718a93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import logging\n",
    "import socket\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from matplotlib import pyplot as plt\n",
    "import seaborn as sns\n",
    "from scipy import stats\n",
    "from scipy import signal\n",
    "from scipy.io import wavfile\n",
    "import IPython.display as ipd\n",
    "from tqdm.autonotebook import tqdm\n",
    "from joblib import Parallel, delayed\n",
    "from datetime import datetime\n",
    "import copy\n",
    "import sys\n",
    "\n",
    "from praatio import textgrid\n",
    "from praatio import audio\n",
    "\n",
    "sys.path.append('/mnt/cube/lo/envs')\n",
    "from plot_sonogram import plot_sonogram as ps\n",
    "\n",
    "sys.path.append('/mnt/cube/lo/envs/ceciestunepipe')\n",
    "from ceciestunepipe.file import bcistructure as et\n",
    "\n",
    "sys.path.append('/mnt/cube/lo/envs/vocalization-segmentation')\n",
    "from vocalseg.continuity_filtering import plot_labelled_elements\n",
    "\n",
    "sys.path.append('/mnt/cube/lo/envs/avgn_paper')\n",
    "from avgn.signalprocessing.filtering import butter_bandpass_filter\n",
    "from avgn.utils.hparams import HParams\n",
    "from avgn.signalprocessing.filtering import prepare_mel_matrix\n",
    "from avgn.signalprocessing.create_spectrogram_dataset import make_spec, mask_spec, log_resize_spec, pad_spectrogram, flatten_spectrograms\n",
    "from avgn.visualization.spectrogram import draw_spec_set\n",
    "from avgn.visualization.quickplots import draw_projection_plots\n",
    "from avgn.visualization.projections import scatter_spec\n",
    "from avgn.visualization.barcodes import plot_sorted_barcodes\n",
    "from avgn.visualization.network_graph import plot_network_graph\n",
    "\n",
    "np.set_printoptions(precision=3, suppress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77909c74-8a61-465c-a42b-21e6bb1bcb1d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# session parameters\n",
    "sess_par = {\n",
    "    'bird':'z_r12r13_21', # bird ID\n",
    "    'sess':'2021-06-27', # session date\n",
    "    'ephys_software':'sglx', # recording software, sglx or oe\n",
    "    'stim_sess':False, # if song stimulus was played during the session, ignore detected bouts\n",
    "    'trim_bouts':True, # manually trim bouts after curation\n",
    "    'sort':'sort_0', # sort index\n",
    "}\n",
    "\n",
    "time_F_in = '00:00:00'\n",
    "time_F_out = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80b3559f-c533-4fde-8bc6-30cbf60c2799",
   "metadata": {},
   "source": [
    "## Load curated acoustics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a1f0553-ea6b-4055-8f6c-4cde8e2d72e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sess_epochs = et.list_ephys_epochs(sess_par)\n",
    "print(f\"Found {len(sess_epochs)} epoch(s):\", sess_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64750695-6bca-40b9-bfb3-f1f94dd030a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "this_epoch = sess_epochs[0] # set epoch index\n",
    "epoch_struct = et.sgl_struct(sess_par,this_epoch,ephys_software=sess_par['ephys_software'])\n",
    "print('Processing epoch', this_epoch)\n",
    "\n",
    "# load bout dataframe\n",
    "bout_df_path = os.path.join(epoch_struct['folders']['derived'],'bout_pd_ap0_curated.pkl')\n",
    "with open(bout_df_path, 'rb') as handle:\n",
    "    bout_df = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bec368f-b883-421d-8c8e-4e4c136420eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get sample rate\n",
    "if len(bout_df.sample_rate.unique()) > 1:\n",
    "    print(f\"{len(bout_df.sample_rate.unique())} sample rates found:\", bout_df.sample_rate.unique())\n",
    "fs = bout_df.sample_rate.unique()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6e94f33-cb19-4f09-9969-8a85036b4e4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get neural sample rate\n",
    "ap_path = os.path.join(epoch_struct['folders']['derived'],'ap_0_sync_dict.pkl')\n",
    "with open(ap_path, 'rb') as handle:\n",
    "    ap_syn_dict = pickle.load(handle)\n",
    "ap_fs = ap_syn_dict['s_f']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ccd5e86-d760-4a13-8584-03e396346fc9",
   "metadata": {},
   "source": [
    "## Add social context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5206ec65-1bad-4057-9e28-8d4e82bfbc4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bout_start(start_ms):\n",
    "    hour = int(np.floor(start_ms/3600000))\n",
    "    minute = int(np.floor((start_ms/60000)-(hour*60)))\n",
    "    second = int(np.floor(start_ms % 60000)/1000)\n",
    "    \n",
    "    bout_start = datetime.strptime(f\"{hour:02}:{minute:02}:{second:02}\", \"%H:%M:%S\").time()\n",
    "    \n",
    "    return bout_start\n",
    "\n",
    "def set_behavior(row, F_in_dt=None, F_out_dt=None):\n",
    "    bout_start = get_bout_start(row['start_ms'])\n",
    "    if F_in_dt and bout_start < F_in_dt:\n",
    "        return 'undirected'\n",
    "    elif F_out_dt and bout_start > F_out_dt:\n",
    "        return 'undirected'\n",
    "    else:\n",
    "        return 'directed'\n",
    "\n",
    "def add_social_context(bout_df_in, time_F_in=None, time_F_out=None):\n",
    "    bout_df_out = bout_df_in.copy()\n",
    "    \n",
    "    if time_F_in:\n",
    "        F_in_dt = datetime.strptime(f\"{time_F_in}\", \"%H:%M:%S\").time()\n",
    "        print('Female introduced at', F_in_dt, '\\n')\n",
    "        bout_df_out['behavior'] = Parallel(n_jobs=-1)(delayed(set_behavior)(row, F_in_dt=F_in_dt) for _, row in bout_df_out.iterrows())\n",
    "        \n",
    "    elif time_F_out:\n",
    "        F_out_dt = datetime.strptime(f\"{time_F_out}\", \"%H:%M:%S\").time()\n",
    "        print('Female removed at', F_out_dt, '\\n')\n",
    "        bout_df_out['behavior'] = Parallel(n_jobs=-1)(delayed(set_behavior)(row, F_out_dt=F_out_dt) for _, row in bout_df_out.iterrows())\n",
    "    \n",
    "    return bout_df_out\n",
    "\n",
    "bout_df = add_social_context(bout_df, time_F_in, time_F_out)\n",
    "\n",
    "print(len(bout_df[bout_df['behavior']=='undirected']), 'undirected bouts')\n",
    "print(len(bout_df[bout_df['behavior']=='directed']), 'directed bouts')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f959703-6b1b-43ac-8caf-b65f342395c4",
   "metadata": {},
   "source": [
    "## Export wav files to label in Praat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "157d1de5-f1ee-493b-9e99-9d3f9544c145",
   "metadata": {},
   "outputs": [],
   "source": [
    "praat_dir = os.path.join(epoch_struct['folders']['derived'],'praat')\n",
    "os.makedirs(praat_dir, exist_ok=True)\n",
    "\n",
    "for idx, row in bout_df.iterrows():\n",
    "    file_path = os.path.join(praat_dir, f\"{idx}-{row['start_ms']}.wav\")\n",
    "    wavfile.write(file_path, fs, row['waveform'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1612d733-0312-48bd-9d1c-e0e226c6e313",
   "metadata": {},
   "source": [
    "## Import TextGrid files from Praat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bca7553c-a1e1-4ccd-b27b-38c7783f13a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "bouts_segmented = bout_df.copy()\n",
    "bouts_segmented['bout_waveform_filt'] = bouts_segmented.apply(lambda r: butter_bandpass_filter(r['waveform'], 300, 12000, r['sample_rate']), axis=1)\n",
    "bouts_segmented.rename(columns={'waveform': 'bout_waveform_raw'}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc99e206-34d7-4618-9d45-5b98a2c8bc6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a dataframe for segmented syllables\n",
    "praat_dir = os.path.join(epoch_struct['folders']['derived'],'praat')\n",
    "dfs = []\n",
    "for index, row in bouts_segmented.iterrows():\n",
    "    \n",
    "    # Info from bouts\n",
    "    file = row['file']\n",
    "    sess = row['sess']\n",
    "    epoch = row['epoch']\n",
    "    sample_rate = row['sample_rate']\n",
    "    bout_index = index\n",
    "    bout_waveform_raw = row['bout_waveform_raw']\n",
    "    bout_waveform_filt = row['bout_waveform_filt']\n",
    "    start_ms_bout = row['start_ms']\n",
    "    start_sample_bout = row['start_sample']\n",
    "    start_ms_ap_0_bout = row['start_ms_ap_0']\n",
    "    start_sample_ap_0_bout = row['start_sample_ap_0']\n",
    "    \n",
    "    # Syllable labels from praat\n",
    "    tg = textgrid.openTextgrid(os.path.join(praat_dir,f\"{index}-{row['start_ms']}.TextGrid\"),\n",
    "                               includeEmptyIntervals=False)\n",
    "    syllables = tg.getTier(tg.tierNames[0])\n",
    "    on_ss = [interval.start for interval in syllables.entries]\n",
    "    off_ss = [interval.end for interval in syllables.entries]\n",
    "    labels = [interval.label for interval in syllables.entries]\n",
    "    \n",
    "    data = []\n",
    "    for syllable_index, (on_s, \n",
    "                         off_s,\n",
    "                         label) in enumerate( zip(on_ss, \n",
    "                                                   off_ss,\n",
    "                                                   labels)\n",
    "                                             ):\n",
    "        on_sample = int(start_sample_bout + on_s*fs)\n",
    "        off_sample = int(start_sample_bout + off_s*fs)\n",
    "        \n",
    "        data.append({\n",
    "            'file': file,\n",
    "            'sess': sess,\n",
    "            'epoch': epoch,\n",
    "            'sample_rate': sample_rate,\n",
    "            'bout_index': bout_index,\n",
    "            'bout_waveform_raw': bout_waveform_raw,\n",
    "            'bout_waveform_filt': bout_waveform_filt,\n",
    "            'start_ms_ap_0': int(start_ms_ap_0_bout + on_s*1000),\n",
    "            'start_sample_ap_0': int(start_sample_ap_0_bout + on_sample/fs*ap_fs),\n",
    "            'syllable_index': syllable_index,\n",
    "            'on_sample': on_sample,\n",
    "            'off_sample': off_sample,\n",
    "            'on_ms': int(start_ms_bout + on_s*1000),\n",
    "            'off_ms': int(start_ms_bout + off_s*1000),\n",
    "            'label': label,\n",
    "            'syllable_waveform': bout_waveform_filt[int(on_s*fs):int(off_s*fs)]})\n",
    "    \n",
    "    df = pd.DataFrame(data)\n",
    "    dfs.append(df)\n",
    "\n",
    "syl_df = pd.concat(dfs, ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f9f6870-3899-46d0-9f3a-813fe496c482",
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalize audio\n",
    "syl_df['syllable_waveform'] = [syll/max(np.min(syll), np.max(syll), key=abs) for i, syll in enumerate(syl_df['syllable_waveform'].values)]\n",
    "syl_df['syllable_waveform'] = [np.nan_to_num(syll) if not np.all(np.isfinite(syll)) else syll for syll in syl_df['syllable_waveform'].values]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7de440c-2043-4087-93b5-3ed9a854aab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot some of the syllables to see how they look\n",
    "nrows = 10\n",
    "ncols = 10\n",
    "zoom = 2\n",
    "fig, axs = plt.subplots(ncols=ncols, nrows = nrows, figsize = (ncols*zoom, nrows+zoom/1.5))\n",
    "for i, syll in tqdm(enumerate(syl_df['syllable_waveform'].values), total = nrows*ncols):\n",
    "    ax = axs.flatten()[i]\n",
    "    ax.plot(syll)\n",
    "    if i == nrows*ncols-1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4aa21f6-d36d-4efb-8999-a82ccb22b34c",
   "metadata": {},
   "source": [
    "## Plot syllable spectrograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d8fd6bb-ac9a-49ca-86fa-b9cdd4482fa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "syllables_wav = syl_df.syllable_waveform.values\n",
    "syllables_rate = syl_df.sample_rate.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a2f8ab-3192-4e16-9902-15b69c6d1141",
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams = HParams(\n",
    "    num_mel_bins = 64,\n",
    "    mel_lower_edge_hertz=300,\n",
    "    mel_upper_edge_hertz=12000,\n",
    "    butter_lowcut = 300,\n",
    "    butter_highcut = 12000,\n",
    "    ref_level_db = 20,\n",
    "    min_level_db = -100,\n",
    "    mask_spec = True,\n",
    "    win_length_ms = 4,\n",
    "    hop_length_ms = 1,\n",
    "    nex = -1,\n",
    "    n_jobs = -1,\n",
    "    verbosity = 0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "657b72c2-90e0-490e-aba4-ac6a04fda08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_jobs = 36\n",
    "verbosity = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "308b883f-7778-4768-8b25-a8c6f382f186",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create spectrograms\n",
    "with Parallel(n_jobs=n_jobs, verbose=verbosity) as parallel:\n",
    "    syllables_spec = parallel(\n",
    "        delayed(make_spec)(\n",
    "            syllable,\n",
    "            rate,\n",
    "            hparams=hparams,\n",
    "            mel_matrix=prepare_mel_matrix(hparams, rate),\n",
    "            use_mel=True,\n",
    "            use_tensorflow=False,\n",
    "        )\n",
    "        for syllable, rate in tqdm(\n",
    "            zip(syllables_wav, syllables_rate),\n",
    "            total=len(syllables_rate),\n",
    "            desc=\"getting syllable spectrograms\",\n",
    "            leave=False,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54386ca0-9dd5-486e-a053-b26b85f7a69b",
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_spec_set(syllables_spec, zoom=1, maxrows=10, colsize=40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41d80b1b-6270-4d05-add9-167f13b88449",
   "metadata": {},
   "outputs": [],
   "source": [
    "# log rescale spectrograms\n",
    "log_scaling_factor = 4\n",
    "\n",
    "with Parallel(n_jobs=n_jobs, verbose=verbosity) as parallel:\n",
    "    syllables_spec = parallel(\n",
    "        delayed(log_resize_spec)(spec, scaling_factor=log_scaling_factor)\n",
    "        for spec in tqdm(syllables_spec, desc=\"scaling spectrograms\", leave=False)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a46ab590-1240-4498-9e8e-a9efc1750496",
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_spec_set(syllables_spec, zoom=1, maxrows=10, colsize=40)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b58e9403-4a68-49c5-b349-23e7b8e594dd",
   "metadata": {},
   "source": [
    "## Plot syllable barcodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d25aaab6-7d90-4116-9670-d19165483552",
   "metadata": {},
   "outputs": [],
   "source": [
    "def song_barcode(start_times, stop_times, labels, label_dict, label_pal_dict, resolution=0.01):\n",
    "    begin = np.min(start_times)\n",
    "    end = np.max(stop_times)\n",
    "    trans_list = (\n",
    "        np.zeros(int((end - begin) / resolution)).astype(\"str\").astype(\"object\")\n",
    "    )\n",
    "    # print(end, begin, end-begin, resolution, len(trans_list))\n",
    "    for start, stop, label in zip(start_times, stop_times, labels):\n",
    "        trans_list[\n",
    "            int((start - begin) / resolution) : int((stop - begin) / resolution)\n",
    "        ] = label_dict[label]\n",
    "\n",
    "    color_list = [\n",
    "        label_pal_dict[i] if i in label_pal_dict else [1, 1, 1] for i in trans_list\n",
    "    ]\n",
    "    color_list = np.expand_dims(color_list, 1)\n",
    "\n",
    "    return trans_list, color_list\n",
    "\n",
    "\n",
    "def indv_barcode(this_df, time_resolution=0.01, label=\"label\", pal=\"tab20\"):\n",
    "    unique_labels = this_df[label].unique()\n",
    "    \n",
    "    # song palette\n",
    "    label_pal = np.random.permutation(sns.color_palette(pal, len(unique_labels)))\n",
    "    label_dict = {lab: str(int(i)).zfill(3) for i, lab in enumerate(unique_labels)}\n",
    "\n",
    "    label_pal_dict = {\n",
    "        label_dict[lab]: color for lab, color in zip(unique_labels, label_pal)\n",
    "    }\n",
    "    sns.palplot(list(label_pal_dict.values()))\n",
    "\n",
    "    # get list of syllables by time\n",
    "    trans_lists = []\n",
    "    color_lists = []\n",
    "    for key in tqdm(this_df.bout_index.unique(), leave=False):\n",
    "        # dataframe of wavs\n",
    "        wav_df = this_df[this_df['bout_index'] == key]\n",
    "        labels = wav_df[label].values\n",
    "        start_times = wav_df.on_ms.values\n",
    "        stop_times = wav_df.off_ms.values\n",
    "        trans_list, color_list = song_barcode(\n",
    "            start_times,\n",
    "            stop_times,\n",
    "            labels,\n",
    "            label_dict,\n",
    "            label_pal_dict,\n",
    "            resolution=time_resolution,\n",
    "        )\n",
    "        color_lists.append(color_list)\n",
    "        trans_lists.append(trans_list)\n",
    "\n",
    "    return color_lists, trans_lists, label_pal_dict, label_pal, label_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa4c9a19-3b57-4f16-be13-b9da2f00e0c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get variables for plotting\n",
    "print('Syllable barcodes: ' + str(syl_df.label.unique()))\n",
    "\n",
    "color_lists, trans_lists, label_pal_dict, label_pal, label_dict = indv_barcode(\n",
    "    syl_df,\n",
    "    time_resolution=12\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "232c03b8-e70b-4f79-9395-1307b49dffa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot syllable barcodes for songs\n",
    "ids = syl_df.bout_index.unique()\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(20, 3))\n",
    "plot_sorted_barcodes(\n",
    "    [color_lists[i] for i in ids],\n",
    "    [trans_lists[i] for i in ids],\n",
    "    max_list_len=600,\n",
    "    seq_len=100,\n",
    "    nex=200,\n",
    "    figsize=(10, 4),\n",
    "    ax=ax,\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36abb351-0d3e-4d52-bdc8-32b998012857",
   "metadata": {},
   "source": [
    "## Save syl_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3654cab3-a70e-46e5-89f3-e4c3e5d4aa2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "syl_df.to_pickle(os.path.join(epoch_struct['folders']['derived'],'syl_df_ap0.pickle'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6542bb4-80ff-427c-9854-d2a5a5e4c680",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "birdsong",
   "language": "python",
   "name": "birdsong"
  },
  "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}