{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Untitled2.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyM0GTSsGEOZSJd2fzWz1C9M",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/rodjfg/los-fabulosos-pixelotls/blob/master/tools_luis.py\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "knDMZTMPdyp-"
      },
      "source": [
        "import os\n",
        "import requests\n",
        "import numpy as np\n",
        "from typing import List\n",
        "from scipy.ndimage import gaussian_filter1d\n",
        "from scipy.ndimage import gaussian_filter  # are they different?\n",
        "\n",
        "\n",
        "def load_raw_data(dat_path='.'):\n",
        "    '''get dataset in its raw form, just like in the notebook.\n",
        "    Returns:\n",
        "    (dict)   : data dictionary with the following keys.\n",
        "    `alldat` contains 39 sessions from 10 mice, data from Steinmetz et al, 2019. Time bins for all measurements are 10ms, starting 500ms before stimulus onset. The mouse had to determine which side has the highest contrast. For each `dat = alldat[k]`, you have the fields below. For extra variables, check out the extra notebook and extra data files (lfp, waveforms and exact spike times, non-binned). \n",
        "    * `dat['mouse_name']`: mouse name\n",
        "    * `dat['date_exp']`: when a session was performed\n",
        "    * `dat['spks']`: neurons by trials by time bins.    \n",
        "    * `dat['brain_area']`: brain area for each neuron recorded. \n",
        "    * `dat['ccf']`: Allen Institute brain atlas coordinates for each neuron. \n",
        "    * `dat['ccf_axes']`: axes names for the Allen CCF. \n",
        "    * `dat['contrast_right']`: contrast level for the right stimulus, which is always contralateral to the recorded brain areas.\n",
        "    * `dat['contrast_left']`: contrast level for left stimulus. \n",
        "    * `dat['gocue']`: when the go cue sound was played. \n",
        "    * `dat['response_times']`: when the response was registered, which has to be after the go cue. The mouse can turn the wheel before the go cue (and nearly always does!), but the stimulus on the screen won't move before the go cue.  \n",
        "    * `dat['response']`: which side the response was (`-1`, `0`, `1`). When the right-side stimulus had higher contrast, the correct choice was `-1`. `0` is a no go response. \n",
        "    * `dat['feedback_time']`: when feedback was provided. \n",
        "    * `dat['feedback_type']`: if the feedback was positive (`+1`, reward) or negative (`-1`, white noise burst).  \n",
        "    * `dat['wheel']`: turning speed of the wheel that the mice uses to make a response, sampled at `10ms`. \n",
        "    * `dat['pupil']`: pupil area  (noisy, because pupil is very small) + pupil horizontal and vertical position.\n",
        "    * `dat['face']`: average face motion energy from a video camera. \n",
        "    * `dat['licks']`: lick detections, 0 or 1.   \n",
        "    * `dat['trough_to_peak']`: measures the width of the action potential waveform for each neuron. Widths `<=10` samples are \"putative fast spiking neurons\". \n",
        "    * `dat['%X%_passive']`: same as above for `X` = {`spks`, `pupil`, `wheel`, `contrast_left`, `contrast_right`} but for  passive trials at the end of the recording when the mouse was no longer engaged and stopped making responses. \n",
        "    * `dat['prev_reward']`: time of the feedback (reward/white noise) on the previous trial in relation to the current stimulus time. \n",
        "    * `dat['reaction_time']`: ntrials by 2. First column: reaction time computed from the wheel movement as the first sample above `5` ticks/10ms bin. Second column: direction of the wheel movement (`0` = no move detected).  \n",
        "    '''\n",
        "\n",
        "    fname = []\n",
        "    for j in range(3):\n",
        "        fname.append(f'{dat_path}/steinmetz_part{j}.npz')\n",
        "    url = [\"https://osf.io/agvxh/download\"]\n",
        "    url.append(\"https://osf.io/uv3mw/download\")\n",
        "    url.append(\"https://osf.io/ehmw2/download\")\n",
        "\n",
        "    for j in range(len(url)):\n",
        "        if not os.path.isfile(fname[j]):\n",
        "            try:\n",
        "                r = requests.get(url[j])\n",
        "            except requests.ConnectionError:\n",
        "                print(\"!!! Failed to download data !!!\")\n",
        "            else:\n",
        "                if r.status_code != requests.codes.ok:\n",
        "                    print(\"!!! Failed to download data !!!\")\n",
        "                else:\n",
        "                    with open(fname[j], \"wb\") as fid:\n",
        "                        fid.write(r.content)\n",
        "\n",
        "    alldat = np.array([])\n",
        "    for j in range(len(fname)):\n",
        "        alldat = np.hstack((alldat, np.load(f'{dat_path}/steinmetz_part{j}.npz', allow_pickle=True)['dat']))\n",
        "\n",
        "    return alldat\n",
        "\n",
        "\n",
        "def select_by_areas(dat, selected_regions: List[str] = ['VISp']):\n",
        "    '''Find indices of neurons belonging to the specified areas in the dataset provided from a single experiment.\n",
        "    Args:\n",
        "        dat (dict): data dictionary from a single experiment. For example dat = alldat[11] in the example notebook.\n",
        "        selected_regions (list of str): list of region names to include. For example ['VISp']\n",
        "    Returns:\n",
        "        indices (1D array): array of indices from neurons belonging to specified regions.'''\n",
        "    return np.where(np.isin(dat['brain_area'], selected_regions))[0]\n",
        "\n",
        "\n",
        "def select_by_contrast(dat, contrast_pair: tuple = (1, 0)):\n",
        "    '''Find indices of trials with the specified contrasts in the dataset provided from a single experiment.\n",
        "    Args:\n",
        "        dat (dict): data dictionary from a single experiment. For example dat = alldat[11] in the example notebook.\n",
        "        contrast_pair (tuple): tuple with the contrast level of Left and Right stimulus. For example (1,0) represents stimulus on left with contrast 1 (max), and stimulus on right with contrast 0 (not shown).\n",
        "    Returns:\n",
        "        indices (1D array): array of indices from trials.'''\n",
        "    return np.where((dat['contrast_left'] == contrast_pair[0])*(dat['contrast_right'] == contrast_pair[1]))[0]\n",
        "\n",
        "\n",
        "def select_by_response(dat, response_type: str = 'to_left'):\n",
        "    '''Find indices of trials with the specified response in the dataset provided from a single experiment.\n",
        "    Args:\n",
        "        dat (dict): data dictionary from a single experiment. For example dat = alldat[11] in the example notebook.\n",
        "        response_type (str): str specifying direction of movement. \"to_left\", \"to_right\", or \"nogo\". For example, responses \"to_left\" are correct when stimulus with higher contrast is on the right side. \"nogo\" is when mouse did not move the wheel enough.\n",
        "    Returns:\n",
        "        indices (1D array): array of indices from trials.'''\n",
        "\n",
        "    if response_type == 'to_left':\n",
        "        indices = np.where(dat['response'] < 0)[0]\n",
        "    elif response_type == 'to_right':\n",
        "        indices = np.where(dat['response'] > 0)[0]\n",
        "    elif response_type == 'nogo':\n",
        "        indices = np.where(dat['response'] == 0)[0]\n",
        "    else:\n",
        "        raise Warning('wrong response_type input, choose: to_left, to_right or nogo')\n",
        "\n",
        "    return indices\n",
        "\n",
        "\n",
        "def select_trials(dat,  contrast_pair: tuple = (0, 1), response_type: str = 'to_left'):\n",
        "    '''Find indices of trials with the specified response and contrast levels in the dataset provided from a single experiment.\n",
        "    Args:\n",
        "        dat (dict): data dictionary from a single experiment. For example dat = alldat[11] in the example notebook.\n",
        "        contrast_pair (tuple): tuple with the contrast level of Left and Right stimulus. For example (1,0) represents stimulus on left with contrast 1 (max), and stimulus on right with contrast 0 (not shown).\n",
        "        response_type (str): str specifying direction of movement. \"to_left\", \"to_right\", or \"nogo\". For example, responses \"to_left\" are correct when stimulus with higher contrast is on the right side. \"nogo\" is when mouse did not move the wheel enough.\n",
        "    Returns:\n",
        "        indices (1D array): array of indices from trials.'''\n",
        "\n",
        "    contrast_indices = select_by_contrast(dat, contrast_pair)\n",
        "\n",
        "    if response_type == 'to_left':\n",
        "        response_indices = np.where(dat['response'] < 0)[0]\n",
        "    elif response_type == 'to_right':\n",
        "        response_indices = np.where(dat['response'] > 0)[0]\n",
        "    elif response_type == 'nogo':\n",
        "        response_indices = np.where(dat['response'] == 0)[0]\n",
        "    else:\n",
        "        raise Warning('wrong response_type input, choose: to_left, to_right or nogo')\n",
        "\n",
        "    # find intersection of both lists using sets and transforming back to a list\n",
        "    indices = list(set.intersection(set(contrast_indices), set(response_indices)))\n",
        "\n",
        "    return indices\n",
        "\n",
        "def select_by_outcome(dat, feedback_type: str = 'correct'):\n",
        "    '''Find indices of trials with the specified outcome (correct, incorrect) in the dataset provided from a single experiment. \n",
        "       Correct = Rewarded, Incorrect = punished. Note that succesful no-go trials indices are also included.\n",
        "    Args:\n",
        "        dat (dict): data dictionary from a single experiment. For example dat = alldat[11] in the example notebook.\n",
        "        response_type (str): str specifying outcome. \"correct\", \"incorrect\"\n",
        "    Returns:\n",
        "        indices (1D array): array of indices from trials.'''\n",
        "\n",
        "    if response_type == 'correct':\n",
        "        indices = np.where(dat['feedback_type'] < 0)[0]\n",
        "    elif response_type == 'incorrect':\n",
        "        indices = np.where(dat['feedback_type'] > 0)[0]\n",
        "  \n",
        "    else:\n",
        "        raise Warning('wrong response_type input, choose: correct or incorrect')\n",
        "\n",
        "    return indices\n",
        "\n",
        "\n",
        "def calculate_mean_firing_rate(spks, dt, mean_across: List[str] = ['population'], gaussfilter: bool = False, gauss_sigma: int = 1):\n",
        "    '''Find mean firing rate from the spikes (dat[\"spks\"]), across the chosen dimensions. Note: spks should contain spikes only of desired neurons, so find indices with the select_* functions, and apply the indices before giving the array of spikes to this function.\n",
        "    Args:\n",
        "        spks (3D array): array of spikes, from the \"spks\" field in the datasets. Dimensions should be (Neurons x Trials x Time)\n",
        "        mean_across (List[str]): list of str specifying type of mean to apply. It can be: A) \"population\", to only average across neurons, keeping trials and time, B) \"trials\" to average across trials, keeping neurons and time, C) \"time\" to average across time, keeping neurons and trials, or D) any combination of 2 from the previous.\n",
        "    Returns:\n",
        "        mean_firing_rates (ND array): array of N dimensions, where N=3-len(mean_across), containing mean firing rates according to parameters.'''\n",
        "\n",
        "    dimension_order = ['population', 'trials', 'time']\n",
        "    axis_to_mean = tuple([dimension_order.index(x) for x in mean_across])\n",
        "    mean_firing_rate = 1/dt * spks.mean(axis=axis_to_mean)\n",
        "    if gaussfilter:\n",
        "        mean_firing_rate = gaussian_filter1d(mean_firing_rate, gauss_sigma)\n",
        "\n",
        "    return mean_firing_rate\n",
        "\n",
        "\n",
        "def find_response_type(contrast):\n",
        "    '''find which is the correct response_type of the mouse given the pair of contrast given. nogo for same level, to_right when contrast is higher on left, and to_left when contrast is higher on right.\n",
        "    Args:\n",
        "        contrast_pair (tuple): tuple with the contrast level of Left and Right stimulus. For example (1,0) represents stimulus on left with contrast 1 (max), and stimulus on right with contrast 0 (not shown).\n",
        "    Returns:\n",
        "        correct_response (str): string corresponding to the response_type that is correct for the given contrast.'''\n",
        "\n",
        "    if contrast[1] == contrast[0]:\n",
        "        correct_response = 'nogo'\n",
        "    elif contrast[0] > contrast[1]:\n",
        "        correct_response = 'to_right'\n",
        "    elif contrast[0] < contrast[1]:\n",
        "        correct_response = 'to_left'\n",
        "    return correct_response\n",
        "\n",
        "\n",
        "def collect_firing_rates(alldat, contrast_pair, selected_regions: List[str] = [\"VISp\"], gaussfilter=True, gauss_sigma=1):\n",
        "    '''\n",
        "    Args:\n",
        "        alldat (dict): data dictionary from full dataset. Output from load_raw_data().\n",
        "        contrast_pair (tuple): tuple with the contrast level of Left and Right stimulus. For example (1,0) represents stimulus on left with contrast 1 (max), and stimulus on right with contrast 0 (not shown).\n",
        "        selected_regions (list of str): list of region names to include. For example ['VISp']\n",
        "    Returns:\n",
        "        all_correct_fr (2D array, N_trials, N_timebins): population average for all correct trials, timebins and experiments, for the given contrast.\n",
        "        all_incorrect_fr (2D array, N_trials, N_timebins): population average for all incorrect trials, timebins and experiments, for the given contrast.\n",
        "    '''\n",
        "\n",
        "    all_response_types = ['nogo', 'to_left', 'to_right']\n",
        "    all_correct_fr, all_incorrect_fr = None, None\n",
        "\n",
        "    for dat in alldat:\n",
        "        # get important variables from specific experiment\n",
        "        dt = dat['bin_size']\n",
        "        correct_response = find_response_type(contrast_pair)\n",
        "\n",
        "        # find neurons indices from selected area\n",
        "        neurons_indices_in_area = select_by_areas(dat, selected_regions=selected_regions)\n",
        "\n",
        "        # only continue if experiment has neurons in the region of interest\n",
        "        if len(neurons_indices_in_area) > 0:\n",
        "            # selected trials per contrast level and response\n",
        "            correct_trials = select_trials(dat, contrast_pair=contrast_pair, response_type=correct_response)\n",
        "            incorrect_trials = []\n",
        "            for response_type in all_response_types:\n",
        "                if response_type != correct_response:\n",
        "                    incorrect_trials.extend(select_trials(dat, contrast_pair=contrast_pair, response_type=response_type))\n",
        "\n",
        "            # store firing rates of correct trials\n",
        "            if len(correct_trials) > 0:\n",
        "                correct_spks = dat['spks'][neurons_indices_in_area][:, correct_trials]\n",
        "                correct_fr = calculate_mean_firing_rate(correct_spks, dt, ['population'], gaussfilter=gaussfilter, gauss_sigma=gauss_sigma)\n",
        "                if all_correct_fr is None:\n",
        "                    all_correct_fr = correct_fr\n",
        "                else:\n",
        "                    all_correct_fr = np.concatenate((all_correct_fr, correct_fr), axis=0)\n",
        "            # store firing rates of incorrect trials\n",
        "            if len(incorrect_trials) > 0:\n",
        "                incorrect_spks = dat['spks'][neurons_indices_in_area][:, incorrect_trials]\n",
        "                incorrect_fr = calculate_mean_firing_rate(incorrect_spks, dt, ['population'], gaussfilter=gaussfilter, gauss_sigma=gauss_sigma)\n",
        "                if all_incorrect_fr is None:\n",
        "                    all_incorrect_fr = incorrect_fr\n",
        "                else:\n",
        "                    all_incorrect_fr = np.concatenate((all_incorrect_fr, incorrect_fr), axis=0)\n",
        "    return all_correct_fr, all_incorrect_fr\n",
        "\n",
        "\n",
        "def choose_time_window(x):\n",
        "    # input: x: vector time series\n",
        "    # output: idxLimits: timestamps for window of size \"size\" after the first peak is skipped\n",
        "    x = gaussian_filter(x, sigma=1)  # if not preprocessed\n",
        "\n",
        "    # find index of the first peak, where our function will start looking for the end of the peak(which will be the left end of our window)\n",
        "    peak = np.argmax(x)\n",
        "    i = peak\n",
        "    while x[i+1] < x[i]:\n",
        "        i += 1\n",
        "    idx0 = i\n",
        "    while x[i] > x[idx0]:\n",
        "        if i > len(x):\n",
        "            break\n",
        "        i += 1\n",
        "    idxf = i\n",
        "\n",
        "    idxLimits = [idx0, idxf]\n",
        "\n",
        "    return idxLimits"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}