02_processing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# processing\n",
    "> Processing the different stream of data to calculate responses of retinal cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from functools import partial\n",
    "import numpy as np\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn import cluster\n",
    "import scipy.ndimage as ndimage\n",
    "import scipy.signal as signal\n",
    "import scipy as sp\n",
    "from cmath import *\n",
    "import itertools\n",
    "import random\n",
    "\n",
    "from theonerig.core import *\n",
    "from theonerig.utils import *\n",
    "from theonerig.modelling import *\n",
    "from theonerig.leddome import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spatial stimulus (checkerboard)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def eyetrack_stim_inten(stim_inten, eye_track, \n",
    "                        upsampling=2,\n",
    "                        eye_calib=[[94 ,8], [ 18, 59]],\n",
    "                        box_w=None, box_h=None, stim_axis=\"x\"):\n",
    "    \"\"\"\n",
    "    From stimulus data and eye tracking, returns a corrected and upsampled stimulus data.\n",
    "    Calibration corresponds to the width and height of the stimulus screen, in \n",
    "    terms of pixels in the eye video: [[WIDTH_x, WIDTH_y], [HEIGHT_x, HEIGHT_y]]\n",
    "    \n",
    "    params:\n",
    "        - stim_inten: Stimulus intensity matrix of shape (t, y, x), or (t, x) or (t, y) depending on stim_axis\n",
    "        - eye_track: Eye tracking data of shape (t, x_pos, y_pos, ...)\n",
    "        - upsampling: Factor for the upsampling (2 will multiply by 2 number of box in width and height)\n",
    "        - eye_calib: Calibration matrix of shape (2,2)\n",
    "        - box_w: Width of a block in pixel (40px in case of a 32 box in width of a checkerboard on a 1280px width)\n",
    "        - box_h: Height of a block in pixel. Both box_x and box_h are calculated from a 1280x720 screen if None\n",
    "        - stim_axis: Specify which direction to shift in case of stim shape different than (t, y, x)\n",
    "        \n",
    "    return:\n",
    "        - Upsampled and shift corrected stimulus intensity\n",
    "    \"\"\"\n",
    "    eye_x, eye_y = eye_track[:,0], eye_track[:,1]\n",
    "    shape_y, shape_x = 1, 1\n",
    "    if len(stim_inten.shape)==2:\n",
    "        if stim_axis==\"x\":\n",
    "            shape_x = stim_inten.shape[1]\n",
    "        elif stim_axis==\"y\":\n",
    "            shape_y = stim_inten.shape[1]\n",
    "    elif len(stim_inten.shape)==3:\n",
    "        shape_y = stim_inten.shape[1]\n",
    "        shape_x = stim_inten.shape[2]\n",
    "    if box_w is None:\n",
    "        box_w = 1280//shape_x\n",
    "    if box_h is None:\n",
    "        box_h = 720//shape_y\n",
    "\n",
    "    if shape_y>1 and shape_x>1:\n",
    "        box_w, box_h = int(box_w/upsampling), int(box_h/upsampling)\n",
    "    elif shape_x > 1:\n",
    "        box_w, box_h = int(box_w/upsampling), box_h\n",
    "    elif shape_y > 1:\n",
    "        box_w, box_h = box_w                , int(box_h/upsampling)\n",
    "\n",
    "    eye_transfo_f = _eye_to_stim_f(eye_calib=eye_calib, \n",
    "                                  box_width=box_w,\n",
    "                                  box_height=box_h)\n",
    "    \n",
    "    if shape_y>1 and shape_x>1:\n",
    "        stim_inten = stim_inten.repeat(upsampling,axis=1).repeat(upsampling,axis=2)\n",
    "    else:\n",
    "        stim_inten = stim_inten.repeat(upsampling,axis=1)\n",
    "        \n",
    "    xpos_avg = np.mean(eye_x)\n",
    "    ypos_avg = np.mean(eye_y)\n",
    "    mean_stim_inten = int((np.max(stim_inten)+np.min(stim_inten))/2)\n",
    "    #After getting the shift of the matrix to apply, we roll the matrix instead of extending it to the shifts\n",
    "    #This seems strange, but from the cell point of view, that is potentially looking at no stimulus,\n",
    "    # the response it gives are uncorrelated with the stimulus, and so shouldn't impact further analysis\n",
    "    # Advantage is that it keeps the data small enough, without loosing regions of the stimulus.\n",
    "    for i in range(len(stim_inten)):\n",
    "        x_eyeShift = eye_x[i]-xpos_avg\n",
    "        y_eyeShift = eye_y[i]-ypos_avg\n",
    "        \n",
    "        stim_shift_x, stim_shift_y = eye_transfo_f(x_eyeShift=x_eyeShift,\n",
    "                                                   y_eyeShift=y_eyeShift) \n",
    "        if shape_y>1 and shape_x>1:\n",
    "            rolled_stim = np.roll(stim_inten[i],stim_shift_y,axis=0)\n",
    "            rolled_stim = np.roll(rolled_stim  ,stim_shift_x,axis=1)\n",
    "        else:\n",
    "            if stim_axis==\"x\":\n",
    "                rolled_stim = np.roll(stim_inten[i],stim_shift_x,axis=0)\n",
    "            else:\n",
    "                rolled_stim = np.roll(stim_inten[i],stim_shift_y,axis=0)\n",
    "        stim_inten[i] = rolled_stim\n",
    "        \n",
    "    return stim_inten\n",
    "\n",
    "def saccade_distances(eye_track):\n",
    "    \"\"\"\n",
    "    Create a mask for the eye position timeserie that indicate how far was the last saccade.\n",
    "    The eye positions need to be smoothed with smooth_eye_position.\n",
    "    \n",
    "    params:\n",
    "        - eye_track: Eye tracking data of shape (t, x_pos, y_pos, ...)\n",
    "        \n",
    "    return:\n",
    "        - Distance to last saccade array\n",
    "    \"\"\"\n",
    "    \n",
    "    x_pos, y_pos = eye_track[:,0], eye_track[:,1]\n",
    "    saccade_pos = np.where((x_pos[1:] != x_pos[:-1]) & (y_pos[1:] != y_pos[:-1]))[0] + 1\n",
    "    \n",
    "    len_chunks = [saccade_pos[0]]+list(saccade_pos[1:]-saccade_pos[:-1])\n",
    "    len_chunks.append(len(x_pos) - saccade_pos[-1])\n",
    "    \n",
    "    saccade_mask = []\n",
    "    for len_chunk in len_chunks:\n",
    "        saccade_mask.extend(list(range(len_chunk)))\n",
    "    \n",
    "    return np.array(saccade_mask)\n",
    "\n",
    "def smooth_eye_position(eye_track, threshold=2):\n",
    "    \"\"\"\n",
    "    Smooth the eye positions. Uses clustering of eye position epochs to smooth.\n",
    "    Clustering is made with sklearn.cluster.dbscan.\n",
    "    \n",
    "    params:\n",
    "        - eye_track: Eye tracking data of shape (t, x_pos, y_pos, ...)\n",
    "        - threshold: Distance threshold in the position clustering\n",
    "        \n",
    "    return:\n",
    "        - Distance to last saccade array\n",
    "    \"\"\"\n",
    "    x_pos, y_pos = eye_position[:,0], eye_position[:,1]\n",
    "\n",
    "    X = np.stack((x_pos,y_pos, np.linspace(0, len(x_pos)/2, len(x_pos)))).T\n",
    "    clusters = cluster.dbscan(X, eps=threshold, min_samples=3, metric='minkowski', p=2)\n",
    "\n",
    "    move_events = np.where(clusters[1][1:] > clusters[1][:-1])[0] + 1\n",
    "    \n",
    "    len_chunks = [move_events[0]]+list(move_events[1:]-move_events[:-1])\n",
    "    len_chunks.append(len(x_pos) - move_events[-1])\n",
    "\n",
    "    eye_x_positions = np.split(x_pos, move_events)\n",
    "    eye_y_positions = np.split(y_pos, move_events)\n",
    "\n",
    "    mean_x_pos = np.array(list(map(np.mean, eye_x_positions)))\n",
    "    mean_y_pos = np.array(list(map(np.mean, eye_y_positions)))\n",
    "    \n",
    "    x_pos_smooth = np.concatenate([[x_pos]*len_chunk for x_pos,len_chunk in zip(mean_x_pos, len_chunks)])\n",
    "    y_pos_smooth = np.concatenate([[y_pos]*len_chunk for y_pos,len_chunk in zip(mean_y_pos, len_chunks)])\n",
    "    \n",
    "    return np.stack((x_pos_smooth, y_pos_smooth)).T\n",
    "\n",
    "def _eye_to_stim_f(eye_calib, box_width, box_height):\n",
    "    \"\"\"\n",
    "    Initialise stimulus shift function parameters and returns a partial function.\n",
    "    \n",
    "    params:\n",
    "        - eye_calib: Calibration matrix of shape (2,2)\n",
    "        - box_width: Width of the elementary block in pixel\n",
    "        - box_height: Width of the elementary block in pixel\n",
    "        \n",
    "    return:\n",
    "        - Partial function to obtain stimulus shift from eye positions\n",
    "    \"\"\"\n",
    "    eye_to_stim     = np.linalg.inv(eye_calib)\n",
    "    box_dim         = np.array([1280/(box_width), 720/(box_height)])\n",
    "    return partial(_linear_transform, box_dim=box_dim, transfo_matrix=eye_to_stim)\n",
    "    \n",
    "def _linear_transform(box_dim, transfo_matrix, x_eyeShift, y_eyeShift):\n",
    "    \"\"\"\n",
    "    Computes the shift of the stimulus from eye position shifts.\n",
    "    \n",
    "    params:\n",
    "        - box_dim: Size in pixel of the box [width, height]\n",
    "        - transfo_matrix: Inverse matrix of the calibration matrix described in eyetrack_stim_inten\n",
    "        - x_eyeShift: Eye shift in x\n",
    "        - y_eyeShift: Eye shift in y\n",
    "        \n",
    "    return:\n",
    "        - Stimulus shift tuple (delta_x, delta_y)\n",
    "    \"\"\"\n",
    "    transform_coord = np.dot(transfo_matrix, np.array([x_eyeShift, y_eyeShift]).T)\n",
    "    stim_vec        = np.round(transform_coord * box_dim).astype(int)\n",
    "    return stim_vec[0], -stim_vec[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def process_sta_batch(stim_inten, spike_counts, Hw=30, Fw=2, return_pval=False, normalisation=\"abs\"):\n",
    "    \"\"\"\n",
    "    Computes the STA and associated pvalues in parallel for a batch of cells.\n",
    "    \n",
    "    params:\n",
    "        - stim_inten: stimulus intensity matrix of shape (t, ...)\n",
    "        - spike_counts: cells activity matrix of shape (t, n_cell)\n",
    "        - Hw: Lenght in frames of the history window, including the 0 timepoint\n",
    "        - Fw: Lenght in frames of the forward window\n",
    "        - return_pval: Flag to signal whether or not to return the pvalues\n",
    "        - normalisation: Normalization applied to the STA. One of [\"abs\", \"L2\", None]\n",
    "        \n",
    "    return:\n",
    "        - stas of shape (n_cell, Hw+Fw, ...)\n",
    "        - stas and pvalues if return_pval=True, both of shape (n_cell, Hw+Fw, ...)\n",
    "    \"\"\"\n",
    "    assert normalisation in [\"abs\", \"L2\", None], \"normalisation must be one of ['abs', 'L2', None]\"\n",
    "    #Preparing the stimulus\n",
    "    orig_shape = stim_inten.shape\n",
    "    stim_inten = stim_inten_norm(stim_inten)\n",
    "    sum_spikes = np.sum(spike_counts, axis=0)\n",
    "    len_stim = len(stim_inten)\n",
    "    \n",
    "    #We just have to calculate one STA over the whole record\n",
    "    stim_inten   = np.reshape(stim_inten, (len(stim_inten),-1))\n",
    "    stim_inten   = np.transpose(stim_inten)\n",
    "    allCells_sta = staEst_fromBins(stim_inten, spike_counts, Hw, Fw=Fw)\n",
    "\n",
    "    if len(orig_shape)==3:\n",
    "        allCells_sta = allCells_sta.reshape((len(allCells_sta),Hw+Fw, orig_shape[-2], orig_shape[-1]))\n",
    "    elif len(orig_shape)==2:\n",
    "        allCells_sta = allCells_sta.reshape((len(allCells_sta),Hw+Fw,-1))\n",
    "    else:\n",
    "        allCells_sta = allCells_sta.reshape((len(allCells_sta),Hw+Fw))\n",
    "        \n",
    "    if allCells_sta.shape[0]==1: #Only one cell, but we need to keep the axis\n",
    "        allCells_sta = np.squeeze(allCells_sta)\n",
    "        allCells_sta = np.expand_dims(allCells_sta, axis=0)\n",
    "    else:\n",
    "        allCells_sta = np.squeeze(allCells_sta)\n",
    "    \n",
    "    if return_pval:\n",
    "        p_values = np.empty(allCells_sta.shape)\n",
    "    for k, cell_sta in enumerate(allCells_sta): #Easy way to do normalization for each cell that works for all possible shapes\n",
    "        if return_pval:\n",
    "            z_scores    = cell_sta/ np.sqrt(1/sum_spikes[k]) #Standard score is calculated as (x-mean)/std\n",
    "            p_values[k] = sp.stats.norm.sf(abs(z_scores))*2\n",
    "            \n",
    "        if normalisation is None:\n",
    "            continue\n",
    "        elif normalisation==\"abs\":\n",
    "            allCells_sta[k] = np.nan_to_num(cell_sta/np.max(np.abs(cell_sta)))\n",
    "        elif normalisation==\"L2\":\n",
    "            allCells_sta[k] = np.nan_to_num(cell_sta/np.sqrt(np.sum(np.power(cell_sta, 2))))\n",
    "        \n",
    "    if return_pval:\n",
    "        return allCells_sta, p_values\n",
    "    else:\n",
    "        return allCells_sta\n",
    "    \n",
    "def staEst_fromBins(stim, spike_counts, Hw, Fw=0):\n",
    "    \"\"\"\n",
    "    Matrix mutliplication to compute the STA. Use the wrapper process_sta_batch to avoid bugs.\n",
    "    \n",
    "    params:\n",
    "        - stim_inten: stimulus intensity matrix of shape (flattened_frame, t)\n",
    "        - spike_counts: cells activity matrix of shape (t, n_cell)\n",
    "        - Hw: Lenght in frames of the history window, including the 0 timepoint\n",
    "        - Fw: Lenght in frames of the forward window\n",
    "        \n",
    "    return:\n",
    "        - STA of shape (n_cell, Hw+Fw, flattened_frame)\n",
    "    \"\"\"\n",
    "    spike_counts[:Hw] = 0\n",
    "    \n",
    "    spike_counts = np.nan_to_num(spike_counts / np.sum(spike_counts,axis=0))\n",
    "    spike_counts = spike_counts - np.mean(spike_counts, axis=0) #Center to 0 the spike counts to include \"inhibitory\" stimulus\n",
    "    sta = np.zeros((Hw+Fw, stim.shape[0], spike_counts.shape[-1]))\n",
    "    for i in range(Hw): #Does the dot product of cell activity and frame intensity matrix\n",
    "        sta[(Hw-1-i),:,:] = np.dot(stim, spike_counts)\n",
    "        spike_counts      = np.roll(spike_counts, -1, axis=0) #And shift the activity to compute the next frame\n",
    "    spike_counts = np.roll(spike_counts, Hw, axis=0)\n",
    "    if Fw != 0:\n",
    "        spike_counts[-Fw:] = 0\n",
    "    for i in range(Fw): #Same thing for Fw\n",
    "        spike_counts  = np.roll(spike_counts, 1, axis=0)\n",
    "        sta[Hw+i,:,:] = np.dot(stim, spike_counts)\n",
    "    spike_counts = np.roll(spike_counts, -Fw, axis=0)\n",
    "    return np.transpose(sta, (2,0,1))\n",
    "\n",
    "def process_sta_batch_large(stim_inten, spike_counts, Hw=30, Fw=2, return_pval=False, normalisation=\"abs\", bs=1000):\n",
    "    \"\"\"\n",
    "    Computes the STA and associated pvalues in parallel for a batch of cells, for a large stimulus.\n",
    "    \n",
    "    params:\n",
    "        - stim_inten: stimulus intensity matrix of shape (t, ...)\n",
    "        - spike_counts: cells activity matrix of shape (t, n_cell)\n",
    "        - Hw: Lenght in frames of the history window, including the 0 timepoint\n",
    "        - Fw: Lenght in frames of the forward window\n",
    "        - return_pval: Flag to signal whether or not to return the pvalues\n",
    "        - normalisation: Normalization applied to the STA. One of [\"abs\", \"L2\", None]\n",
    "        - bs: batch size to compute partial STA\n",
    "        \n",
    "    return:\n",
    "        - stas of shape (n_cell, Hw+Fw, ...)\n",
    "        - stas and pvalues if return_pval=True, both of shape (n_cell, Hw+Fw, ...)\n",
    "    \"\"\"\n",
    "    orig_shape = stim_inten.shape\n",
    "    n_spatial_dim = orig_shape[1]*orig_shape[2]\n",
    "    \n",
    "    sum_spikes = np.sum(spike_counts, axis=0)\n",
    "    len_stim = len(stim_inten)\n",
    "    allCells_sta = np.zeros((n_spatial_dim, spike_counts.shape[1], Hw+Fw))\n",
    "    stim_inten = stim_inten.reshape((len_stim,-1))\n",
    "    print(\"Computing the STA part by part:\")\n",
    "    for i, batch_pos in enumerate(range(0, n_spatial_dim, bs)):\n",
    "        print(str(round(100*batch_pos/n_spatial_dim,2))+\"%      \", end=\"\\n\", flush=True)\n",
    "        #Computing STA on partial portions of the screen sequentially\n",
    "        stim_part = stim_inten_norm(stim_inten[:, batch_pos:batch_pos+bs]).T\n",
    "        sub_sta = staEst_fromBins(stim_part, spike_counts, Hw, Fw=Fw)\n",
    "        allCells_sta[batch_pos:batch_pos+bs] = np.transpose(sub_sta, (2,0,1))#(ncell,Hw,stim_len) to (stim_len,ncell,Hw)\n",
    "    allCells_sta = np.transpose(allCells_sta, (1,2,0))\n",
    "    print(\"100%      \")\n",
    "    \n",
    "    if len(orig_shape)==3:\n",
    "        allCells_sta = allCells_sta.reshape((len(allCells_sta),Hw+Fw, orig_shape[-2], orig_shape[-1]))\n",
    "    elif len(orig_shape)==2:\n",
    "        allCells_sta = allCells_sta.reshape((len(allCells_sta),Hw+Fw,-1))\n",
    "    else:\n",
    "        allCells_sta = allCells_sta.reshape((len(allCells_sta),Hw+Fw))\n",
    "\n",
    "    if allCells_sta.shape[0]==1: #Only one cell, but we need to keep the axis\n",
    "        allCells_sta = np.squeeze(allCells_sta)\n",
    "        allCells_sta = np.expand_dims(allCells_sta, axis=0)\n",
    "    else:\n",
    "        allCells_sta = np.squeeze(allCells_sta)\n",
    "\n",
    "    if return_pval:\n",
    "        p_values = np.empty(allCells_sta.shape)\n",
    "    for k, cell_sta in enumerate(allCells_sta): #Easy way to do normalization for each cell that works for all possible shapes\n",
    "        if return_pval:\n",
    "            z_scores    = cell_sta/ np.sqrt(1/sum_spikes[k]) #Standard score is calculated as (x-mean)/std\n",
    "            p_values[k] = sp.stats.norm.sf(abs(z_scores))*2\n",
    "\n",
    "        if normalisation is None:\n",
    "            continue\n",
    "        elif normalisation==\"abs\":\n",
    "            allCells_sta[k] = np.nan_to_num(cell_sta/np.max(np.abs(cell_sta)))\n",
    "        elif normalisation==\"L2\":\n",
    "            allCells_sta[k] = np.nan_to_num(cell_sta/np.sqrt(np.sum(np.power(cell_sta, 2))))\n",
    "\n",
    "    if return_pval:\n",
    "        return allCells_sta, p_values\n",
    "    else:\n",
    "        return allCells_sta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def cross_correlation(spike_counts, tail_len=100):\n",
    "    \"\"\"\n",
    "    From calculate the cross correlation of the cells over a time window.\n",
    "    \n",
    "    params:\n",
    "        - spike_counts of shape (t, n_cell)\n",
    "        - tail_len: time correlation window size \n",
    "        \n",
    "    return:\n",
    "        - cross correlation between the cells\n",
    "    \"\"\"\n",
    "    n_dpoints, n_cell = spike_counts.shape\n",
    "    corr_arr = np.zeros((n_cell,n_cell,tail_len+1))\n",
    "    spike_counts = (spike_counts / np.max(spike_counts, axis=0)) #Independant normalization of the cells\n",
    "    spike_counts_edged = np.concatenate((np.zeros((tail_len,n_cell)), \n",
    "                    spike_counts, \n",
    "                    np.zeros((tail_len,n_cell)))) #Creating an array with 0 tails on both side to use the valid mode\n",
    "                                                  # of numpy.correlate\n",
    "    for i in range(n_cell):\n",
    "        for j in range(i, n_cell):\n",
    "            corr_arr[i,j] = np.correlate(spike_counts_edged[:,i],\n",
    "                                         spike_counts[:,j], \n",
    "                                         mode=\"valid\")\n",
    "            corr_arr[j,i] = corr_arr[i,j]\n",
    "    return corr_arr/n_dpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def corrcoef(spike_counts):\n",
    "    \"\"\"\n",
    "    Computes correlation coefficient between the cells\n",
    "    \n",
    "    params:\n",
    "        - spike_counts: Cells activity of shape (t, n_cell)\n",
    "        \n",
    "    return:\n",
    "        - Correlation matrix\n",
    "    \"\"\"\n",
    "    return np.corrcoef(spike_counts.T)\n",
    "\n",
    "def flatten_corrcoef(corrcoef_matrix):\n",
    "    \"\"\"\n",
    "    Retrieves a one dimensional array of the cells correlations\n",
    "    \n",
    "    params:\n",
    "        - corrcoef_matrix: Correlation matrix from `corrcoef`\n",
    "    \n",
    "    return:\n",
    "        - flattened correlation matrix\"\"\"\n",
    "    shp = corrcoef_matrix.shape\n",
    "    return np.array([corrcoef_matrix[i,j] for i in range(shp[0]) for j in range(i+1, shp[0])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def stimulus_ensemble(stim_inten, Hw=30, x=0, y=0, w=None, h=None):\n",
    "    \"\"\"\n",
    "    Generate the stimulus ensemble used to compute the nonlinearity\n",
    "\n",
    "    params:\n",
    "        - stim_inten: stimulus intensity matrix of shape (t, ...)\n",
    "        - Hw: Lenght in frames of the history window, including the 0 timepoint\n",
    "        - x: Left position of the window where to get the ensemble from\n",
    "        - y: Up position of the window where to get the ensemble from\n",
    "        - w: Width of the window where to get the ensemble from. If None, is set to stim_inten.shape[2]\n",
    "        - h: Height of the window where to get the ensemble from. If None, is set to stim_inten.shape[1]\n",
    "\n",
    "    return:\n",
    "        - Flatten stimulus ensemble of size (len(stim_inten)-(Hw-1), w*h*Hw). To obtain the corresponding cell activity,\n",
    "        slice it like so: slice(Hw-1, None)\n",
    "    \"\"\"\n",
    "    stim_inten = stim_inten_norm(stim_inten)\n",
    "    if len(stim_inten.shape) == 1:\n",
    "        stim_inten = stim_inten[..., np.newaxis, np.newaxis]\n",
    "    elif len(stim_inten.shape) == 2:\n",
    "        stim_inten = stim_inten[..., np.newaxis]\n",
    "\n",
    "    if w is None:\n",
    "        w = stim_inten.shape[2]\n",
    "    if h is None:\n",
    "        h = stim_inten.shape[1]\n",
    "    xmin, xmax = max(0,x-w), min(stim_inten.shape[2], x+w+1)\n",
    "    ymin, ymax = max(0,y-h), min(stim_inten.shape[1], y+h+1)\n",
    "    dtype = stim_inten.dtype\n",
    "    if np.all(np.in1d(stim_inten, [-1,0,1])):\n",
    "        dtype = \"int8\"\n",
    "    stim_ensmbl = np.zeros((len(stim_inten)-(Hw-1), (xmax-xmin)*(ymax-ymin)*Hw), dtype=dtype)\n",
    "    for i in range(0, len(stim_inten)-(Hw-1)):\n",
    "        flat_stim         = np.ndarray.flatten(stim_inten[i:i+Hw,\n",
    "                                                          ymin:ymax,\n",
    "                                                          xmin:xmax])\n",
    "        stim_ensmbl[i] = flat_stim\n",
    "    return stim_ensmbl\n",
    "\n",
    "def process_nonlinearity(stim_inten, spike_counts, bins, stas, p_norm=2):\n",
    "    \"\"\"\n",
    "    Computes the nonlinearity of a single cell. The STA of the cell is in L2 normalization, which\n",
    "    should restrict the histogram values.\n",
    "\n",
    "    params:\n",
    "        - stim_inten: stimulus intensity in shape (t, y, x)\n",
    "        - spike_counts: cells activity in shape (t, n_cell)\n",
    "        - bins: bins in which the transformed stimuli ensembles are set. (usually between -6 and 6)\n",
    "        - stas:  The STAs to convolve with stim_inten in shape (n_cell, Hw, ...)\n",
    "        - p_norm: Power for the normalization. https://en.wikipedia.org/wiki/Norm_(mathematics)#p-norm :\n",
    "            1 -> can compare nonlinearites of stimuli with different dimensionality\n",
    "            2 -> Common L2 normalization for STAs\n",
    "\n",
    "    return:\n",
    "        - nonlinearity of the cell.\n",
    "    \"\"\"\n",
    "    assert len(stim_inten)==len(spike_counts)\n",
    "    stim_inten    = stim_inten_norm(np.array(stim_inten))\n",
    "    \n",
    "    nonlins = np.empty((len(stas), len(bins)-1))\n",
    "    \n",
    "    compute_with_dotprod = False\n",
    "    try:\n",
    "        stim_ensemble = stimulus_ensemble(stim_inten, Hw=stas.shape[1])\n",
    "        compute_with_dotprod = True\n",
    "    except MemoryError as err:\n",
    "        print(\"Not enough memory to generate the stimulus ensemble, \"+\n",
    "              \"computing nonlinearity with cross-correlation (will be slower)\")\n",
    "\n",
    "    #In the case of calcium imaging, we have probabilities, not spike counts. Need to make it integers\n",
    "    # The discretisation of the calcium imaging is done here globally (for all cells together)\n",
    "    # If it's not what you want, either do the discretisation inside the loop bellow, or discretise the S_matrix\n",
    "    # before passing it to this function\n",
    "    \n",
    "    if np.max(spike_counts)<1:\n",
    "        mask         = np.where(spike_counts > 0)\n",
    "        nonzero_min  = np.min(spike_counts[mask])\n",
    "        discretized  = spike_counts/nonzero_min\n",
    "        spike_counts = ((10*discretized)/(np.max(discretized))).astype(int)\n",
    "\n",
    "    spike_counts  = np.array(spike_counts)[stas.shape[1]-1:].astype(int)\n",
    "\n",
    "    for i, (sta, sp_count) in enumerate(zip(stas, spike_counts.T)):\n",
    "        sta /= np.power(np.sum(np.power(np.abs(sta), p_norm)), 1/p_norm) # p-norm\n",
    "\n",
    "        if not sp_count.any(): #No spikes\n",
    "            continue\n",
    "\n",
    "        if compute_with_dotprod:\n",
    "            #This one is faster, but requires the stim_ensemble to fit the computer memory\n",
    "            filtered_stim = stim_ensemble@sta.reshape(-1)\n",
    "        else:\n",
    "            filtered_stim = np.squeeze(sp.signal.correlate(stim_inten, sta, mode=\"valid\"))\n",
    "\n",
    "        filtered_sptrigg = np.repeat(filtered_stim, sp_count)\n",
    "\n",
    "        hist_all   = np.histogram(filtered_stim, bins=bins)[0]\n",
    "        hist_trigg = np.histogram(filtered_sptrigg, bins=bins)[0]\n",
    "\n",
    "        nonlin = hist_trigg/hist_all\n",
    "\n",
    "        if np.count_nonzero(~np.isnan(nonlin))<2: #Less than two values in the nonlin, cannot fill the gaps\n",
    "            nonlin = np.nan_to_num(nonlin)\n",
    "        else:\n",
    "            nonlin = np.nan_to_num(fill_nan(nonlin))\n",
    "\n",
    "        nonlins[i] = nonlin\n",
    "\n",
    "    return nonlins\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def activity_histogram(spike_counts):\n",
    "    \"\"\"\n",
    "    Retrieve an histogram of the individual cells activity.\n",
    "    \n",
    "    params:\n",
    "        - spike_counts: cells activity matrix of shape (t, n_cell)\n",
    "        \n",
    "    return:\n",
    "        - Cells activity histogram\n",
    "    \"\"\"\n",
    "    flat_spikes = spike_counts.reshape(-1)\n",
    "    flat_cell = np.array([[i]*spike_counts.shape[0] for i in range(spike_counts.shape[1])]).reshape(-1)\n",
    "    hist = np.histogram2d(flat_spikes, flat_cell, bins=[100,spike_counts.shape[1]])[0] / spike_counts.shape[0]\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def cross_distances(masks):\n",
    "    \"\"\"\n",
    "    Computes cross distances from the center of mass of a list of mask. \n",
    "    \n",
    "    params:\n",
    "        - masks: cells mask of shape (n_mask, y, x)\n",
    "        \n",
    "    return:\n",
    "        - cross distances matrix\n",
    "    \"\"\"\n",
    "    center_mass = np.array([ndimage.measurements.center_of_mass(mask) for mask in masks])\n",
    "    cross_distances = np.zeros((len(masks),len(masks)))\n",
    "    for i in range(len(masks)):\n",
    "        for j in range(i,len(masks)):\n",
    "            cross_distances[i,j] = np.linalg.norm(center_mass[i]-center_mass[j])\n",
    "            cross_distances[j,i] = cross_distances[i,j]\n",
    "    return cross_distances\n",
    "\n",
    "def cross_distances_sta(fits, sta_shape, f):\n",
    "    \"\"\"\n",
    "    Computes cross distances between STAs \n",
    "    \n",
    "    params:\n",
    "        - fits: fits of STA to compare\n",
    "        - sta_shape: shape of the STA\n",
    "        - f: function used to compute the stas from the fits parameters\n",
    "        \n",
    "    return:\n",
    "        - cross distances matrix\n",
    "    \"\"\"\n",
    "    sta_masks = np.array([img_2d_fit(sta_shape, fit, f) for fit in fits])\n",
    "    for i,sta_mask in enumerate(sta_masks):\n",
    "        if abs(np.min(sta_mask)) > np.max(sta_mask):\n",
    "            sta_masks[i] = sta_mask < -.5\n",
    "        else:\n",
    "            sta_masks[i] = sta_mask > .5\n",
    "    return cross_distances(sta_masks)\n",
    "\n",
    "def paired_distances(masks_1, masks_2):\n",
    "    \"\"\"\n",
    "    Compute the distance between two masks\n",
    "    \n",
    "    params:\n",
    "        - masks_1: Mask of the first cell\n",
    "        - masks_2: Mask of the second cell\n",
    "        \n",
    "    return:\n",
    "        - distance between the two masks\n",
    "    \"\"\"\n",
    "    center_mass_1 = np.array([ndimage.measurements.center_of_mass(mask) for mask in masks_1])\n",
    "    center_mass_2 = np.array([ndimage.measurements.center_of_mass(mask) for mask in masks_2])\n",
    "    paired_distances = np.zeros(len(masks_1))\n",
    "    for i, (center_1, center_2) in enumerate(zip(masks_1, masks_2)):\n",
    "        paired_distances[i] = np.linalg.norm(center_1-center_2)\n",
    "    return paired_distances\n",
    "\n",
    "def paired_distances_sta(sta_fits_1, sta_fits_2, sta_shape, f):\n",
    "    \"\"\"\n",
    "    Compute the distance between two STAs.\n",
    "    \n",
    "    params:\n",
    "        - sta_fits_1: fits of STA of the first cell\n",
    "        - sta_fits_2: fits of STA of the second cell\n",
    "        - sta_shape: shape of the STA\n",
    "        - f: function used to compute the stas from the fits parameters\n",
    "        \n",
    "    return:\n",
    "        - distance between the two STAs\n",
    "    \"\"\"\n",
    "    sta_masks_1 = np.array([img_2d_fit(sta_shape, fit, f) for fit in sta_fits_1])\n",
    "    for i,sta_mask in enumerate(sta_masks_1):\n",
    "        if abs(np.min(sta_mask)) > np.max(sta_mask):\n",
    "            sta_masks_1[i] = sta_mask < -.5\n",
    "        else:\n",
    "            sta_masks_1[i] = sta_mask > .5\n",
    "    sta_masks_2 = np.array([img_2d_fit(sta_shape, fit, f) for fit in sta_fits_2])\n",
    "    for i,sta_mask in enumerate(sta_masks_2):\n",
    "        if abs(np.min(sta_mask)) > np.max(sta_mask):\n",
    "            sta_masks_2[i] = sta_mask < -.5\n",
    "        else:\n",
    "            sta_masks_2[i] = sta_mask > .5\n",
    "    return paired_distances(sta_masks_1, sta_masks_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direction and orientation selectivity with imaginary numbers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have 8 angles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.         0.78539816 1.57079633 2.35619449 3.14159265 3.92699082\n",
      " 4.71238898 5.49778714]\n"
     ]
    }
   ],
   "source": [
    "n_angle = 8\n",
    "x = np.linspace(0, (n_angle-1)/4*np.pi, num=n_angle)\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exponentiating those and multiplying with 0+1i will wind up the angles around the origin in imag number space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.00000000e+00+0.00000000e+00j  7.07106781e-01+7.07106781e-01j\n",
      "  6.12323400e-17+1.00000000e+00j -7.07106781e-01+7.07106781e-01j\n",
      " -1.00000000e+00+1.22464680e-16j -7.07106781e-01-7.07106781e-01j\n",
      " -1.83697020e-16-1.00000000e+00j  7.07106781e-01-7.07106781e-01j]\n"
     ]
    }
   ],
   "source": [
    "vectors = np.exp(x*1j)\n",
    "print(vectors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And so multiplying by two before exponentiation double the angle of the arrow. It gives to originally opposite vectors the same direction -> good to calculate orientation selectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'OS vectors')"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXwAAAEICAYAAABcVE8dAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAgAElEQVR4nO3deZgcZdX38e/pZZZM1slC9gUIgYS4kAAmikZZwg4PiIYHQYX3jaAouILiyiuYR4LK4iOiLBqQBGSXBAiIAkIwYZPEJCRCIJMEsq+z9XK/f1Ql6Zn0LJnp6eru+n1y1ZXuqpqq0zOnT1ffXafanHOIiEjpiwQdgIiI5IcKvohISKjgi4iEhAq+iEhIqOCLiISECr6ISEio4IuIhIQKfg6Z2SozqzOzHWa21cxeMLOLzSySsc5QM7vfzDaa2TYze8PMvpDHGKeYWU2+9ielw8y+4OdrrZm9Z2a/MbPeGct7m9nt/rIdZvammV2Rx/hGmpkzs1i+9llsVPBz7zTnXA9gBDADuAK4LWP5LGC1v7wvcAHwfr6D7Cg9mcLJzL4J/A/wbaAX8BG8HJ5vZmX+ar8EugOH+eucDvwn/9F2TChy2zmnKUcTsAo4rtm8o4A0cLh/fyfwoXZubylwasb9GLAROMK//xHgBWAr8DowJWPdauAOYC2wBXgIqALq/Hh2+tNgoBz4lb/uWv92ub+dKUAN3gvXe3gvWP2Av/j73Qw8B0SC/v1r6poJ6Onnymeaze8OrAcu9O8vBs5s5zYfBy5tNu914Cz/9qHAfD+/lmfuG6gErgfeAbYBz/vz3gVcRm5Pwjuo/b6/7nrgj0Avfzsj/fUv8n/2WaACuAvY5Of3QuCAoP8GOftbBh1AKU3ZCr4//13gEv/2U8A/gGnA8Da290Pg7oz7pwDL/NtD/KQ82U/q4/37/f3ljwFzgD5AHPiEP38KUNNsP1cDC4ABQH+8F5H/l7F+Eu/ortx/Yv0MuMXfbhw4BrCgf/+aumYCTvRzIJZl2R+Ae/zbvweWAF8ERrexzQuAf2TcH+sX2HK8A5PV/nZiwBF4Bzrj/HV/DfzNfw5Egcn+z+0u4LGM7V4IrAQOxHuBegCY5S/bvf4f/X1WAl8CHgW6+dueAPQM+m+Qs79l0AGU0tRKwV8AXOXf7oM31LMESAGvAUe2sL2DgR1AN//+3cAP/dtX7E7cjPWfAD4PDMI7iu+TZZvZCv5/gJMz7k8FVmWs3whUZCy/GngYODjo37mmrp+AzwHvtbBsBjDfv10JfA94GUj4hfakFn6uB7ALGOHfvwa43b/9WeC5Zuv/FvgR3sFNHfDBLNvMVvCfBr6ccX+MH1ssY/0DM5ZfiHfA84Ggf+9dMWkMPz+G4L01xTm3xTl3pXNuHHAAXsF/yMys+Q8551biDeucZmbd8MZE/+QvHgGc4384vNXMtgIfwyv2w4DNzrkt7YxvMN5b3t3e8efttsE5V59x/zq8J/OTZvaWmV3Zzv1IcdoI9GthjHuQvxznXJ1z7lrn3AS8z6fuBe4zs+rmP+Sc24H3LnSaP2sa3gENeLl9dLPcPg8YiDecWEH7PxvIltsxvOfebqszbs/CO3CabWZrzeznZhZv574Kngp+FzOzI/EK/vPNlznnNgIz8ZJynyeF7x7gXOAM4N/+iwB4STrLOdc7Y6pyzs3wl1VnnkGRudss89biPcl2G+7Py/ozzrkdzrlvOucOBE4DvmFmx7YQvxS/F4EG4KzMmWZWBZyEdxTdhHNuO3At3lDJqBa2ew9wrplNwnt38Iw/fzXw92a53d05dwnei0s9cFCW7bU3t5M0PVFiz8855xLOuZ8458biDRWdijf8VBJU8LuImfU0s1OB2cBdzrk3/Pn/Y2aHm1nMzHoAlwArnXObWtjUbOAEf70/Zcy/C+/If6qZRc2swj/lcqhzbh0wD/hfM+tjZnEz+7j/c+8Dfc2sV8a27gG+b2b9zawf3mcHd7Xy2E41s4P9dyXb8YamUvv1C5Ki4ZzbBvwEuMnMTvTzaSRwH94H+rMAzOwHZnakmZWZWQVwGd64/PIWNj0XrxhfDcxxzqX9+X8BDjGz8/19xf3tHuavczvwCzMb7Of+JDMrBzbgDWUemLGPe4Cvm9koM+uO9yI0xzmXzBaQmX3SzMabWRQvtxOUUm4HPaZUShPeGH4d3rj7Nrwjo68A0Yx1bgJW4J1FsAEvuQ9rY7tP4x2VDGw2/2jg73jDRRvw3iIP95dV432g9j7eWToPZPzc7ew9C2Ew3lvkG4F1/nQj/pg92cf8v+4/1l14T/gfBP2719T1E97ZLIv9HH8fb1y9T8by7/vLt/s5+TdgchvbvA3vCPvIZvPH+Pm8wc/Vv+Kf3Yb3buBXwBr/efYsUOkvu9r/ma14Z7FF8A5gVvvz79odM9nH/M/Fe4Ha5T/GG8nyYXWxTuY/SBERKXEa0hERCYmcFHy/nXq9mS1uYfkU/zICr/nTD3OxX5GupLyWUpOrVuI7gZvxGhha8pxz7tQc7U8kH+5EeS0lJCdH+M65Z/HPMxcpFcprKTX5vFjQJDN7He+82G8555ZkW8nMpgPTAaqqqiYceuiheQxRwuTll1/e6Jzr38nNKK+loLSW1zk7S8c/L/cvzrnDsyzrCaSdczvN7GTgBufc6La2OXHiRLdo0aKcxCfSnJm97Jyb2MY6I1FeSxFpLa/zcpaOc267c26nf3suEPcbfESKlvJaik1eCr6ZDdx9rRgzO8rfb0udpSJFQXktxSYnY/hmdg9eR2Y//9uUfoR32Vycc7cAnwYuMbMkXpfeNKeOLylwymspNTkp+M65c9tYfjPe6W0iRUN5LaVGnbYiIiGhgi8iEhIq+CIiIaGCLyISEir4IiIhoYIvIhISKvgiIiGhgi8iEhIq+CIiIaGCLyISEir4IiIhoYIvIhISKvgiIiGhgi8iEhIq+CIiIaGCLyISEir4IiIhoYIvIhISKvgiIiGhgi8iEhIq+CIiIaGCLyISEir4IiIhoYIvIhISKvgiIiGhgi8iEhIq+CIiIZGTgm9mt5vZejNb3MJyM7MbzWylmf3LzI7IxX5FupLyWkpNro7w7wRObGX5ScBof5oO/CZH+xXpSneivJYSkpOC75x7FtjcyipnAH90ngVAbzMblIt9i3QV5bWUmnyN4Q8BVmfcr/HnSQ7UJeu47K+X8fqG14MOJWyU113t7Wfh3gsg2Rh0JCUhXwXfssxzWVc0m25mi8xs0YYNG7o4rNIw46UZPFvzLJc/czm1idqgwwkT5XVXqt3sFfvlc+HpHwcdTUnIV8GvAYZl3B8KrM22onPuVufcROfcxP79++cluGL2XM1zzH17LkmXZHvjdq5+8eqgQwoT5XVXcQ4eugQad0EqAQtvg3deCDqqopevgv8IcIF/VsNHgG3OuXV52nfJ2lK/hSueu4L6VD0AjalGnn73af767l8Djiw0lNdd5fXZ3nBOyh/KSdZ7R/v124ONq8jl6rTMe4AXgTFmVmNmF5nZxWZ2sb/KXOAtYCXwO+DLudhvmDnn+O5z36UuWddkfn2qnquev4qNdRsDiqx0KK8DsvVdeOyb0Hx4smE7PPq1YGIqEbFcbMQ5d24byx3wlVzsSzwPrnyQl99/mWQ6uc+y+mQ93/n7d7ht6m2YZRtmlvZQXgcgnYI5F3hH9M0lG+DNx2HJQzDuzPzHVgLUaVuEVu9YzYx/ztgzlNNc0iVZvHExs5fPznNkIp30/C9h43JwqezLE3Xw8Fdgu0bOOkIFvwjN+OcMGlONVMWqqIpVURGtACAeie+ZlybNzIUz9xnyESlYtZvhmWvBpaGsuzdF4t6yeOXeecl6mP+DYGMtUjkZ0pH8uuyIyzhl1Cl77v90wU+pT9VTHi3novEXMbhqMADdy7pTGasMKkyR/dOtGj47yzuKB9i4Ahb8LzQkvEJ/4oy96w78QDAxFjkV/CJ0SJ9DOKTPIXvuX7foOnYkdgBwzJBjGFM9JqjQRDrn0L0HMry7wCv4AOU9YPyng4mphGhIR0QkJFTwRURCQgVfRCQkVPBFREJCBV9EClYibWxu0JlmuaKzdESkYD357jCWb+3Nf3ffwMCggykBOsIXkYK04o1lrNzWE4fx8LJqEg3ZO8ul/VTwRaTg7Nq6hcfnPELSRQGoT0b46523BhxV8VPBF5GC4pzj0V/9D4nE3gsDJtMRlj3/d956dVGAkRU/FXwRKSivPvEX3n9rBS6dbjI/2djA3Juuo3b7toAiK34q+CJSMDavreG5u+8g2dCQdXmioYF5N1+Pd2Vq2V86S6eDtj74INsefTTrMrMIg352LfEBA/IcVZ49ejlsWZV9Wa9hcNoNENExhbRPKpnk4Zk/JdnY8heWp5NJ3l3yLxb/bT7jP3lCHqMrDSr4HbT1vj9T98or2RfGYrjaEv8yceegZiG8v4Ss39ttETj55xDROdTSPqlEIy6dpqpPtXe/oY5U/U4S6SgRc1T27rtn3a3rsn51sLRBBT/HrLKSAd/8BmUjRwYdStcyg8/eBb+ZvO9X0e1eLrIfyiq7ceGv9p6Js+bv9/HALbcD0LM8xUW3/DGo0EqG3m/nUixG5eGH0+e884KOJD+qR3nXKI93CzoSEWkHFfwcipSXM+QX14fre2SPuACGH733m4lEpGCp4OeQS6VI7dgZdBj5lU7Bzg0tfwepiBQMFfwccg0N1Fx6KS6RCDqU/Hn2Otj8lvc9pCJS0FTwcyESwcrLwTkSa9ey4YYbg44oP9a8Av+4Ye+HthaBWEWwMYlIi1TwO8nKyxl+2++J9u4NgKuvZ/OsWdS+8mrAkXWxxlqYcx4k/S+cjlXA+Q9CZXWwcYlIi1TwOyMSofpzn6Nq0iSG3nyTd5SPP7Tzta+R2rkr4AC70LxvQ+1m73a8Eo76Ehw4Babd7RX/tIZ4RAqNCn5HmVE2ciT9L78MgMrx4+l70UVYpddolN6+nfd+/KMgI+w6bz4Jb9wPyXrAvK7aY3/gLRtyBEz+GlmbsUQkUCr4HTTo2msY/vvfYfG9pyP2+/IllA0fBpEIrrGRHU89zfb58wOMsgvs2gQP/J+9QznxCpj2J4hmnJb5iSvgi/M0ni9SYFTwO6h8xAjigwc3mWexGENvyhjaqa9n3Xe/R3LDhiBCzD3nvGLf6H9IG+8Gx/0E+o1uul40BiMmq9tWpMDkpOCb2YlmttzMVprZlVmWTzGzbWb2mj/9MBf7LURlw4dzwHev3Du0U1/Pmm98szSu7vfqLHh3AaQTEInB4CPgqOlBR9WllNtSSjpd8M0sCvwaOAkYC5xrZmOzrPqcc+5D/nR1Z/dbyHqfcw7djjgCYjFIJqlbvJgtd98ddFids/ltmHfF3lMw45Xw6dtL+iheuS2lJhdH+EcBK51zbznnGoHZwBk52G7RMjMGX/dzIv5RvqurY/3M62l46+2AI+ugVBLmfM7/kBav2J/5G+hxQLBxdT3ltpSUXBT8IcDqjPs1/rzmJpnZ62Y2z8zGtbQxM5tuZovMbNGGIh77jlVXM/i6n2MV3geXRd2Fm9lNGy2HQ0+Fw04LOqp8yFlul0peS3HLRcHP9p6++YD1K8AI59wHgZuAh1ramHPuVufcROfcxP79++cgvOD0mDKFniefVNxduGtegX/8au9QTmVvOPWXwcaUPznL7VLKayleuSj4NcCwjPtDgSbfTuCc2+6c2+nfngvEzaxfDvZd8Ab+4AdNu3DvuovaV4ukC7extulQTqzCuwZ+eY9g48of5baUlFwU/IXAaDMbZWZlwDTgkcwVzGyg+dcMNrOj/P1uysG+C16ksrJpF259PTVfu6w4unDnfQdq/T/T7m7aYUcFG1N+KbelpHS64DvnksClwBPAUuBe59wSM7vYzC72V/s0sNjMXgduBKa5kjhPsX326cLdtq3wu3DffBLe+HP2btqQUG5LqcnJVxz6b2XnNpt3S8btm4Gbc7GvYtXvy5ew4+mnaFixskkXbs/jjw86tH21p5s2JJTbUkrUaZsnRdOF295uWhEpOir4eVQUXbgh7KYVCQsV/DzL2oX7p3uCDssTwm5akTBRwc+zrF24110XfBduKglzzg9jN61IaKjgB6Agu3CfvQ42/2dvN+2YU8LSTSsSGir4ASmoLtzm3bQVveC0XwUTi4h0GRX8ABVEF+4+3bSVYeumFQkNFfwAtdSFm96Vxy7cfbppp8Pwo/O3fxHJGxX8gFWOH0/1hRc2+S7cdT/+SX52vmI+LA53N61ImKjgF4D+X/kyZcP878JtaGDH/Plsf+qprt3prk1w/0WQUDetSFio4BcAi8X2GdpZd+V3u64LV920IqGkgl8g8tqFu0837YfVTSsSAir4BSQvXbhZu2nvUDetSAio4BeQFrtw385RF246Bfeqm1YkrFTwC0z2Ltyv5qYL99nrYJO6aUXCSgW/AO3ThbtmDRtuvKlzG137KjyvblqRMFPBL1D7dOHOmtXxLtzGWph93t4vNAnfd9OKCCr4BWtPF+7uoZ3OdOFmdtPG1E0rElYq+AWscvx4qr/4xc514Tbvpu09DI79Ye6DFZGCp4Jf4LJ24T79dPt+WN20IpJBBb/AZe3CveJKkhs3tv6DzsED/1fdtCKyhwp+EehQF+6rd8G7L6qbVkT2UMFvZs2KLWzbUBt0GPvYpwv3jTda7sLd/Lb3QW2BfzdtXWOKuf9aV1hf4l6i0i7N428/TmOqsc11a5dsJF2fzENUxef9VdvZvK7tEydWrt/Bq+9uyUNE+0cFP8P2jXU8esNrLH1hXdCh7KPdXbguvW837Rn/Cz0G5jnitq3atIsv/+kV7l7wTtChlLzbF9/Ot5/9Nutr17e6Xu2SjWyetZTGd3fkKbLi8vx9K3j0xtdINqZaXe+Of6zivN+/xPrt9XmKrH1U8H3ptGPeb98glSzco81YdTWDf75vF24klRHzy3fu20079vRgAm6HeNS4Zu4y3t6Yxy99CZnlm5dzy+u3EIvEWl0vtaORLfe+CdHCeidYaHZtbeC5e1e0uV5dY4pL73m1oN7BquD7Xp63iq3vF95QTnM9PrlvF+6pf/WPxtIp70qYRdRNWx6LUp9MMf2Pi0ik0kGHU3IaUg1c9sxlNKQaKIuUtbiec45N9yzDJVJYXGWhNc7Bmy+9xztLWj9xwgFv1GzjDy+syktc7aG/LLDh3R288vg7JBuLo+A078L91Iu1jF7jvEKf9Mdoi6ib1jmo2VLHL+e/GXQoJWfmwplsrGvjjC5g14vrSKzeAcXxFAhcMpFm/m3/pn5n69e4qkukmPH4clau35mnyFqXk4JvZiea2XIzW2lmV2ZZbmZ2o7/8X2Z2RC72mwvJxhRzf/MvkoniyfTmXbhlSfjW/SnKd79gFWE3bV0ixe3Pv11wH3QVc26/tO4lHlr5EA2phlbXS2yoZdu8t3FF9BwoBImGFPPvWNLmkE1DIsWXZi2iMRn879c6O75kZlHgTeB4oAZYCJzrnPt3xjonA18FTgaOBm5wzrVZjSZOnOgWLVrUqfja8sxdy1j+0nukMpI9Fo8QK4926X5zIV1Xh6uvJ53aRiSdoNHdzXGH/ovyRD8+Y78gaa2P2QYtlU6TTDtqMz4AG9CjnGe+NYWq8q6P3cxeds5NbGV5l+R2PvJ6W8M2Tn3wVLY2bN0zrzJWSTwSJ2J7j/OiLsKMZV9lcEN/Iv7xX12kAYcjaa1/MNnVzEE8YazZtYKXt7xAus9FgcYDUJ6GCHs/40iY44WejhVVTetobWOShowCXxmPcP6kkXzv5MO6PMbW8joXz6qjgJXOubf8nc0GzgD+nbHOGcAfnffqssDMepvZIOdcoKfDrF66mTebFXvw3q4VxxF/DBeN0iP5T3YO6MfG7sfwdqqWmQ1fZINLA22fgldottUl+NEjS5h5zgeDDgWKOLd//MKP2ZVo+kF4XbKOOuqazDt//an0a+y9p9gDVKbL8xJjW9I4VpStY2eknGRdJVXpwvswOe6MSdtgqWtkR6Tlg+e6RJo/vriKkw4fyIeH98lfgM3kYkhnCLA6436NP29/1wHAzKab2SIzW7Shq77T1Ve/MwGFl0Pt5pyjcdc8Dj7gY2yu8h7IS6kTWOMGBBxZx6XSjo07Wh+CyKOc5XY+8xpgU/0mUq7tI/Q+yZ5YgT4J1ts2/hFbxusVa6kfPDLocFpkQLydAyXbanPwvRadkIuCny1bmj/89qzjzXTuVufcROfcxP79+3c6uNYcPHEAgw7qRaTZaWgWKY4plViMS9ZQZ4k9v+A0xpT4f4gaRM0KfupW1nTorCIe5efnfKBL/+77IWe5nc+8BphxzAzKo02P1CtjlUQt2mS6feBD1Eaanitebw2kCuBfmjRRIqTNgcVJ4wKfUs3+tAkciyqSbIs1f741/XuUxSJ86tABTDk02IOxXAzp1ADDMu4PBdZ2YJ28MzOOv3Acd/1wAY11ezsLP3TccI46bVSAkbVt2/r3uOu7vwaSDE9V83Z6M+siW8AcwyrqmXVsP46Y0OLwdEFY/t4Ozvv9S3vuV8QjXH/OBxjQoyLAqJoo2twe3H0wVx19Fde8dA11/vcgGMYDZzzAkO5N34Ak3tnB1tuXgd+DUlnejR7nHEjZQb3yHnemXc88A8/XQzxOt8RWvviLjwcaD8DcG19jwyr/NGiDAYO6c+t3Juxz0PijR5YwZ+HeN349ymPMODv4A5lcFPyFwGgzGwWsAaYB/91snUeAS/0x0KOBbUGPce5W2aOM4y8cyxO3Lt4zbh+JGrF44X5om06neOyGGaQavbeHZsaUxDjuK3+RRpIkEgmeefopDjn4IPr16xdwtC0ri+19g1keizB17ECmHj4owIj2UdS5ffpBp/PkO0/y4toXSaS9XCmLlO1z5F9+YDnJyUPY+eJa8J8DZfFyKioq8x5zplgs86quju7dgr/KazS6N2dj8QinXDKebhX7ltFYZO8LQEUswm8+N4GeFcHH3+khHedcErgUeAJYCtzrnFtiZheb2cX+anOBt4CVwO+AL3d2v7k0cnw/Rh95ANEiaThZ8MActqxbi3NpovE4sXgZlZQxhcOJRb3kSyQSzJkzh1Qq2DMt2qtnRZxrzhofdBhNFHtumxnXfuxaKmNtF+5eU0cQ610YH9YWg1hZhI995hB69e/W6noV8QjnfWQER42qzlNkrctJhXPOzXXOHeKcO8g5d40/7xbn3C3+beec+4q/fLxzrmvPSeuAY6YdQkVV8K/AbXn/rZUsfPjPJBu9DzbLu1VR3q0KgBEM4LCDxhCLeUV/69atPPPMM4HF2h6JVJqKeIRbzp9A9zycirm/ij23e5X34rpPXEdFtIJkuuULolk0Qt/zx2LxCK4AzhcvaAaDDurF2I+2/W50UK9Krjjx0DwE1T7FcUibB/GyKCdfMr6gz9pJNNTz8Myfkmz0TreMlZVx+je+1yTmqR87jkr/AmuJRIIFCxawevXqbJsrCA3JNF+YPJIJI4I7Va3UTR48mdMOOo3GdOun6cYHdKPniSP3jOVLdmUVUY6/aBzWxtVnoxHj1vMnNBm6DFrhHVIFaMCInnzq/EPpM6gq6FCyeubOW6nbvh2AWFk5H5p6CkMOHctaFuxZpyxexrRp07jjjjtIJpMkk0nmzJnDV7/6VcrLC+st++BelVx+7Gi+8qmDgw6l5H3nyO9QXVFNdUXrQwvdJw/GNaSID2x9qCKsPvDJoVT2iFPZveXrEgGc+oFBHH1gNaMPKKxLmxTOS0+BOGzyYAaOCvbshGzeenURS5//G8lEI2D07Nefj027IOu6Q4YMYfLkycTj3hBVfX09jz32WB6jbZ9e3eJcfvwhxKNKw65WEavg0g9fSrd464XczOj5qeFEexbWwUGhGD3xAIaOaXs8ftJB/Tj9g1lbjQKlZ1oRqN2+jbk3Xdd0KOdbVxGNtfwG7ROf+AR9+vTBzEgmkyxdupRly5blK2QRKUAq+AXOOce8m68n0eB9SBsvL+eYcz9P3yHDWv25aDTKtGnT9nyAm0gkePDBB9m5szCu2ici+aeCX+AW/20+NcuWkE4msUiUAaMO5sMnndaun62urmbq1Kl7hnYSiQR//vOfC+oLGUQkf1TwC9i29e/x1zt+S3LP0X0Zp15+RZtnB2SaMGECw4YNIxKJkE6nWbNmDQsXLuyqkEWkgKngF6h0OsXDM6/Z000bKy9n6iVfp3uf/WvgMDPOPvvsJkf58+fPZ+PGtr8UQ0RKiwp+gWreTXvQhKM55OjJHdpWVVUVZ511VpPx/GLqwhWR3FDBL0DZumlPmH5pp7Y5ZswYxo0bV1RduCKSWyr4Baalbtqyys43wpxyyilF1YUrIrmlgl9g9ummPcHrps2FsrIyPvvZz+45yt/dhdvQUDBfOCIiXUgFv4B43bR/b9pNe272btqOGjp0aFF04YpI7qngF4i93bTe0fbebtrcX8FTXbgi4aSCXwA62k3bUerCFQknFfwC0Jlu2o5SF65I+KjgBywX3bQdNWHCBIYOHaouXJGQUMEPUPZu2sv3u5u2o9SFKxIuKvgBWvDAvWx5r3k37UfzGkP37t3VhSsSEir4AfG6ae/bM5STi27ajlIXrkg4qOAHoCu7aTsqWxduTU1NYPGISO6p4AegK7tpOypbF+7s2bPVhStSQlTw8+zt115u0k3bo1+/nHfTdtTQoUOZNGmSunBFSpQKfh7Vbt/GYzf+vEk37Rnf/H6XdNN21JQpU9SFK1KiVPDzpHk3bWx3N+3Qrumm7ahoNNpkaEdduCKlQwU/T5p30x6Qh27ajurbt6+6cEVKkAp+HgTZTdtR2bpwFy1aFHRYItIJnSr4ZlZtZvPNbIX/f58W1ltlZm+Y2WtmFqqqEXQ3bUdl68J98skn2bRpU8CR5YdyW0pRZ4/wrwSeds6NBp7277fkk865DznnJnZyn0Vln27aI47KezdtR+3uws0s+rNnzw5LF65yW0pOZwv+GcAf/Nt/AM7s5PZKSrZu2uOnfzXgqPbPmDFjGDt2bBi7cJXbUnI6W/APcM6tA/D/H9DCeg540sxeNrPprW3QzKab2SIzW7Rhw4ZOhtKnm3gAAApaSURBVBecbN20p33ju5R3C66btqNC2oWb09wulbyW4tZmwTezp8xscZbpjP3Yz0edc0cAJwFfMbOPt7Sic+5W59xE59zE/v3778cuCku2btqhh44LOKqOKeEu3EPyldulktdS3Nos+M6545xzh2eZHgbeN7NBAP7/61vYxlr///XAg8BRuXsIhaeQu2k7qkS7cN9UbkuYdHZI5xHg8/7tzwMPN1/BzKrMrMfu28AJwOJO7rdg1e3YzmM3XlfQ3bQdFbIuXOW2lJzOFvwZwPFmtgI43r+PmQ02s7n+OgcAz5vZ68A/gcecc493cr+BSzQ20FBb22Sec465N80k0VAPeKdgfuzcCwqum7aj2tuF65wrhc7c0Oa2lK5OFXzn3Cbn3LHOudH+/5v9+Wudcyf7t99yzn3Qn8Y5567JReBBe/BnP+bhmT9t0n26bzftQRxx0ukBRpl7ffv25YQTTmi1C/fVV19l5syZJBKJoMLstDDntpQuddp2UCqVYvW/3+D1+fOAFrppLyvsbtqOmjhxYotduFu2bGHevHkBRygi2ajgd4Zz/P2u29i0ZnX2btrqvgEH2DVa6sLdsGEDc+bMIZlMEokotUQKjZ6VnZRsbOTu7329aLtpOypbF+7vfvc7Nm3apIusiRQoFfzOco5kQ0NRd9N2VPMu3EQiUdTj9iKlTgU/B3Yf0RZzN21HZXbh6shepLCp4OdQLF7GgJEHBh1GXpWVlTFq1Cii0WjQoYhIG1TwcyjR2MDTt/0m6DDyasWKFSxdujQsV9AUKWoq+DmUSiR4c8E/WLlwQdCh5EVtbS3333+/xu1FikQs6ABKTbKxgbk3X89FN9xKVe+s35lREpxzPPDAAzT6VwPNtlxkfySTSe677749F+Wr27KF3VlUV1bGnXfeuWfdcePGceSRR+Y/yCKngt9BI8Z/kO0bsl5PC4sY9bt2lnTBB2hsbKR79+5Zl/Xs2VPn4st+SaVS1NTUsGvXrr0z/dN+U9Eoq1atAiASiTBixIgAIix+KvgdNPmc85h8znlBhxEYM+PCCy8MOgwpIeXl5XzmM59h1qxZJJPJrOuYGf369ePjH2/xCuvSCh2CiUjBGDFiBEceeeSehr7mYrEY06ZN01lhHaSCLyIF5dhjj6Vnz577zI/H40ydOpXq6uoAoioNKvgiUlB2H8XHMj4DsnSaYcOGMWHChAAjK34q+CJScPr3788xhx1G1B/Lj6XTnH322SV59dl8UsEXkYI08cADqd6+HZzjoytWUFVVFXRIRU9n6YhIQTIzPvryK2wqK2NoCx/iyv7REb6IFKzKxkYGr1sXdBglQwVfRCQkVPBFREJCBV9EJCRU8EVEQkJn6RSh1M5GUtszrlKZ2ntlyuSGWvCvVGnxCPH+4fn2LSl+Df/5D86/Amvjqnf2zHeNjdQvXbrnfmzgQGJ9SvvihF1BBb8Ibb7vTRpWbMXi/hs0v8C7tGPL/Su8eWmHSzsG/2gSkTJdd0QKX2rrVt467XSsvBzzr5Xj/Mar1PbtvHP+BQCk6+vpecIJDPnF9YHFWqw0pFOEeh03AqKGa0h5U2PaW5BI750HdP/oEBV7KRrR3r3pffbZkE6T3rmT9M6duPp6AFxt7Z55FovR75KLA462OKngF6GyYT3occyQvUf4WcR6l9Nrqq4ZLsXlgKu+R7SVi6NZZSX9v/Y1ykePzmNUpUMFv0j1PHYE0eoKyHJpEYtH6Hv+WCyqP68Ul0hFBUNvvgkrL993YTRKxZgxVH/h8/kPrER0qiKY2TlmtsTM0mY2sZX1TjSz5Wa20syu7Mw+xWNRo98FY7FY0z+hxSP0PHEk8QH6sLYzlNvBqRw3jr5fmo5VVjaZH6moYMgNv8L0TWod1tnf3GLgLODZllYwsyjwa+AkYCxwrpmN7eR+BYj1raTXKQfuHdqJQHxId7pPHhxsYKVBuR2gfl/6EmUjR4Jf3K2igoFXX038gAOCDazIdargO+eWOueWt7HaUcBK59xbzrlGYDZwRmf2K3tVHT2QshE9wcDiUfr+92G6hGwOKLeDZdEoQ2+60RvaicXofswx9Drl5KDDKnr5eG80BFidcb/Gnyc5YGZUTxtDrF8lfc4ZTbRnWdAhhYlyuwuVDR3KwB/9kLJRIxl07TVBh1MS2jwP38yeAgZmWXSVc+7hduwj2+GmyzJv9/6mA9MBhg8f3o7NS7R7GQO/2eIws7TsEDNbnGV+znNbed0xvc88k95nnhl0GCWjzYLvnDuuk/uoAYZl3B8KrG1lf7cCtwJMnDixxRcGkRx40znXmVfKdue28loKQT6GdBYCo81slJmVAdOAR/KwX5GuptyWotLZ0zL/y8xqgEnAY2b2hD9/sJnNBXDOJYFLgSeApcC9zrklnQtbpGspt6UUdepaOs65B4EHs8xfC5yccX8uMLcz+xLJJ+W2lCJ1MIiIhIQKvohISKjgi4iEhAq+iEhIqOCLiISECr6ISEio4IuIhIQKvohISKjgi4iEhAq+iEhIqOCLiISECr6ISEio4IuIhIQKvohISKjgi4iEhAq+iEhIqOCLiISECr6ISEio4IuIhIQKvohISKjgi4iEhAq+iEhIqOCLiISECr6ISEio4IuIhIQKvohISKjgi4iEhAq+iEhIdKrgm9k5ZrbEzNJmNrGV9VaZ2Rtm9pqZLerMPkXyQbktpSjWyZ9fDJwF/LYd637SObexk/sTyRfltpScThV859xSADPLTTQiBUK5LaWos0f47eWAJ83MAb91zt3a0opmNh2Y7t/daWbLuzi2fkCpHJ2VymPJ1+MYkYNttCu3A8hrUD4Uonw8lhbzus2Cb2ZPAQOzLLrKOfdwOwP4qHNurZkNAOab2TLn3LPZVvSfMC2+IOSamS1yzrU4RltMSuWx5OtxmNlTZrY4y6Kc53a+8xqUD4Uo6MfSZsF3zh3X2Z0459b6/683sweBo4CsBV8kX5TbEjZdflqmmVWZWY/dt4ET8D4QEylqym0pNp09LfO/zKwGmAQ8ZmZP+PMHm9lcf7UDgOfN7HXgn8BjzrnHO7PfHMvr2+wuViqPJfDHodwuKKXyOCDgx2LOuSD3LyIieaJOWxGRkFDBFxEJCRV82t9GX6jM7EQzW25mK83syqDj6Sgzu93M1rdwqqTsp2LPa1Bu55oKvmd3G33RnU5nZlHg18BJwFjgXDMbG2xUHXYncGLQQZSQos1rUG53BRV8vDZ651w+Oh+7wlHASufcW865RmA2cEbAMXWI37C0Oeg4SkWR5zUot3NOBb/4DQFWZ9yv8eeJFDvldo7l61o6gcvRJSIKUbare+lc25Ao4bwG5XbOhabg56KNvkDVAMMy7g8F1gYUi+RZCec1KLdzTkM6xW8hMNrMRplZGTANeCTgmERyQbmdYyr4tNxGXwycc0ngUuAJYClwr3NuSbBRdYyZ3QO8CIwxsxozuyjomIpZMec1KLe7JA5dWkFEJBx0hC8iEhIq+CIiIaGCLyISEir4IiIhoYIvIhISKvgiIiGhgi8iEhL/HwguNznmSf26AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "colors = plt.rcParams[\"axes.prop_cycle\"].by_key()['color']\n",
    "\n",
    "fig, axs = plt.subplots(1,2)\n",
    "vectors = np.exp(x*1j)\n",
    "for dx,dy, c in zip (vectors.real, vectors.imag, colors):\n",
    "    axs[0].arrow(0, 0, dx, dy, head_width=0.2, head_length=0.1,length_includes_head=True, color=c, width=.04)\n",
    "axs[0].set_xlim(-1.5,1.5)\n",
    "axs[0].set_ylim(-1.5,1.5)\n",
    "axs[0].set_title(\"DS vectors\")\n",
    "\n",
    "vectors = np.exp(x*1j*2)\n",
    "for i,(dx,dy,c) in enumerate(zip(vectors.real, vectors.imag, colors)):\n",
    "    if i>=4:\n",
    "        dx = dx*0.8\n",
    "        dy = dy*0.8\n",
    "    axs[1].arrow(0, 0, dx, dy, head_width=0.2, head_length=0.1, length_includes_head=True, color=c, width=.04)\n",
    "axs[1].set_xlim(-1.5,1.5)\n",
    "axs[1].set_ylim(-1.5,1.5)\n",
    "axs[1].set_title(\"OS vectors\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def direction_selectivity(grouped_spikes_d, n_bootstrap=1000):\n",
    "    \"\"\"\n",
    "    Compute the direction selectivity index of cells in the given dict containing for each condition as\n",
    "    the keys, an array of shape (n_angle, n_repeat, trial_len, n_cell). Such dictionnary can be obtained\n",
    "    by using `utils.group_direction_response`.\n",
    "    \n",
    "    params:\n",
    "        - grouped_spikes_d: Results of the group_direction_response of shape (n_angle, n_repeat, t, n_cell)\n",
    "        - n_bootstrap: Number of bootstrap iteration to calculate the p-value\n",
    "    \n",
    "    return:\n",
    "        - A dictionnary with a key for each condition retrieving a tuple containing list of the cells\n",
    "        (spike sum, direction pref, DS idx, orientation pref, OS idx, pval DS, pval OS)\n",
    "    \"\"\"\n",
    "    \n",
    "    res_d = {}\n",
    "    for cond, sp_count in grouped_spikes_d.items():\n",
    "        n_angle = sp_count.shape[0]\n",
    "        sum_rep_spike = np.sum(sp_count, axis=(1,2)).T\n",
    "\n",
    "        x         = np.linspace(0, (n_angle-1)/4*np.pi, num=n_angle)\n",
    "\n",
    "        #Direction selectivity\n",
    "        vect_dir  = np.exp(x*1j)#np.array([np.cos(x) + np.sin(x)*1j])\n",
    "        dir_pref  = np.nan_to_num((vect_dir * sum_rep_spike).sum(axis=1) / sum_rep_spike.sum(axis=1))\n",
    "        ds_idx    = abs(dir_pref)\n",
    "\n",
    "        #Orientation selectivity\n",
    "        vect_ori  = np.exp(x*1j*2)#np.concatenate((vect_dir[:,:n_angle//2], vect_dir[:,:n_angle//2]), axis=1)\n",
    "        ori_pref  = np.nan_to_num((vect_ori * sum_rep_spike).sum(axis=1) / sum_rep_spike.sum(axis=1))\n",
    "        ori_idx   = abs(ori_pref)\n",
    "\n",
    "        #Generating direction and orientation index from shuffled trials\n",
    "        axtup_l = list(itertools.product(range(sp_count.shape[0]), range(sp_count.shape[1])))\n",
    "        random.seed(1)\n",
    "        axtup_l_shuffled = axtup_l.copy()\n",
    "        rand_ori_idx_l = np.empty((n_bootstrap, sp_count.shape[3]))\n",
    "        rand_dir_idx_l = np.empty((n_bootstrap, sp_count.shape[3]))\n",
    "        for i in range(n_bootstrap):\n",
    "            random.shuffle(axtup_l_shuffled)\n",
    "            shuffled_sp_count = np.empty(sp_count.shape)\n",
    "            for axtup, axtup_shuff in zip(axtup_l, axtup_l_shuffled):\n",
    "                shuffled_sp_count[axtup] = sp_count[axtup_shuff]\n",
    "            rand_sum_rep_spike      = np.sum(shuffled_sp_count, axis=(1,2)).T\n",
    "            rand_dir_pref     = np.nan_to_num((vect_dir * rand_sum_rep_spike).sum(axis=1) / rand_sum_rep_spike.sum(axis=1))\n",
    "            rand_dir_idx_l[i] = abs(rand_dir_pref)\n",
    "\n",
    "            rand_ori_pref     = np.nan_to_num((vect_ori * rand_sum_rep_spike).sum(axis=1) / rand_sum_rep_spike.sum(axis=1))\n",
    "            rand_ori_idx_l[i] = abs(rand_ori_pref)\n",
    "\n",
    "        #Same calculation of pval as in Baden et al 2016\n",
    "        p_val_dir = 1 - (np.sum(rand_dir_idx_l<ds_idx, axis=0)/n_bootstrap)\n",
    "        p_val_ori = 1 - (np.sum(rand_ori_idx_l<ori_idx, axis=0)/n_bootstrap)\n",
    "\n",
    "        #Finally we have to transform the orientation selectivity vectors to put them back in their\n",
    "        # original orientation, by divinding the phase of the vector by two\n",
    "        tau = np.pi*2\n",
    "        polar_ori_pref = np.array(list((map(polar, ori_pref))))\n",
    "        polar_ori_pref[:,1] = ((polar_ori_pref[:,1]+tau)%tau)/2 #Convert to positive radian angle and divide by two\n",
    "        ori_pref = np.array([rect(pol[0], pol[1]) for pol in polar_ori_pref])\n",
    "\n",
    "        res_d[cond] = (sum_rep_spike, dir_pref, ds_idx, ori_pref, ori_idx, p_val_dir, p_val_ori)\n",
    "        \n",
    "    return res_d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def wave_direction_selectivity(wave_array, spike_counts, moving_distance_th=1, looming_distance_th=.3, n_bootstrap=1000):\n",
    "    \"\"\"\n",
    "    Computes the direction, orientation and looming/shrinking indexes of the cells in response to the wave stimulus (in LED dome).\n",
    "    \n",
    "    params:\n",
    "        - wave_array: The indexes of the waves from the record master.\n",
    "        - spike_counts: The cells response to the waves from the record master.\n",
    "        - moving_distance_th: Distance threshold in radians above which a stimulus can be considered for the OS and DS.\n",
    "        - looming_distance_th: Distance threshold in radians bellow which a stimulus can be considered for the looming/shrinking response\n",
    "        - n_bootstrap: Number of repetition in the bootstraping method to compute the pvalues\n",
    "        \n",
    "    return:\n",
    "        - A tuple containing lists of the cells response, in order:\n",
    "            * [0] summed_responses: Summed response of a cell to each wave condition\n",
    "            * [1] dir_pref_l: Direction preference vector\n",
    "            * [2] dir_idx_l : Direction indexes\n",
    "            * [3] dir_pval_l: Direction p_values\n",
    "            * [4] ori_pref_l: Orientation preference vector\n",
    "            * [5] ori_idx_l : Orientation indexes\n",
    "            * [6] ori_pval_l: Orientation p_values\n",
    "            * [7] loom_idx_l : Looming indexes \n",
    "            * [8] loom_pval_l: Looming p_values\n",
    "            * [9] stas_position_l: Position (theta, phi) of the cells receptive fields with the water stimulus\n",
    "            * [10] waves_position_l: Relative positions of the waves to the cells STA.\n",
    "    \"\"\"\n",
    "    tau = np.pi*2\n",
    "    \n",
    "    indexes, order = np.unique(wave_array, return_index=True)\n",
    "    epoch_sequence = indexes[1:][np.argsort(order[1:])]\n",
    "    wave_inten     = build_wave_stimulus_array(epoch_sequence)\n",
    "    stas_wave      = process_sta_batch(wave_inten, spike_counts, Hw=1, Fw=0, return_pval=False)\n",
    "    #Hw of 1 because there is a high temporal correlation in the stimulus, so that's enough to find the RF\n",
    "    \n",
    "    summed_responses = np.zeros((100, spike_counts.shape[1]))\n",
    "    \n",
    "    for i in indexes[1:]: #Iterate from 0 to n_wave-1\n",
    "        where = np.where(wave_array==i)[0]\n",
    "        summed_responses[i,:] = np.sum(spike_counts[where,:], axis=0)\n",
    "        \n",
    "    summed_responses = summed_responses.T\n",
    "      \n",
    "    dome_positions = get_dome_positions(mode=\"spherical\")\n",
    "    \n",
    "    ori_pref_l, dir_pref_l              = [], []\n",
    "    ori_idx_l, dir_idx_l, loom_idx_l    = [], [], []\n",
    "    ori_pval_l, dir_pval_l, loom_pval_l = [], [], []\n",
    "    stas_position_l, waves_position_l   = [], []\n",
    "    \n",
    "    for sta, cell_responses  in zip(stas_wave, summed_responses):\n",
    "        maxidx_sta     = np.argmax(np.abs(sta))\n",
    "        theta_led      = dome_positions[maxidx_sta//237,maxidx_sta%237,1]\n",
    "        phi_led        = dome_positions[maxidx_sta//237,maxidx_sta%237,2]\n",
    "        relative_waves = get_waves_relative_position((theta_led, phi_led), mode=\"spherical\")\n",
    "        \n",
    "        stas_position_l.append((theta_led, phi_led))\n",
    "        waves_position_l.append(relative_waves)\n",
    "    \n",
    "        waves_distance = relative_waves[:,1]\n",
    "        waves_angle    = (relative_waves[:,2]+tau)%(tau) #Set the angle in the (0,2pi) range\n",
    "\n",
    "        looming_mask   = (waves_distance<looming_distance_th)\n",
    "        shrink_mask    = (waves_distance>np.pi-looming_distance_th)\n",
    "        waves_mask     = (waves_distance>moving_distance_th) & (waves_distance<np.pi-moving_distance_th)\n",
    "        \n",
    "        vectors_dir  = np.exp(waves_angle*1j)  #Create vectors using imaginary numbers\n",
    "        vectors_ori  = np.exp(waves_angle*1j*2)#x2 gather the vectors with opposite directions\n",
    "        dir_pref     = np.nan_to_num((vectors_dir[waves_mask] * cell_responses[waves_mask]).sum() / cell_responses[waves_mask].sum())\n",
    "        ori_pref     = np.nan_to_num((vectors_ori[waves_mask] * cell_responses[waves_mask]).sum() / cell_responses[waves_mask].sum())\n",
    "        ds_idx       = abs(dir_pref)\n",
    "        os_idx       = abs(ori_pref)\n",
    "        \n",
    "        looming_response   = (cell_responses[looming_mask]).sum()\n",
    "        shrinking_response = (cell_responses[shrink_mask]).sum()\n",
    "        looming_idx        = (looming_response-shrinking_response)/(looming_response+shrinking_response)\n",
    "        \n",
    "        ori_pref_l.append(ori_pref); dir_pref_l.append(dir_pref)\n",
    "        ori_idx_l.append(os_idx); dir_idx_l.append(ds_idx)\n",
    "        loom_idx_l.append(looming_idx)\n",
    "        \n",
    "        np.random.seed(1)\n",
    "        rand_ori_idx_l  = np.empty(n_bootstrap)\n",
    "        rand_dir_idx_l  = np.empty(n_bootstrap)\n",
    "        rand_loom_idx_l = np.empty(n_bootstrap)\n",
    "        for i in range(n_bootstrap):\n",
    "            shuffled_response = cell_responses.copy()\n",
    "            np.random.shuffle(shuffled_response)\n",
    "            \n",
    "            rand_dir_pref     = np.nan_to_num((vectors_dir[waves_mask] * shuffled_response[waves_mask]).sum() / shuffled_response[waves_mask].sum())\n",
    "            rand_ori_pref     = np.nan_to_num((vectors_ori[waves_mask] * shuffled_response[waves_mask]).sum() / shuffled_response[waves_mask].sum())\n",
    "            rand_dir_idx_l[i] = abs(rand_dir_pref)\n",
    "            rand_ori_idx_l[i] = abs(rand_ori_pref)\n",
    "            \n",
    "            rand_looming_response   = (shuffled_response[looming_mask]).sum()\n",
    "            rand_shrinking_response = (shuffled_response[shrink_mask]).sum()\n",
    "            rand_loom_idx_l[i]      = (rand_looming_response-rand_shrinking_response)/(rand_looming_response+rand_shrinking_response)\n",
    "\n",
    "        #Same calculation of pval as in Baden et al 2016\n",
    "        p_val_dir  = 1 - (np.sum(rand_dir_idx_l<ds_idx)/n_bootstrap)\n",
    "        p_val_ori  = 1 - (np.sum(rand_ori_idx_l<os_idx)/n_bootstrap)\n",
    "        p_val_loom = 1 - (np.sum(np.abs(rand_loom_idx_l)<abs(looming_idx))/n_bootstrap)\n",
    "        \n",
    "        ori_pval_l.append(p_val_ori); dir_pval_l.append(p_val_dir); loom_pval_l.append(p_val_loom)\n",
    "        \n",
    "        # original orientation, by divinding the phase of the vector by two\n",
    "        polar_ori_pref = polar(ori_pref)\n",
    "        new_vector     = ((polar_ori_pref[1]+tau)%tau)/2 #Convert to positive radian angle and divide by two\n",
    "        ori_pref       = rect(polar_ori_pref[0], new_vector)\n",
    "        \n",
    "    return (summed_responses,\n",
    "            dir_pref_l, dir_idx_l, dir_pval_l, \n",
    "            ori_pref_l, ori_idx_l, ori_pval_l,\n",
    "            loom_idx_l, loom_pval_l,\n",
    "            stas_position_l, waves_position_l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def peri_saccadic_response(spike_counts, eye_track, motion_threshold=5, window=15):   \n",
    "    \"\"\"\n",
    "    \n",
    "\n",
    "    Computes the cell average response around saccades.\n",
    "    \n",
    "    params:\n",
    "        - spike_counts: cells activity matrix of shape (t, n_cell)\n",
    "        - eye_track: Eye tracking data of shape (t, x_pos, y_pos, ...)\n",
    "        - motion_threshold: Amount of motion in pixel to account for a saccade\n",
    "        - window: Size of the window before and after the saccade on which to average the cell response\n",
    "        \n",
    "    return:\n",
    "        - peri saccadic response of cells of shape (n_cell, window*2+1)\n",
    "    \"\"\"\n",
    "    eye_shifts = np.concatenate(([0],\n",
    "                                 np.linalg.norm(eye_tracking[1:,:2]-eye_tracking[:-1,:2], axis=1)))\n",
    "    \n",
    "    #Because eye tracking is usually upsampled from 15 to 60Hz, it sums the shift, and smooth the peak\n",
    "    # detection\n",
    "    summed_shifts = np.convolve(eye_shifts, [1,1,1,1,1,1,1], mode=\"same\")\n",
    "    \n",
    "    peaks, res = signal.find_peaks(eye_shifts, height=motion_threshold, distance=10)\n",
    "    heights = res[\"peak_heights\"] #Not used for now\n",
    "    \n",
    "    psr = np.zeros((window*2, spike_bins.shape[1]))\n",
    "    for peak in peaks:\n",
    "        if peak<window or (peak+window)>len(spike_bins):\n",
    "            continue #Just ignoring peaks too close to the matrix edges\n",
    "        psr += spike_bins[peak-window:peak+window]\n",
    "    psr /= len(peaks)\n",
    "    return psr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_core.ipynb.\n",
      "Converted 01_utils.ipynb.\n",
      "Converted 02_processing.ipynb.\n",
      "Converted 03_modelling.ipynb.\n",
      "Converted 04_plotting.ipynb.\n",
      "Converted 05_database.ipynb.\n",
      "Converted 06_eyetrack.ipynb.\n",
      "Converted 10_synchro.io.ipynb.\n",
      "Converted 11_synchro.extracting.ipynb.\n",
      "Converted 12_synchro.processing.ipynb.\n",
      "Converted 13_leddome.ipynb.\n",
      "Converted 99_testdata.ipynb.\n",
      "Converted index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "from nbdev.export import *\n",
    "notebook2script()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}