ST_features_to_etho_LDA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Copyright (c) 2019 Primoz Ravbar UCSB Licensed under BSD 2-Clause [see LICENSE for details] Written by Primoz Ravbar\n",
    "\n",
    "#This script will train a LDA model with ST-features and human labels. It will next extract behaviors (produce ethograms) from all \n",
    "# ST-feature files in the folder specified. \n",
    "\n",
    "import numpy as np\n",
    "import scipy\n",
    "from scipy import ndimage\n",
    "from scipy import misc\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import os\n",
    "\n",
    "from ABRS_modules import discrete_radon_transform\n",
    "from ABRS_modules import etho2ethoAP\n",
    "from ABRS_modules import smooth_1d\n",
    "from ABRS_modules import create_LDA_training_dataset\n",
    "from ABRS_modules import removeZeroLabelsFromTrainingData\n",
    "from ABRS_modules import computeSpeedFromPosXY \n",
    "\n",
    "\n",
    "from sklearn.discriminant_analysis import LinearDiscriminantAnalysis\n",
    "\n",
    "from ABRS_data_vis import create_colorMat\n",
    "from ABRS_data_vis import cmapG\n",
    "from ABRS_data_vis import cmapAP\n",
    "\n",
    "\n",
    "pathToABRS_GH_folder = 'INSERT PATH TO ABRS MAIN FOLDER HERE'\n",
    "\n",
    "dirPathFeatures = 'INSERT PATH TO FOLDER CONTAINING ST-FEATURES HERE'\n",
    "\n",
    "dirPathLabel = pathToABRS_GH_folder + '\\\\Labels' #PATH TO FOLDER CONTAINING THE LABEL FILES\n",
    "\n",
    "idxLabelDirPathFileName = dirPathLabel + '\\\\' + 'combLabelLiplusCS7fb3' #NAME OF THE LABEL FILE\n",
    "\n",
    "\n",
    "dirPathLabelFeatures = pathToABRS_GH_folder + '\\\\LabelsFeatures\\STF_comb' #PATH TO ST-FEATURES FOR TRAINING (MUST CORRESPOND TO THE LABEL)\n",
    "\n",
    "numbFilesFeatureLabel = 7; #number of ST feature files used for LDA training\n",
    "\n",
    "outputFolderEtho = pathToABRS_GH_folder + '\\\\Etho';\n",
    "\n",
    "smoothingWindow = 89; #89 works for fly grooming behavior at 30 Hz\n",
    "halfWindowSpeed = 15; #16 or 15 for 30 Hz #adjust the time-window to extract speed of body displacement\n",
    "\n",
    "\n",
    "labelShift = 10; # if the label is not fully aligned with the ST-features, shift it; this is a patch\n",
    "#labelShift = 0; # 10 works with auto\n",
    "\n",
    "labelShiftVal = 10; # \n",
    "\n",
    "with open(idxLabelDirPathFileName, \"rb\") as f:\n",
    "     idxLabel = pickle.load(f)\n",
    "        \n",
    "with open(dirPathLabel + '\\\\' + 'labelCS1fb1_SS', \"rb\") as f:\n",
    "\t idxLabelVal = pickle.load(f)   \n",
    "        \n",
    "shLVal = np.shape(idxLabelVal);\n",
    "labelShftRightVal = np.hstack((np.zeros((1,labelShiftVal)),idxLabelVal[:,0:shLVal[1]-labelShiftVal])); # works with janelia data 11/16/2018 # shift 15 works too\n",
    "idxLabelVal = labelShftRightVal;          \n",
    "\n",
    "shL = np.shape(idxLabel);\n",
    "labelShftRight = np.hstack((np.zeros((1,labelShift)),idxLabel[:,0:shL[1]-labelShift])); # works with janelia data 11/16/2018 # shift 15 works too\n",
    "idxLabel = labelShftRight;  \n",
    "idxLabel[idxLabel==0]=7\n",
    "\n",
    "idxLabelAP=etho2ethoAP(idxLabel)\n",
    "\n",
    "featureMat, posMat, maxMovementMat = create_LDA_training_dataset (dirPathLabelFeatures,numbFilesFeatureLabel); #\n",
    "\n",
    "speedVect = computeSpeedFromPosXY (posMat,halfWindowSpeed)\n",
    "\n",
    "featureMatSm = smooth_1d (featureMat, smoothingWindow, axis = 1);\n",
    "speedVectSm = smooth_1d (speedVect, smoothingWindow, axis = 1);\n",
    "\n",
    "inputMatComb = np.vstack((featureMat,speedVect));\n",
    "inputMatAP = np.vstack((featureMatSm,speedVectSm));\n",
    "\n",
    "x_train = np.transpose(inputMatComb)\n",
    "y_train = idxLabel[0,:]\n",
    "\n",
    "model = LinearDiscriminantAnalysis();\n",
    "model.fit(x_train, y_train);\n",
    "    \n",
    "#prX = model.transform(x_train);\n",
    "prX = model.predict_proba(x_train);\n",
    "\n",
    "###########################\n",
    "x_trainAP = np.transpose(inputMatAP)\n",
    "y_trainAP = idxLabelAP[0,:]\n",
    "\n",
    "modelAP = LinearDiscriminantAnalysis();\n",
    "modelAP.fit(x_trainAP, y_trainAP);\n",
    "\n",
    "prAP = modelAP.predict_proba(x_trainAP);\n",
    "\n",
    "###########################\n",
    "\n",
    "x_trainComb = smooth_1d(np.hstack((prAP,prX)),9,axis = 0)\n",
    "y_trainComb = idxLabel[0,:]\n",
    "\n",
    "modelComb = LinearDiscriminantAnalysis();\n",
    "modelComb.fit(x_trainComb, y_trainComb);\n",
    "\n",
    "prComb = modelComb.predict_proba(x_trainComb);\n",
    "\n",
    "##############################################\n",
    "\n",
    "fileList = sorted(os.listdir(dirPathFeatures));\n",
    "sh = np.shape(fileList);\n",
    "numbFiles = sh[0] #LOOP AROUND ALL ST-FEATURE FILES (NUMBER OF ETHOGRAMS)\n",
    "#numbFiles = 1; # just the first STF file\n",
    "#numbFiles = -5; # just run the training\n",
    "\n",
    "\n",
    "for fl in range(0, numbFiles, 1):\n",
    "\n",
    "    inputFileName = fileList[fl];\n",
    "    print (inputFileName);\n",
    "\n",
    "    featureMatDirPathFileName = dirPathFeatures + '\\\\' + inputFileName;\n",
    "\n",
    "    with open(featureMatDirPathFileName, \"rb\") as f:\n",
    "         STF_30_posXY_dict = pickle.load(f)\n",
    "         \n",
    "         \n",
    "    featureMatCurrent =  STF_30_posXY_dict[\"featureMat\"];\n",
    "    \n",
    "    posMat = STF_30_posXY_dict[\"posMat\"];\n",
    "    maxMovementMat = STF_30_posXY_dict[\"maxMovementMat\"];\n",
    "    featureMatSm = smooth_1d (featureMatCurrent, smoothingWindow, axis = 1);\n",
    "    \n",
    "    speedVect = computeSpeedFromPosXY (posMat,halfWindowSpeed);\n",
    "    speedVectSm = smooth_1d (speedVect, smoothingWindow, axis = 1);\n",
    "    \n",
    "    inputMatCombCurrent = np.vstack((featureMatCurrent,speedVect));\n",
    "    inputMatAPCurrent = np.vstack((featureMatSm,speedVectSm));\n",
    "\n",
    "\n",
    "    if fl == 0:\n",
    "        \n",
    "        speedVectRec = speedVect\n",
    "        maxMovementRec = maxMovementMat\n",
    "        \n",
    "        inputMatComb = inputMatCombCurrent;\n",
    "        inputMatAP = inputMatAPCurrent;\n",
    "        \n",
    "        x_pred = np.transpose(inputMatComb);\n",
    "        \n",
    "        x_predAP = np.transpose(inputMatAP);      \n",
    "        \n",
    "        predictions = model.predict_proba(x_pred)\n",
    "        predictionsAP = modelAP.predict_proba(x_predAP)\n",
    "        \n",
    "        x_predComb = smooth_1d(np.hstack((predictionsAP,predictions)),9,axis=0)\n",
    "        \n",
    "        predictionsComb = modelComb.predict_proba(x_predComb);\n",
    "                                                          \n",
    "    if fl > 0:\n",
    "\n",
    "        speedVectRec = np.hstack((speedVectRec,speedVect))\n",
    "        maxMovementRec = np.hstack((maxMovementRec,maxMovementMat))\n",
    "        \n",
    "        inputMatComb = inputMatCombCurrent\n",
    "        inputMatAP = inputMatAPCurrent\n",
    "        \n",
    "        x_pred = np.transpose(inputMatComb);\n",
    "        \n",
    "        x_predAP = np.transpose(inputMatAP);\n",
    "        \n",
    "        predictionsCurrent = model.predict_proba(x_pred)\n",
    "        predictionsAPCurrent = modelAP.predict_proba(x_predAP)\n",
    "        \n",
    "        predictions = np.vstack((predictions,predictionsCurrent))\n",
    "        predictionsAP = np.vstack((predictionsAP,predictionsAPCurrent))\n",
    "     \n",
    "        x_predComb = smooth_1d(np.hstack((predictionsAPCurrent,predictionsCurrent)),9,axis=0)\n",
    "        \n",
    "        predictionsCombCurrent = modelComb.predict_proba(x_predComb);\n",
    "        \n",
    "        predictionsComb = np.vstack((predictionsComb,predictionsCombCurrent))\n",
    "        \n",
    "\n",
    "predictions=np.transpose(predictions)\n",
    "predictionsAP=np.transpose(predictionsAP)\n",
    "predictionsComb=np.transpose(predictionsComb)\n",
    "        \n",
    "thrX = smooth_1d (predictions, 89, axis = 1)\n",
    "thrComb = smooth_1d (predictionsComb, 89, axis = 1)\n",
    "\n",
    "\n",
    "thrX_left = np.hstack((thrX[:,30:np.shape(thrX)[1]], np.zeros((np.shape(thrX)[0],30)) )) \n",
    "thrX = thrX_left\n",
    "\n",
    "thrComb_left = np.hstack((thrComb[:,30:np.shape(thrComb)[1]], np.zeros((np.shape(thrComb)[0],30)) )) \n",
    "thrComb = thrComb_left\n",
    "\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this code applies user-specified thresholds to the LDA features to produce ethograms\n",
    "\n",
    "protocolName = 7 \n",
    "\n",
    "shPred = np.shape(predictions);\n",
    "\n",
    "etho = np.zeros((1,shPred[1]));\n",
    "ethoAP = np.zeros((1,shPred[1]));\n",
    "lengthEtho=shPred[1]\n",
    "\n",
    "if protocolName == 7:\n",
    "\n",
    "\n",
    "    for i in range(0, lengthEtho):\n",
    "        \n",
    "        if (predictionsComb[1,i]-predictionsComb[0,i])<(thrComb[1,i]-thrComb[0,i]) and (predictionsAP[1,i]-predictionsAP[2,i])>0:   \n",
    "            etho[0,i] = 1\n",
    "        \n",
    "        if (predictionsComb[1,i]-predictionsComb[0,i])>(thrComb[1,i]-thrComb[0,i]) and (predictionsAP[1,i]-predictionsAP[2,i])>0:   \n",
    "            etho[0,i] = 2\n",
    "            \n",
    "        if (predictionsComb[3,i]-predictionsComb[2,i])<(thrComb[3,i]-thrComb[2,i]) and (predictionsAP[1,i]-predictionsAP[2,i])<0:   \n",
    "            etho[0,i] = 3\n",
    "        \n",
    "        if (predictionsComb[3,i]-predictionsComb[2,i])>(thrComb[3,i]-thrComb[2,i]) and (predictionsAP[1,i]-predictionsAP[2,i])<0:    \n",
    "            etho[0,i] = 4   \n",
    "            \n",
    "        if (predictionsComb[4,i])>0.01 and (predictionsComb[4,i]-predictionsComb[3,i])>0 and (predictionsAP[1,i]-predictionsAP[2,i])<0 and \\\n",
    "            maxMovementRec[0,i]>100:    \n",
    "            etho[0,i] = 5    \n",
    "                     \n",
    "        if (predictionsAP[3,i]+predictionsComb[5,i]>1.0 or speedVectRec[0,i]>25) and ethoAP[0,i]==0: \n",
    "            etho[0,i] = 6\n",
    "            \n",
    "        if  maxMovementRec[0,i]<100:    \n",
    "            etho[0,i] = 7   \n",
    "              \n",
    "sizeEtho = 50000            \n",
    "ethoZ = etho;ethoZ[0,5]=7;ethoZ[0,6]=0  \n",
    "\n",
    "ethoMat = np.reshape(etho[0,:],(numbFiles,sizeEtho));  \n",
    "            \n",
    "plt.matshow(ethoMat, interpolation=None, aspect='auto',cmap=cmapG);plt.show()  "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}