data_exploration_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 0.0: Load packages and customised functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## a useful function to generate a data list for further analysis\n",
    "import os,json,sys,itertools\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from locustvr_converter import preprocess_matrex_data\n",
    "from behavioural_classification import classify_heading_direction\n",
    "##need to add this additional cell because useful tools are in another folder. Need to integrate these two folders one day\n",
    "current_working_directory = Path.cwd()\n",
    "parent_dir = current_working_directory.resolve().parents[0]\n",
    "sys.path.insert(0, str(parent_dir) + \"\\\\utilities\")\n",
    "from useful_tools import select_animals_gpt,find_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "from matplotlib import cm\n",
    "class MplColorHelper:\n",
    "    def __init__(self, cmap_name, start_val, stop_val):\n",
    "        self.cmap_name = cmap_name\n",
    "        self.cmap = plt.get_cmap(cmap_name)\n",
    "        self.norm = mpl.colors.Normalize(vmin=start_val, vmax=stop_val)\n",
    "        self.scalarMap = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)\n",
    "\n",
    "    def get_rgb(self, val):\n",
    "        return self.scalarMap.to_rgba(val)\n",
    "colormap_name = \"viridis\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 0.1: Load analysis methods in python dictionary form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_file = \"./analysis_methods_dictionary.json\"\n",
    "with open(json_file, \"r\") as f:\n",
    "    analysis_methods = json.loads(f.read())\n",
    "    \n",
    "\n",
    "#Put the folder of your Unity experiment below\n",
    "#thisDataset =\"D:/MatrexVR_Swarm_Data/RunData\"\n",
    "#thisDataset =\"D:/MatrexVR_blackbackground_Data/RunData\"\n",
    "thisDataset =\"D:/MatrexVR_grass1_Data/RunData\"\n",
    "#thisDataset =\"D:/MatrexVR_navigation_Data/RunData\"\n",
    "#thisDataset =\"D:/MatrexVR_2024_Data/RunData\"\n",
    "#parameter name means independent variable in the experiment\n",
    "#parameter_name='kappa' \n",
    "parameter_name='mu'\n",
    "#parameter_name='agent_speed'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis_methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 0.2: Load animals' experiment directory into a list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## this cell searches for a folder with csv files, usually that is the folder saving the tracking data.\n",
    "## Since data from the 4 VRs are saved in the same folder, this command will return that one folder for the 4 experiment\n",
    "dir_list = []\n",
    "file_type=\".csv\"\n",
    "for root, dirs, files in os.walk(thisDataset):\n",
    "    for folder in dirs:\n",
    "        folder_path=os.path.join(root,folder)\n",
    "        if any(name.endswith(file_type) for name in os.listdir(folder_path)):\n",
    "            dir_list.append(folder_path.replace(\"\\\\\", \"/\"))\n",
    "\n",
    "print(f\"these directories are found {dir_list}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_list[7:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_list=dir_list[7:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Session 0.3: pass temperature information into each folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##This cell is used to move data of the thermo-humidity logger to animals'folder\n",
    "import shutil\n",
    "#tmp_file_name='matrexVR240824-240901.txt'\n",
    "#tmp_file_name='DL220THP_Thermo2_241012_241014.csv'\n",
    "tmp_file_name='DL220THP_Thermo3_241207_241210.csv'\n",
    "tmp_source=os.path.join('Z:/Users/chiyu',tmp_file_name)\n",
    "for this_dir in dir_list:\n",
    "    tmp_new_dir = os.path.join(this_dir,tmp_file_name)\n",
    "    if os.path.isfile(tmp_new_dir):\n",
    "        print(\"Found EL USB temperature file in the new directory already\")\n",
    "    else:\n",
    "        shutil.copy(tmp_source, tmp_new_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 1.0: Create curated dataset based on a list of experiment directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This function receives directory path that contains the 4-VR data and save the tracking + stimulus information as h5 file\n",
    "pattern=\"VR*.h5\"\n",
    "for this_dir in dir_list:\n",
    "    if \"archive\" in this_dir:\n",
    "        print(f\"skip archive folder for {this_dir}\")\n",
    "        continue\n",
    "    if any(Path(this_dir).glob(pattern)) and analysis_methods.get(\"overwrite_curated_dataset\")==False:\n",
    "        print(f\"curated matrexvr h5 database found in {this_dir}. Skip this file\")\n",
    "        continue\n",
    "    else:\n",
    "        print(f\"no curated matrexvr h5 database in {this_dir}. Create curated file\")\n",
    "        preprocess_matrex_data(this_dir,analysis_methods)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.0: select animal based on condition and return which a directory list and a list of vr rig number to specify which animal to analyse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build up dir_list of animals based on condition.\n",
    "dir_list = []\n",
    "file_type=\".h5\"\n",
    "using_google_sheet=True\n",
    "sheet_name = \"Unity_MatrexVR\"\n",
    "scene_name=analysis_methods.get(\"experiment_name\")\n",
    "if analysis_methods.get(\"load_individual_data\") == True:\n",
    "    if using_google_sheet==True:\n",
    "        database_id = \"1UL4eEUrQMapx9xz11-IyOSlPBcep1I9vBJ2uGgVudb8\"\n",
    "                #https://docs.google.com/spreadsheets/d/1UL4eEUrQMapx9xz11-IyOSlPBcep1I9vBJ2uGgVudb8/edit?usp=sharing\n",
    "        url = f\"https://docs.google.com/spreadsheets/d/{database_id}/gviz/tq?tqx=out:csv&sheet={sheet_name}\"\n",
    "        #df = pd.read_excel(url, engine='openpyxl')## use this function if the file is not google sheet but uploaded excel file\n",
    "        df = pd.read_csv(url)\n",
    "    else:\n",
    "        excel_file_path = \"Z:/DATA/experiment_trackball_Optomotor/Locusts Management.xlsx\"\n",
    "        print(f\"using a database {excel_file_path} from the server but this file might be outdated\")\n",
    "        # Create a 'with' statement to open and read the Excel file\n",
    "        with pd.ExcelFile(excel_file_path) as xls:\n",
    "            # Read the Excel sheet into a DataFrame with the sheet name (folder name)\n",
    "            df = pd.read_excel(xls, sheet_name)\n",
    "        ##list up the conditions and answers as strings for input argument to select animal. One condition must pair with one answer\n",
    "    if analysis_methods.get(\"select_animals_by_condition\") == True:\n",
    "       #animal_of_interest=select_animals_gpt(df,\"Independent variable (list up all of them in the experiment)\",\"gregarious_leader_grass\",\"Excluding this animal from analysis (Usually when animals die or molt, T/F)\",\"F\")\n",
    "       #animal_of_interest=select_animals_gpt(df,\"Independent variable (list up all of them in the experiment)\",\"gregarious_leader_black\")\n",
    "        animal_of_interest=select_animals_gpt(df,\"Independent variable (list up all of them in the experiment)\",parameter_name,\"Excluding this animal from analysis (Usually when animals die or molt, T/F)\",\"F\")\n",
    "        #animal_of_interest=select_animals_gpt(df,\"Independent variable (list up all of them in the experiment)\",parameter_name)\n",
    "    else:\n",
    "        animal_of_interest=df\n",
    "    folder_name=animal_of_interest[\"folder name\"].values\n",
    "    dir_tile=np.tile(thisDataset, (len(folder_name), 1))\n",
    "    vr_no=animal_of_interest[\"VR number\"].values\n",
    "    dir_list = [''.join([x[0], '/', y]) for x,y in zip(dir_tile,folder_name)]\n",
    "    dir_dict = zip(dir_list, vr_no.tolist())\n",
    "else:\n",
    "    for root, dirs, files in os.walk(thisDataset):\n",
    "        for folder in dirs:\n",
    "            folder_path=os.path.join(root,folder)\n",
    "            if any(name.endswith(file_type) for name in os.listdir(folder_path)):\n",
    "                dir_list.append(folder_path.replace(\"\\\\\", \"/\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.2: connecting information between two H5 files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.2.1: loading meta info from the json file (work in progress)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#work in progress, this read condition directly from Control scene json file, which will be useful when not saving trial information into the simulated-agent file\n",
    "'''\n",
    "def read_condition_json(dir_list):\n",
    "    j_pattern = f\"*sequenceConfig.json\"\n",
    "    conditions_list=[]\n",
    "    for this_dir in dir_list:\n",
    "        found_result = find_file(Path(this_dir), j_pattern) \n",
    "        with open(found_result) as f:\n",
    "\n",
    "            d = json.load(f)\n",
    "            print(d[\"sequence\"])\n",
    "\n",
    "        for i in d[\"sequence\"]:\n",
    "            print(\"Mu:\", i['Mu'])\n",
    "            print(\"kappa:\", i['kappa'])\n",
    "            print(\"LocustSpeed:\", i['LocustSpeed'])\n",
    "            density = int(n_locusts.split(\":\")[1]) / (\n",
    "            int(boundary_size.split(\":\")[1]) ** 2 / 10000\n",
    "        )\n",
    "            conditions = {\n",
    "            \"Density\": density,\n",
    "            mu.split(\":\")[0]: int(mu.split(\":\")[1]),\n",
    "            \"kappa\": d[\"sequences\"][0]['parameters']['kappa'],\n",
    "            agent_speed.split(\":\")[0]: d[\"sequences\"][0]['parameters']['LocustSpeed'],\n",
    "        }\n",
    "        conditions_list.append(conditions)\n",
    "    return conditions_list\n",
    "dir_list[0]\n",
    "j_pattern = f\"*sequenceConfig.json\"\n",
    "found_result = find_file(Path(this_dir), j_pattern)\n",
    "test=pd.read_json(found_result)\n",
    "print(test)\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.2.2: combine tables in the two lists in to 2 big tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# When using spatial discretization, information about tracking quality is not logged in the dfXY, \n",
    "# hence there is a need to pass that information from df\n",
    "def connect_two_tables(dir_list,analysis_methods,test_parameter='kappa',vr_no=[]):\n",
    "    scene_name=analysis_methods.get(\"experiment_name\")\n",
    "    df_all=[]\n",
    "    dfxy_all=[]\n",
    "    dir_iterator=[]\n",
    "    if len(vr_no)>0:\n",
    "        print(\"i am using list\")\n",
    "        dir_iterator=zip(dir_list,vr_no)\n",
    "    elif type(dir_list)==dict:\n",
    "        print(\"i am using dictionary\")\n",
    "        dir_iterator=dir_dict\n",
    "    else:\n",
    "        print(\"there is a bug\")\n",
    "        return df_all,dfxy_all\n",
    "    for this_dir,this_vr in dir_iterator:\n",
    "        if Path(this_dir).is_dir()==False:\n",
    "            continue\n",
    "        summary_pattern = f\"VR{this_vr}*score.h5\"\n",
    "        xy_pattern = f\"VR{this_vr}*XY.h5\"\n",
    "        found_result = find_file(Path(this_dir), summary_pattern)        \n",
    "        df = pd.read_hdf(found_result)\n",
    "        df['VR'] = np.tile(f\"VR{this_vr}\", (len(df), 1))\n",
    "        df['VR'] =df[\"VR\"]+\"_\"+df[\"fname\"]\n",
    "        #COL = MplColorHelper(colormap_name, 0, num_independent_variable)\n",
    "        found_result = find_file(Path(this_dir), xy_pattern)\n",
    "        dfxy = pd.read_hdf(found_result)\n",
    "        dfxy['VR'] = np.tile(f\"VR{this_vr}\", (len(dfxy), 1))\n",
    "        dfxy['VR'] =dfxy[\"VR\"]+\"_\"+dfxy[\"fname\"]\n",
    "        #df.loc[(df[\"distTotal\"]<10.0) | (df[\"loss\"]> 0.05), \"distTotal\"] = np.nan\n",
    "        ##hardcode color code here for scatter plot\n",
    "        if test_parameter == 'kappa':\n",
    "            color_code={0.1: 0.2, 1.0: 0.4, 10.0: 0.6,100000.0:1}\n",
    "        elif test_parameter == 'mu':\n",
    "            if scene_name.lower()=='choice':\n",
    "                #color_code={0: 0.1,45: 0.4,315: 0.7}\n",
    "                color_code={0: 0.1,135: 0.4,225: 0.7}\n",
    "            elif scene_name.lower()=='swarm':\n",
    "                color_code={0: 0.1, 45: 0.2, 90: 0.3,135:0.4,180: 0.5, 225: 0.6, 270: 0.7,315:0.8}\n",
    "            elif scene_name.lower()=='band':\n",
    "                color_code={0: 0.1, 45: 0.2, 90: 0.3, 270: 0.7,315:0.8}\n",
    "            else:\n",
    "                return Warning('scene name not found')\n",
    "        elif test_parameter == 'agent_speed':\n",
    "            color_code={1.0: 0.2,2.0: 0.4, 4.0: 0.6,8.0:1}\n",
    "        else:\n",
    "            return Warning('test parameter not found')\n",
    "        df['color_code'] = df[test_parameter].map(color_code)\n",
    "        df_all.append(df)\n",
    "        dfxy_all.append(dfxy)\n",
    "    return df_all,dfxy_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'vr_no' in locals():\n",
    "    df_all,dfxy_all=connect_two_tables(dir_list,analysis_methods,parameter_name,vr_no)\n",
    "else:\n",
    "    df_all,dfxy_all=connect_two_tables(dir_dict,analysis_methods,parameter_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3: make summary plots of animals' response with customised plotting functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.1: introduce customised plotting functions used in Sercan's paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#introduce customised functions\n",
    "def plot_sercansincos(df,analysis_methods,parameters,parameter_name,vr_num='all'):\n",
    "    save_output= analysis_methods.get(\"save_output\")\n",
    "    scene_name=analysis_methods.get(\"experiment_name\")\n",
    "    if analysis_methods.get(\"active_trials_only\"):\n",
    "        active_trial='active_trials'\n",
    "    else:\n",
    "        active_trial=''\n",
    "    \n",
    "    cos = df[\"cos\"]\n",
    "    sin = df[\"sin\"]\n",
    "    if 'density' in df.columns:\n",
    "        density=df[\"density\"].unique()[0]\n",
    "        cos_fig_name=f\"{vr_num}_cos_{scene_name}_{parameter_name}_{parameters}_density_{int(density)}{active_trial}.svg\"\n",
    "        sin_fig_name=f\"{vr_num}_sin_{scene_name}_{parameter_name}_{parameters}_density_{int(density)}{active_trial}.svg\" \n",
    "    elif scene_name=='choice':\n",
    "        cos_fig_name=f\"{vr_num}_cos_{scene_name}_{parameter_name}_{parameters}_single_target_{df['type'].values[0]}{active_trial}.svg\"\n",
    "        sin_fig_name=f\"{vr_num}_sin_{scene_name}_{parameter_name}_{parameters}_single_target_{df['type'].values[0]}{active_trial}.svg\"\n",
    "\n",
    "    fig, ax = plt.subplots(dpi=300, figsize=(1.1,0.25))\n",
    "    #plt.rcParams.update(plt.rcParamsDefault)\n",
    "    plt.rcParams.update({'font.size': 8})\n",
    "    plt.rcParams['font.family'] = 'Arial'\n",
    "    plt.set_cmap('cividis')\n",
    "    # Set the axis line width to 2\n",
    "    plt.rcParams['ytick.major.width'] = 2\n",
    "    plt.rcParams['xtick.major.width'] = 2\n",
    "    plt.rcParams['axes.linewidth'] = 2\n",
    "    plt.subplots_adjust(bottom=0.4)\n",
    "    sns.kdeplot(cos, cut=0, color=\"#21918c\", fill=True, alpha=0.9)\n",
    "    plt.xlim(-1,1)\n",
    "    plt.title(\"r cos\\u03F4\")\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['left'].set_visible(False)\n",
    "    ax.set_yticks([])\n",
    "    plt.ylabel(\"\")\n",
    "    plt.xlabel(\"\")    \n",
    "    if save_output==True:\n",
    "        plt.savefig(cos_fig_name)\n",
    "    fig, ax = plt.subplots(dpi=300, figsize=(1.1,0.25))\n",
    "    plt.subplots_adjust(bottom=0.4)\n",
    "    sns.kdeplot(sin, cut=0, color=\"#21918c\",  fill=True, alpha=0.9)#),lw=1,\n",
    "    plt.xlim(1,-1)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['left'].set_visible(False)\n",
    "    plt.xticks(rotation = 90)\n",
    "    ax.set_yticks([])\n",
    "    plt.ylabel(\"\")\n",
    "    plt.xlabel(\"\")\n",
    "    plt.title(\"r sin\\u03F4\")\n",
    "    if save_output==True:\n",
    "        plt.savefig(sin_fig_name)\n",
    "    plt.show()\n",
    "def plot_sercantrajec(dfXY,analysis_methods,parameters,parameter_name,trajec_lim=1000,vr_num='all'):\n",
    "    save_output= analysis_methods.get(\"save_output\")\n",
    "    scene_name=analysis_methods.get(\"experiment_name\")\n",
    "    if analysis_methods.get(\"active_trials_only\"):\n",
    "        active_trial='active_trials'\n",
    "    else:\n",
    "        active_trial=''\n",
    "    \n",
    "    a = dfXY.groupby('VR')\n",
    "    if 'density' in dfXY.columns:\n",
    "        density=dfXY[\"density\"].unique()[0]\n",
    "        print(density)\n",
    "        fig_name=f\"{vr_num}_summary_trajectory_{scene_name}_{parameter_name}_{parameters}_density_{int(density)}{active_trial}.png\"\n",
    "    else:\n",
    "        fig_name=f\"{vr_num}_summary_trajectory_{scene_name}_{parameter_name}_{parameters}_single_target_{dfXY['type'].values[0]}{active_trial}.png\"\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(3,3), dpi=300) \n",
    "    plt.rcParams.update(plt.rcParamsDefault)\n",
    "    plt.rcParams.update({'font.size': 8})\n",
    "    plt.rcParams['ytick.major.width'] = 2\n",
    "    plt.rcParams['xtick.major.width'] = 2\n",
    "    plt.rcParams['axes.linewidth'] = 2\n",
    "    # Get the colormap\n",
    "    cmap = plt.get_cmap('viridis')\n",
    "\n",
    "    #plt.style.use('dark_background') \n",
    "    for i, (key2, grp2) in enumerate(a):\n",
    "        color = cmap(i/ len(a))\n",
    "        plt.plot(grp2[\"X\"].values, grp2[\"Y\"].values, color=color, linewidth=1)\n",
    "\n",
    "    if type(parameters)==tuple:\n",
    "        this_variable=parameters[0]\n",
    "    else:\n",
    "        this_variable=parameters\n",
    "\n",
    "\n",
    "        # Here plot agent's trajectory with hardcode parameters\n",
    "    if scene_name=='choice' and dfXY['type'].values[0] !='empty_trial':\n",
    "        agent_speed=2\n",
    "        radial_distance=8\n",
    "        duration=60\n",
    "        travel_direction=this_variable*-np.pi/180#the radian circle is clockwise in Unity, so 45 degree should be used as -45 degree in the regular radian circle\n",
    "        radial_distance_b=np.cos(travel_direction)*radial_distance\n",
    "        delta_cos=np.cumsum(np.repeat(np.cos(travel_direction)*agent_speed, duration))\n",
    "        agent_cos=radial_distance_b+delta_cos\n",
    "        radial_distance_b=np.sin(travel_direction)*radial_distance\n",
    "        delta_sin=np.cumsum(np.repeat(np.sin(travel_direction)*agent_speed, duration))\n",
    "        agent_sin=radial_distance_b+delta_sin\n",
    "        plt.plot(agent_cos, agent_sin, color='k', linewidth=1)\n",
    "\n",
    "    plt.xlim(-1*trajec_lim, trajec_lim)\n",
    "    plt.ylim(-1*trajec_lim, trajec_lim)\n",
    "    plt.yticks([-1*trajec_lim, 0, trajec_lim])\n",
    "    plt.xticks([-1*trajec_lim, 0, trajec_lim])\n",
    "        # Set the aspect ratio to be equal\n",
    "    plt.gca().set_aspect('equal')                                                  \n",
    "    if save_output==True:    \n",
    "        plt.savefig(fig_name)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.2: introduce a pilot scatter plot functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_travel_distance_set(df_all,analysis_methods,parameter_name,y_axis_lim=[0.1,1000]):\n",
    "    colormap_name = \"nipy_spectral\"\n",
    "    COL = MplColorHelper(colormap_name, 0, 10)\n",
    "    \n",
    "    fig, (ax1, ax2,ax3) = plt.subplots(\n",
    "    nrows=1, ncols=3, figsize=(18, 6), tight_layout=True\n",
    ")\n",
    "    colour_code=analysis_methods.get(\"graph_colour_code\")\n",
    "    if parameter_name=='kappa':\n",
    "        ax1.set_xscale('log')\n",
    "    #ax1.set_yscale('log')\n",
    "    ax1.set_ylim([y_axis_lim[0],y_axis_lim[1]])\n",
    "    ax1.set(\n",
    "        yticks=[y_axis_lim[0], y_axis_lim[1]],\n",
    "        ylabel=\"Travel distance per second\",\n",
    "        xticks=sorted(df_all[0][parameter_name].unique()),\n",
    "        xlabel=parameter_name,\n",
    "    )\n",
    "    ax2.set(\n",
    "        ylabel=\"trial n\",\n",
    "        xlabel=\"interval prior to trial n\",\n",
    "        xlim=[y_axis_lim[0],y_axis_lim[1]],\n",
    "        ylim=[y_axis_lim[0],y_axis_lim[1]],\n",
    "        title=\"travel distance per second\",\n",
    "    )\n",
    "    ax3.set(\n",
    "        ylabel=\"trial n\",\n",
    "        xlabel=\"trial n-1\",\n",
    "        xlim=[y_axis_lim[0],y_axis_lim[1]],\n",
    "        ylim=[y_axis_lim[0],y_axis_lim[1]],\n",
    "        title=\"travel distance per second\",\n",
    "    )\n",
    "    #ax1.gca().set_aspect('equal')\n",
    "    for id in np.arange(len(df_all)):\n",
    "        df=df_all[id]\n",
    "        viridis_code=df['color_code']*10#normalise the code from 1 to 10\n",
    "        viridis_code=viridis_code.astype('int')\n",
    "        #set some thresholds to remove bad tracking \n",
    "        #df.loc[(df[\"distTotal\"]<50.0) | (df[\"loss\"]> 0.05), \"distTotal\"] = np.nan\n",
    "        if df.iloc[0][\"VR\"].startswith('VR1'):\n",
    "            vr_color=colour_code[0]\n",
    "        elif df.iloc[0][\"VR\"].startswith('VR2'):\n",
    "            vr_color=colour_code[1]\n",
    "        elif df.iloc[0][\"VR\"].startswith('VR3'):\n",
    "            vr_color=colour_code[2]\n",
    "        else:\n",
    "            vr_color=colour_code[3]\n",
    "\n",
    "        #ax1.scatter(df.iloc[::2][parameter_name], df.iloc[::2][\"distTotal\"]/df.iloc[::2][\"duration\"],c='k')\n",
    "\n",
    "        ax1.scatter(df.iloc[1::2][parameter_name], df.iloc[1::2][\"distTotal\"]/df.iloc[1::2][\"duration\"],c=COL.get_rgb(viridis_code[1::2]))\n",
    "        if df.shape[0] % 2 == 0:\n",
    "            ax2.scatter(df.iloc[::2][\"distTotal\"]/df.iloc[::2][\"duration\"],df.iloc[1::2][\"distTotal\"]/df.iloc[1::2][\"duration\"],c=COL.get_rgb(viridis_code[1::2]))\n",
    "        else:\n",
    "            ax2.scatter(df.iloc[:-1:2][\"distTotal\"]/df.iloc[:-1:2][\"duration\"],df.iloc[1::2][\"distTotal\"]/df.iloc[1::2][\"duration\"],c=COL.get_rgb(viridis_code[1::2]))\n",
    "        ax3.scatter(df.iloc[1:-2:2][\"distTotal\"]/df.iloc[1:-2:2][\"duration\"],df.iloc[3::2][\"distTotal\"]/df.iloc[3::2][\"duration\"],c=COL.get_rgb(viridis_code[3::2]))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot responses (mean angle and travel distance) from individual experiments (usually every 4 animal an experiment; different colour mark different animals in that experiment) \n",
    "#or comparing trial by trial response through normalised response (e.g. ratio to previous trial) or scatter plot (each dot means a comparison, \n",
    "#different colour means data from different rigs, different independent variables is marked with different kappa value)\n",
    "## 1st: plots with independent variables such as kappa or mu against travel distance or angle\n",
    "if len(df_all)>0:\n",
    "    plot_travel_distance_set(df_all,analysis_methods,parameter_name,y_axis_lim=[0,12])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.3: combine pandas dataframe across animals and filter out trials with bad tracking (these are preprocessing steps to use Sercan's functions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Firstly, concatenate every animal's dataframe into a big table and then sort them based on conditions.\n",
    "if len(dfxy_all)>0:\n",
    "    dfxy_con = pd.concat(dfxy_all)\n",
    "if len(df_all)>0:\n",
    "    df_con = pd.concat(df_all)\n",
    "good_tracking=df_con['loss']< 0.05\n",
    "active_trials=(df_con['loss'] < 0.05) & (df_con[\"distTotal\"]>50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis_methods.update({\"save_output\": True})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.4: plot trial by trial trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plotting trajectory\n",
    "#differentiate between stim and ISI based on columns in dfxy_con\n",
    "if analysis_methods.get(\"active_trials_only\"):\n",
    "    df_interest=df_con[active_trials]\n",
    "else:\n",
    "    df_interest=df_con[good_tracking]\n",
    "\n",
    "\n",
    "if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "    stim_or_isi=dfxy_con['radial_distance']\n",
    "elif analysis_methods.get(\"experiment_name\")==\"swarm\" or analysis_methods.get(\"experiment_name\")==\"band\":\n",
    "    stim_or_isi=dfxy_con['density']\n",
    "df_stim=dfxy_con.loc[(dfxy_con['VR'].isin(df_interest['VR'])) & (stim_or_isi>0)]\n",
    "for key, grp in df_stim.groupby([parameter_name,'type']):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercantrajec(grp,analysis_methods,key[0],parameter_name,500)\n",
    "df_isi=dfxy_con.loc[(dfxy_con['VR'].isin(df_interest[\"VR\"])) & (stim_or_isi==0)]\n",
    "for key, grp in df_isi.groupby([parameter_name,'type']):\n",
    "#for key, grp in df_isi.groupby(parameter_name):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercantrajec(grp,analysis_methods,key[0],parameter_name,500)\n",
    "#xy_lim at around 2000 is good for trial lasts around 4 or 5 min\n",
    "#xy_lim at around 500 is good for trial lasts around 1 min"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.5: pool mean angle together to make KDE plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Visualise the distribution of mean angle at sin and cos using seaborn kernel density estimation plot\n",
    "#differentiate between stim and ISI based on columns in df_con\n",
    "if analysis_methods.get(\"active_trials_only\"):\n",
    "    df_interest=df_con[active_trials]\n",
    "else:\n",
    "    df_interest=df_con\n",
    "\n",
    "if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "    stim_or_isi=df_interest['radial_distance']\n",
    "elif analysis_methods.get(\"experiment_name\")==\"swarm\" or analysis_methods.get(\"experiment_name\")==\"band\":\n",
    "    stim_or_isi=df_interest['density']\n",
    "df_stim=df_interest[stim_or_isi>0]\n",
    "\n",
    "for key, grp in df_stim.groupby([parameter_name,'type']):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercansincos(grp,analysis_methods,key[0],parameter_name)\n",
    "df_isi=df_interest[stim_or_isi==0]\n",
    "for key, grp in df_isi.groupby([parameter_name,'type']):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercansincos(grp,analysis_methods,key[0],parameter_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_circular_histrogram(df,analysis_methods,parameters,parameter_name,vr_num='all'):\n",
    "    save_output= analysis_methods.get(\"save_output\")\n",
    "    scene_name=analysis_methods.get(\"experiment_name\")\n",
    "    if analysis_methods.get(\"active_trials_only\"):\n",
    "        active_trial='active_trials'\n",
    "    else:\n",
    "        active_trial=''\n",
    "    \n",
    "    angles = df[\"mean_angle\"]\n",
    "    if 'density' in df.columns:\n",
    "        density=df[\"density\"].unique()[0]\n",
    "        hist_fig_name=f\"{vr_num}_circular_hist_{scene_name}_{parameter_name}_{parameters}_density_{int(density)}{active_trial}.svg\"\n",
    "    elif scene_name=='choice':\n",
    "        hist_fig_name=f\"{vr_num}_circular_hist_{scene_name}_{parameter_name}_{parameters}_single_target_{df['type'].values[0]}{active_trial}.svg\"\n",
    "\n",
    "    ax = plt.subplot(111, polar=True)\n",
    "    ax.hist(angles, bins=24, alpha=0.75)\n",
    "    #ax.set_xticks([])\n",
    "    ax.set_yticks([5,10])\n",
    "    ax.set_xticklabels([])\n",
    "    #ax.set_title(f'banding direction: {parameters}')\n",
    "    if save_output==True:\n",
    "        plt.savefig(hist_fig_name)\n",
    "    plt.show()\n",
    "def plot_travel_histrogram(df,analysis_methods,parameters,parameter_name,vr_num='all'):\n",
    "    save_output= analysis_methods.get(\"save_output\")\n",
    "    scene_name=analysis_methods.get(\"experiment_name\")\n",
    "    if analysis_methods.get(\"active_trials_only\"):\n",
    "        active_trial='active_trials'\n",
    "    else:\n",
    "        active_trial=''\n",
    "    distTotal = df[\"distTotal\"]\n",
    "    if 'density' in df.columns:\n",
    "        density=df[\"density\"].unique()[0]\n",
    "        hist_fig_name=f\"{vr_num}_travel_hist_{scene_name}_{parameter_name}_{parameters}_density_{int(density)}{active_trial}.svg\"\n",
    "    elif scene_name=='choice':\n",
    "        hist_fig_name=f\"{vr_num}_travel_hist_{scene_name}_{parameter_name}_{parameters}_single_target_{df['type'].values[0]}{active_trial}.svg\"\n",
    "\n",
    "    #fig, ax = plt.subplots(dpi=300, figsize=(1,1))\n",
    "    fig, ax = plt.subplots(figsize=(5,2))\n",
    "    ax.hist(distTotal, bins=24, alpha=0.75)\n",
    "    ax.set_xticks([50,100])\n",
    "    #ax.set_yticks([5,10])\n",
    "    #ax.set_xticklabels([])\n",
    "    #ax.set_title(f'banding direction: {parameters}')\n",
    "    if save_output==True:\n",
    "        plt.savefig(hist_fig_name)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Visualise the distribution of  mean angle using seaborn kernel density estimation plot\n",
    "#differentiate between stim and ISI based on columns in df_con\n",
    "if analysis_methods.get(\"active_trials_only\"):\n",
    "    df_interest=df_con[active_trials]\n",
    "else:\n",
    "    df_interest=df_con\n",
    "if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "    stim_or_isi=df_interest['radial_distance']\n",
    "elif analysis_methods.get(\"experiment_name\")==\"swarm\" or analysis_methods.get(\"experiment_name\")==\"band\":\n",
    "    stim_or_isi=df_interest['density']\n",
    "#df_stim=df_con[stim_or_isi>0]\n",
    "df_stim=df_interest[stim_or_isi>0]\n",
    "for key, grp in df_stim.groupby(parameter_name):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_travel_histrogram(grp,analysis_methods,key,parameter_name)\n",
    "    #plot_circular_histrogram(grp,analysis_methods,key,parameter_name)\n",
    "#df_isi=df_con[stim_or_isi==0]\n",
    "df_isi=df_interest[stim_or_isi==0]\n",
    "for key, grp in df_isi.groupby(parameter_name):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_travel_histrogram(grp,analysis_methods,key,parameter_name)\n",
    "    #plot_circular_histrogram(grp,analysis_methods,key,parameter_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.5: plot individual animal's trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_iterator=[]\n",
    "if len(vr_no)>0:\n",
    "    print(\"i am using list\")\n",
    "    dir_iterator=zip(dir_list,vr_no)\n",
    "elif type(dir_list)==dict:\n",
    "    print(\"i am using dictionary\")\n",
    "    dir_iterator=dir_dict\n",
    "else:\n",
    "    print(\"there is a bug\")\n",
    "animal_count=0\n",
    "for this_dir,this_vr in dir_iterator:\n",
    "    if Path(this_dir).is_dir()==False:\n",
    "        #print(f'no such a dir exist {this_dir}')\n",
    "        continue\n",
    "    locust_pattern = f\"VR{this_vr}*XY.h5\"\n",
    "    found_result = find_file(Path(this_dir), locust_pattern)        \n",
    "    print(found_result)\n",
    "    #if str(found_result).endswith('VR2_2024-10-14_134515_XY.h5'):\n",
    "    if str(found_result).endswith('VR4_2024-10-13_184515_XY.h5'):\n",
    "        print('use this animal as example')\n",
    "        dfxy = pd.read_hdf(found_result)\n",
    "        dfxy['VR'] = np.tile(f\"VR{this_vr}\", (len(dfxy), 1))\n",
    "        dfxy['VR'] =dfxy[\"VR\"]+\"_\"+dfxy[\"fname\"]\n",
    "        summary_pattern = f\"VR{this_vr}*score.h5\"\n",
    "        found_result = find_file(Path(this_dir), summary_pattern)        \n",
    "        df = pd.read_hdf(found_result)\n",
    "        df['VR'] = np.tile(f\"VR{this_vr}\", (len(df), 1))\n",
    "        df['VR'] =df[\"VR\"]+\"_\"+df[\"fname\"]\n",
    "        if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "            stim_or_isi=dfxy['radial_distance']\n",
    "        elif analysis_methods.get(\"experiment_name\")==\"swarm\" or analysis_methods.get(\"experiment_name\")==\"band\":\n",
    "            stim_or_isi=dfxy['density']\n",
    "        df_stim=dfxy.loc[(dfxy['VR'].isin(df['VR'])) & (stim_or_isi>0)]\n",
    "        for key, grp in df_stim.groupby(parameter_name):\n",
    "            print(f\"kappa:{key},animal_id:{animal_count}\")\n",
    "            plot_sercantrajec(grp,analysis_methods,key,parameter_name,300)\n",
    "        df_isi=dfxy.loc[(dfxy['VR'].isin(df['VR'])) & (stim_or_isi==0)]\n",
    "        for key, grp in df_isi.groupby(parameter_name):\n",
    "            print(f\"kappa:{key}\")\n",
    "            plot_sercantrajec(grp,analysis_methods,key,parameter_name,300)\n",
    "        animal_count=animal_count+1\n",
    "    else:\n",
    "        animal_count=animal_count+1\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#this is used to define active animals \n",
    "travel_distance_across_animals=np.ones(len(df_all))\n",
    "for id in np.arange(len(df_all)):\n",
    "    this_df=df_all[id]\n",
    "    travel_distance_across_animals[id]=this_df[this_df['loss']< 0.05]['distTotal'].sum()\n",
    "plt.hist(travel_distance_across_animals,bins=30)\n",
    "plt.show()\n",
    "## if an animal dont move at all, default analysis pipeline will out put 3 body length (12 cm) per trial, hence the minimum travel distance for an experiment is 12*(trial+ISI number)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#note PI and OI during ISI does not make sense \n",
    "##next step, check animals' response across stim type, across trials\n",
    "save_output= analysis_methods.get(\"save_output\")\n",
    "num_example_animal=22\n",
    "all_OIs=np.ones((2,len(dfxy_all)))\n",
    "all_PIs=np.ones((2,len(dfxy_all)))\n",
    "all_PIs_follow_only=np.ones((2,len(dfxy_all)))\n",
    "all_tortuosity=np.ones((4,len(dfxy_all)))## return NaN, if animals make one or less than one move; and if inactive animals are excluded from analysis\n",
    "# if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "#     stim_or_isi=dfxy_con['radial_distance']\n",
    "# elif analysis_methods.get(\"experiment_name\")==\"swarm\" or analysis_methods.get(\"experiment_name\")==\"band\":\n",
    "#     stim_or_isi=dfxy_con['density']\n",
    "active_animal_threshold=2500\n",
    "animal_count=0\n",
    "for id in np.arange(len(dfxy_all)):\n",
    "    if analysis_methods.get(\"active_animals_only\",False):\n",
    "        this_animal=df_all[id]\n",
    "        if this_animal[this_animal['loss']< 0.05]['distTotal'].sum()>active_animal_threshold:\n",
    "            pass\n",
    "        else:\n",
    "            all_OIs[:,animal_count]=np.nan\n",
    "            all_PIs[:,animal_count]=np.nan\n",
    "            all_tortuosity[:,animal_count]=np.nan\n",
    "            animal_count=animal_count+1\n",
    "            continue\n",
    "    else:\n",
    "        pass\n",
    "    optomotor_stim=[]\n",
    "    preference_stim=[]\n",
    "    labels_stim=[]\n",
    "    optomotor_isi=[]\n",
    "    preference_isi=[]\n",
    "    labels_isi=[]\n",
    "    tortuosity_stim=[]\n",
    "    tortuosity_isi=[]\n",
    "    preference_stim_follow_of=[]\n",
    "    preference_isi_follow_of=[]\n",
    "    trial_id=0\n",
    "    for key, grp in dfxy_all[id].groupby('fname'):\n",
    "        if grp['density'][0]>0 or trial_id==0:## during trial or during pre-stim background, use default mu\n",
    "            this_mu=grp['mu'].unique()[0]\n",
    "        else:## during ISI, continue to use mu from previous trial to check how long navigation direction persist \n",
    "            pass\n",
    "        this_hdf_file=dfxy_all[id].iloc[0]['VR']\n",
    "        #generate labels of classification, index for that trials\n",
    "        l,oi,pi,pi_follow_of_only=classify_heading_direction(grp['heading'].values,this_mu)\n",
    "        #calculate tortuosity\n",
    "        arc=np.sqrt(np.square(np.diff(grp['X'].values)) + np.square(np.diff(grp['Y'].values)))\n",
    "        chord=np.sqrt(np.square(grp['X'].values[-1]-grp['X'].values[0])+np.square(grp['Y'].values[-1]-grp['Y'].values[0]))\n",
    "        if chord>0:\n",
    "            this_tortuosity=np.sum(arc)/chord\n",
    "        else:\n",
    "            this_tortuosity=np.nan\n",
    "        if grp['density'][0]>0:\n",
    "            optomotor_stim.append(oi)\n",
    "            preference_stim.append(pi)\n",
    "            preference_stim_follow_of.append(pi_follow_of_only)\n",
    "            tortuosity_stim.append((np.sum(arc),chord))\n",
    "            labels_stim.append(l)\n",
    "            fig_title=f'{this_hdf_file}_trial{trial_id}_stim_{this_mu}'\n",
    "        else:\n",
    "            optomotor_isi.append(oi)\n",
    "            preference_isi.append(pi)\n",
    "            preference_isi_follow_of.append(pi_follow_of_only)\n",
    "            tortuosity_isi.append((np.sum(arc),chord))\n",
    "            labels_isi.append(l)\n",
    "            fig_title=f'{this_hdf_file}_trial{trial_id}_isi_{this_mu}'\n",
    "        if analysis_methods.get(\"plotting_trajectory\",False) and animal_count==num_example_animal:\n",
    "            print(dfxy_all[id].iloc[0]['VR'],this_mu)\n",
    "            if grp['heading'].shape[0]>1:\n",
    "                ax = plt.subplot(111, polar=True)\n",
    "                ax.hist(grp['heading'].values, bins=24, alpha=0.75)\n",
    "                fig2, ax2 = plt.subplots(\n",
    "                nrows=1, ncols=1, figsize=(18, 6), tight_layout=True\n",
    "            )\n",
    "                xy=np.column_stack((grp[\"X\"].values,grp[\"Y\"].values))\n",
    "                seg_no=1\n",
    "                for start, stop in zip(xy[:-1], xy[1:]):\n",
    "                    x, y = zip(start, stop)\n",
    "                    if l[seg_no]=='for_of':\n",
    "                        this_color='b'\n",
    "                    elif l[seg_no]==\"target_ob\":\n",
    "                        this_color='r'\n",
    "                    elif l[seg_no]==\"against_of\":\n",
    "                        this_color='c'\n",
    "                    else:\n",
    "                        this_color='k'\n",
    "                    ax2.plot(x, y, color=this_color,linewidth=1)\n",
    "                    ax2.set(xlim=(-250,250),ylim=(-250,250),aspect=('equal'))\n",
    "                    seg_no=seg_no+1\n",
    "                fig2.suptitle(fig_title)\n",
    "                fig2_name=fig_title+'.png'\n",
    "                if save_output:\n",
    "                    fig2.savefig(fig2_name)\n",
    "                plt.show()\n",
    "\n",
    "        trial_id=trial_id+1\n",
    "\n",
    "    if analysis_methods.get(\"plotting_trajectory\",False) and animal_count==num_example_animal:\n",
    "        decision_count = np.array([these_labels.size for these_labels in labels_stim])\n",
    "        alpha_values=np.ones(decision_count.size)\n",
    "        alpha_values[(decision_count < 10)] = 0.2\n",
    "        alpha_values[(decision_count >= 10) & (decision_count < 20)] = 0.6\n",
    "        alpha_values[(decision_count >= 20) & (decision_count < 30)] = 0.8\n",
    "        alpha_values[(decision_count >= 30)] = 1\n",
    "        fig1, axes = plt.subplots(\n",
    "            nrows=4, ncols=3, figsize=(12, 9), tight_layout=True\n",
    "        )\n",
    "        ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9,ax10,ax11,ax12= axes.flatten()\n",
    "        ax1.hist(optomotor_stim,color='k')\n",
    "        ax1.hist(optomotor_isi,alpha = 0.05,color='k',ls='dashed',lw=3)\n",
    "        ax2.hist(preference_stim_follow_of,color='k')\n",
    "        ax2.hist(preference_isi_follow_of,alpha = 0.05,color='k',ls='dashed',lw=3)\n",
    "        ax3.hist(preference_stim,color='k')\n",
    "        ax3.hist(preference_isi,alpha = 0.05,color='k',ls='dashed',lw=3)\n",
    "        ax4.scatter(np.array(optomotor_isi),np.array(optomotor_stim),color='k',alpha=alpha_values)\n",
    "        ax5.scatter(np.array(preference_isi_follow_of),np.array(preference_stim_follow_of),color='k',alpha=alpha_values)\n",
    "        ax6.scatter(np.array(preference_isi),np.array(preference_stim),color='k',alpha=alpha_values)\n",
    "        ax7.scatter(np.array(optomotor_stim)[::2],np.array(optomotor_stim)[1::2],color='k',alpha=alpha_values[1::2])\n",
    "        ax8.scatter(np.array(preference_stim_follow_of)[::2],np.array(preference_stim_follow_of)[1::2],color='k',alpha=alpha_values[1::2])\n",
    "        ax9.scatter(np.array(preference_stim)[::2],np.array(preference_stim)[1::2],color='k',alpha=alpha_values[1::2])\n",
    "        ax10.scatter(np.array(optomotor_stim)[1:-1:2],np.array(optomotor_stim)[3::2],color='k',alpha=alpha_values[1:-1:2])\n",
    "        ax11.scatter(np.array(preference_stim_follow_of)[1:-1:2],np.array(preference_stim_follow_of)[3::2],color='k',alpha=alpha_values[1:-1:2])\n",
    "        ax12.scatter(np.array(preference_stim)[1:-1:2],np.array(preference_stim)[3::2],color='k',alpha=alpha_values[1:-1:2])\n",
    "        ax1.set(\n",
    "        xlim=(-1.2,1.2),\n",
    "        xlabel='optomotor index',\n",
    "        ylabel='count of decision')\n",
    "        ax2.set(\n",
    "        xlim=(-1.2,1.2),\n",
    "        xlabel='preference index (follow of vs. target ob)',\n",
    "        ylabel='count of decision')\n",
    "        ax3.set(\n",
    "        xlim=(-1.2,1.2),\n",
    "        xlabel='preference index',\n",
    "        ylabel='count of decision')\n",
    "        ax4.set(\n",
    "        xlabel='optomotor index ISI',\n",
    "        ylabel='optomotor index following Stim',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax5.set(\n",
    "        xlabel='preference index (follow of vs. target ob) ISI',\n",
    "        ylabel='preference index (follow of vs. target ob) following Stim',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax6.set(\n",
    "        xlabel='preference index ISI',\n",
    "        ylabel='preference index following Stim',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax7.set(\n",
    "        xlabel='optomotor index Stim n',\n",
    "        ylabel='optomotor index Stim n+1',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax8.set(\n",
    "        xlabel='preference index (follow of vs. target ob) Stim n',\n",
    "        ylabel='preference index (follow of vs. target ob) Stim n+1',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax9.set(\n",
    "        xlabel='preference index Stim n',\n",
    "        ylabel='preference index Stim n+1',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax10.set(\n",
    "        xlabel='optomotor index Stim n',\n",
    "        ylabel='optomotor index following ISI',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax11.set(\n",
    "        xlabel='preference index (follow of vs. target ob) Stim n',\n",
    "        ylabel='preference index (follow of vs. target ob) following ISI',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        ax12.set(\n",
    "        xlabel='preference index Stim n',\n",
    "        ylabel='preference index following ISI',\n",
    "        xlim=(-1.2,1.2),\n",
    "        ylim=(-1.2,1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'))\n",
    "        #fig1_title=f'animal_no{animal_count}'\n",
    "        #fig1.suptitle(fig1_title)\n",
    "        fig1_name=f'{this_hdf_file}_tbt_index_hist.svg'\n",
    "        if save_output:\n",
    "            fig1.savefig(fig1_name)\n",
    "        plt.show()\n",
    "    ## The OI and PI here are calculated based on responses pooled from every movement across trials\n",
    "    labels_across_trials=np.concat(labels_stim)\n",
    "    arc_chord_aba=np.concat(tortuosity_stim)\n",
    "    all_tortuosity[0,animal_count]=np.nanmean(arc_chord_aba[::2]/arc_chord_aba[1::2])\n",
    "    all_tortuosity[1,animal_count]=np.nanstd(arc_chord_aba[::2]/arc_chord_aba[1::2])/np.sqrt(np.count_nonzero(arc_chord_aba[1::2]!=0))\n",
    "    of_responses=sum(labels_across_trials==\"for_of\")+sum(labels_across_trials==\"against_of\")\n",
    "    all_OIs[0,animal_count]=(sum(labels_across_trials==\"for_of\")-sum(labels_across_trials==\"against_of\"))/of_responses\n",
    "    all_PIs[0,animal_count]=(of_responses-sum(labels_across_trials==\"target_ob\"))/(of_responses+sum(labels_across_trials==\"target_ob\"))\n",
    "    all_PIs_follow_only[0,animal_count]=(sum(labels_across_trials==\"for_of\")-sum(labels_across_trials==\"target_ob\"))/(sum(labels_across_trials==\"for_of\")+sum(labels_across_trials==\"target_ob\"))\n",
    "    labels_across_trials=np.concat(labels_isi)\n",
    "    arc_chord_aba=np.concat(tortuosity_isi)\n",
    "    all_tortuosity[2,animal_count]=np.nanmean(arc_chord_aba[::2]/arc_chord_aba[1::2])\n",
    "    all_tortuosity[3,animal_count]=np.nanstd(arc_chord_aba[::2]/arc_chord_aba[1::2])/np.sqrt(np.count_nonzero(arc_chord_aba[1::2]!=0))\n",
    "    of_responses=sum(labels_across_trials==\"for_of\")+sum(labels_across_trials==\"against_of\")\n",
    "    all_OIs[1,animal_count]=(sum(labels_across_trials==\"for_of\")-sum(labels_across_trials==\"against_of\"))/of_responses\n",
    "    all_PIs[1,animal_count]=(of_responses-sum(labels_across_trials==\"target_ob\"))/(of_responses+sum(labels_across_trials==\"target_ob\"))\n",
    "    all_PIs_follow_only[1,animal_count]=(sum(labels_across_trials==\"for_of\")-sum(labels_across_trials==\"target_ob\"))/(sum(labels_across_trials==\"for_of\")+sum(labels_across_trials==\"target_ob\"))\n",
    "    animal_count=animal_count+1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_pi_oi_comparison(all_PIs,all_OIs,all_tortuosity,all_PIs_follow_only,analysis_methods,travel_distance_across_animals)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_pi_oi_comparison(all_PIs,all_OIs,all_tortuosity,all_PIs_follow_only,analysis_methods,travel_distance=np.ones(all_OIs.shape[1]),parameters='',parameter_name='mu',vr_num='all'):\n",
    "    save_output= analysis_methods.get(\"save_output\")\n",
    "    scene_name=analysis_methods.get(\"experiment_name\")\n",
    "    if analysis_methods.get(\"active_trials_only\"):\n",
    "        active_trial='active_trials'\n",
    "    else:\n",
    "        active_trial=''\n",
    "    alpha_values=np.ones(all_OIs.shape[1])\n",
    "    if np.sum(travel_distance)>travel_distance.shape[0]:\n",
    "        alpha_values[(travel_distance < 2500)] = 0.2\n",
    "        alpha_values[(travel_distance >= 5000) & (travel_distance < 15000)] = 0.6\n",
    "        #travel_distance[(travel_distance >= 10000) & (travel_distance < 15000)] = 0.8\n",
    "        alpha_values[(travel_distance >= 15000)] = 1\n",
    "    fig, axes = plt.subplots(nrows=3, ncols=2, figsize=(18, 14))\n",
    "    ax1,ax2,ax3,ax4,ax5,ax6= axes.flatten()\n",
    "    ax1.scatter(all_PIs[1],all_tortuosity[2],c='g',alpha=alpha_values,label='ISIs')\n",
    "    ax1.errorbar(all_PIs[1],all_tortuosity[2],yerr=all_tortuosity[3],c='g',fmt='o')\n",
    "    ax1.scatter(all_PIs[0],all_tortuosity[0],c='k',alpha=alpha_values,label='Trials')##stim\n",
    "    ax1.errorbar(all_PIs[0],all_tortuosity[0],yerr=all_tortuosity[1],c='k',fmt='o')\n",
    "    ax2.scatter(all_OIs[1],all_tortuosity[2],c='g',label='ISIs')\n",
    "    ax2.errorbar(all_OIs[1],all_tortuosity[2],yerr=all_tortuosity[3],c='g',fmt='o')\n",
    "    ax2.scatter(all_OIs[0],all_tortuosity[0],c='k',label='Trials')##stim\n",
    "    ax2.errorbar(all_OIs[0],all_tortuosity[0],yerr=all_tortuosity[1],c='k',fmt='o')\n",
    "    ax3.scatter(all_OIs[1],all_PIs[1],c='g',alpha=alpha_values,label='ISIs')\n",
    "    ax3.scatter(all_OIs[0],all_PIs[0],c='k',alpha=alpha_values,label='Trials')##stim\n",
    "    ax4.scatter(all_PIs[0],all_PIs[1],c='b',alpha=alpha_values,label='preference index')\n",
    "    ax4.scatter(all_OIs[0],all_OIs[1],c='r',alpha=alpha_values,label='optomotor index')\n",
    "    ax5.scatter(all_OIs[1],all_PIs_follow_only[1],c='g',alpha=alpha_values,label='ISIs')\n",
    "    ax5.scatter(all_OIs[0],all_PIs_follow_only[0],c='k',alpha=alpha_values,label='Trials')##stim\n",
    "    ax6.scatter(all_PIs_follow_only[0],all_PIs_follow_only[1],c='b',alpha=alpha_values,label='(follow OF vs. Track OB)')\n",
    "    ax6.scatter(all_OIs[0],all_OIs[1],c='r',alpha=alpha_values,label='optomotor index')\n",
    "    ax1.set(\n",
    "        xlim=(-1.2, 1.2),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        xlabel='preference index',\n",
    "        ylabel='tortuosity',\n",
    "    )\n",
    "    ax2.set(\n",
    "        xlim=(-1.2, 1.2),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        xlabel='optomotor index',\n",
    "        ylabel='tortuosity',\n",
    "    )\n",
    "    ax3.set(\n",
    "        xlim=(-1.2, 1.2),\n",
    "        ylim=(-1.2, 1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        xlabel='optomotor index',\n",
    "        ylabel='preference index',\n",
    "        aspect=('equal'),\n",
    "    )\n",
    "    ax4.set(\n",
    "        xlim=(-1.2, 1.2),\n",
    "        ylim=(-1.2, 1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'),\n",
    "        xlabel='index during trials',\n",
    "        ylabel='index during ISIs',\n",
    "    )\n",
    "    ax5.set(\n",
    "        xlim=(-1.2, 1.2),\n",
    "        ylim=(-1.2, 1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        xlabel='optomotor index',\n",
    "        ylabel='preference index (follow OF vs. Track OB)',\n",
    "        aspect=('equal'),\n",
    "    )\n",
    "    ax6.set(\n",
    "        xlim=(-1.2, 1.2),\n",
    "        ylim=(-1.2, 1.2),\n",
    "        yticks=([-1, 0, 1]),\n",
    "        xticks=([-1, 0, 1]),\n",
    "        aspect=('equal'),\n",
    "        xlabel='index during trials',\n",
    "        ylabel='index during ISIs',\n",
    "    )\n",
    "    ax1.legend(loc='upper left')\n",
    "    ax2.legend(loc='upper left')\n",
    "    ax3.legend(loc='upper left')\n",
    "    ax4.legend(loc='upper left')\n",
    "    ax5.legend(loc='upper left')\n",
    "    ax6.legend(loc='upper left')\n",
    "    fig.suptitle('+ pi means prefer optic flow, - means prefer target,; + oi means follow optic flow, - means go against optic flow')\n",
    "    fig_name=f\"{vr_num}_oi_pi_comparision_{scene_name}_{parameter_name}_{parameters}{active_trial}.svg\"\n",
    "    if save_output:\n",
    "        fig.savefig(fig_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.6: plot individual animal's kernel density estimation plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_iterator=[]\n",
    "if len(vr_no)>0:\n",
    "    print(\"i am using list\")\n",
    "    dir_iterator=zip(dir_list,vr_no)\n",
    "elif type(dir_list)==dict:\n",
    "    print(\"i am using dictionary\")\n",
    "    dir_iterator=dir_dict\n",
    "else:\n",
    "    print(\"there is a bug\")\n",
    "\n",
    "for this_dir,this_vr in dir_iterator:\n",
    "    if Path(this_dir).is_dir()==False:\n",
    "        print(f'no such a dir exist {this_dir}')\n",
    "        continue\n",
    "    locust_pattern = f\"VR{this_vr}*score.h5\"\n",
    "    found_result = find_file(Path(this_dir), locust_pattern)        \n",
    "    df = pd.read_hdf(found_result)\n",
    "    if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "        stim_or_isi=df['radial_distance']\n",
    "    elif analysis_methods.get(\"experiment_name\")==\"swarm\" or analysis_methods.get(\"experiment_name\")==\"band\":\n",
    "        stim_or_isi=df['density']\n",
    "    df_stim=df[stim_or_isi>0]\n",
    "    for key, grp in df_stim.groupby(parameter_name):\n",
    "        print(f\"{parameter_name}:{key}\")\n",
    "        plot_sercansincos(grp,analysis_methods,key,parameter_name)\n",
    "    df_isi=df[stim_or_isi==0]\n",
    "    for key, grp in df_isi.groupby(parameter_name):\n",
    "        print(f\"{parameter_name}:{key}\")\n",
    "        plot_sercansincos(grp,analysis_methods,key,parameter_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Optional] Analyse data with multi-engines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##this cell start the multi-engines. Make sure to run only once\n",
    "import time\n",
    "import ipyparallel as ipp\n",
    "def show_clusters():\n",
    "    clusters = ipp.ClusterManager().load_clusters() \n",
    "    print(\"{:15} {:^10} {}\".format(\"cluster_id\", \"state\", \"cluster_file\")) \n",
    "    for c in clusters:\n",
    "        cd = clusters[c].to_dict()\n",
    "        cluster_id = cd['cluster']['cluster_id']\n",
    "        controller_state = cd['controller']['state']['state']\n",
    "        cluster_file = getattr(clusters[c], '_trait_values')['cluster_file']\n",
    "        print(\"{:15} {:^10} {}\".format(cluster_id, controller_state, cluster_file))\n",
    "    return cluster_id\n",
    "\n",
    "cluster = ipp.Cluster(n=6)\n",
    "await cluster.start_cluster()\n",
    "cluster_neuropc=show_clusters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##input cluster_id from previous cell\n",
    "rc = ipp.Client(cluster_id=cluster_neuropc)\n",
    "\n",
    "# Create a DirectView for parallel execution\n",
    "dview = rc.direct_view()\n",
    "\n",
    "# Define a function for parallel processing\n",
    "def process_directory(this_dir, analysis_methods):\n",
    "    from pathlib import Path\n",
    "    import sys\n",
    "    current_working_directory = Path.cwd()\n",
    "    parent_dir = current_working_directory.resolve().parents[0]\n",
    "    sys.path.insert(0, str(parent_dir) + \"\\\\utilities\")\n",
    "    from locustvr_converter import preprocess_matrex_data\n",
    "    preprocess_matrex_data(this_dir,analysis_methods)\n",
    "\n",
    "# Define analysis_methods\n",
    "\n",
    "# Use parallel execution to process directories\n",
    "dview.map_sync(process_directory, dir_list, [analysis_methods] * len(dir_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rc.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "matrexvr_analysis",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}