data_exploration_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 0.0: Load packages and customised functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "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",
    "##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\n",
    "#from data_cleaning import preprocess_fictrac_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "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 = \"coolwarm\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 0.1: Load analysis methods in python dictionary form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "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\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 0.2: Load animals' experiment directory into a list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "these directories are found ['D:/MatrexVR_grass1_Data/RunData/20240907_142802', 'D:/MatrexVR_grass1_Data/RunData/20240907_170446', 'D:/MatrexVR_grass1_Data/RunData/20240907_190839', 'D:/MatrexVR_grass1_Data/RunData/20240908_125638', 'D:/MatrexVR_grass1_Data/RunData/20240908_150715', 'D:/MatrexVR_grass1_Data/RunData/20240908_174232', 'D:/MatrexVR_grass1_Data/RunData/20240908_193840']\n"
     ]
    }
   ],
   "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",
    "\n",
    "print(f\"these directories are found {dir_list}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "tmp_file_name='DL220THP_Thermo3_240904_240908.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 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: introduce customised plotting functions used in Sercan's paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "#introduce customised functions\n",
    "def plot_sercansincos(df,parameters,parameter_name,vr_num='all'):\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_{parameter_name}_{parameters}_density_{int(density)}.svg\"\n",
    "        sin_fig_name=f\"{vr_num}_sin_{parameter_name}_{parameters}_density_{int(density)}.svg\" \n",
    "    else:\n",
    "        cos_fig_name=f\"{vr_num}_cos_{parameter_name}_{parameters}_single_target.svg\"\n",
    "        sin_fig_name=f\"{vr_num}_sin_{parameter_name}_{parameters}_single_target.svg\"\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.set_cmap('cividis')\n",
    "\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.rcParams['font.family'] = 'Arial'\n",
    "    sns.kdeplot(cos, cut=0, color=\"#21918c\", fill=True, alpha=0.9)#)#, lw=1\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",
    "    plt.rcParams.update({'font.size': 8})\n",
    "    plt.subplots_adjust(bottom=0.4)\n",
    "    plt.savefig(cos_fig_name)\n",
    "    plt.show()   \n",
    "\n",
    "    fig, ax = plt.subplots(dpi=300, figsize=(1.1,0.25))\n",
    "    plt.rcParams['ytick.major.width'] = 2\n",
    "    plt.rcParams['xtick.major.width'] = 2\n",
    "    plt.rcParams['axes.linewidth'] = 2\n",
    "    plt.rcParams['font.family'] = 'Arial'\n",
    "    sns.kdeplot(sin, cut=0, color=\"#21918c\",  fill=True, alpha=0.9)#),lw=1,\n",
    "    plt.xlim(1,-1)\n",
    "\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['left'].set_visible(False)\n",
    "\n",
    "    plt.xticks(rotation = 90)\n",
    "    ax.set_yticks([])\n",
    "    plt.ylabel(\"\")\n",
    "    plt.xlabel(\"\")\n",
    "    plt.title(\"r sin\\u03F4\")\n",
    "\n",
    "    plt.rcParams.update({'font.size': 8})\n",
    "\n",
    "    plt.subplots_adjust(bottom=0.4)\n",
    "    plt.savefig(sin_fig_name)\n",
    "    plt.show()\n",
    "def plot_sercantrajec(dfXY,parameters,parameter_name,trajec_lim=1000,vr_num='all'):\n",
    "    dfXY[parameter_name].unique()\n",
    "    \n",
    "    a = dfXY.groupby('VR')\n",
    "    print(a)\n",
    "    if 'density' in dfXY.columns:\n",
    "        density=dfXY[\"density\"].unique()[0]\n",
    "        print(density)\n",
    "        fig_name=f\"{vr_num}_summary_trajectory_{parameter_name}_{parameters}_density_{int(density)}.png\"\n",
    "    else:\n",
    "        fig_name=f\"{vr_num}_summary_trajectory_{parameter_name}_{parameters}_single_target.png\"\n",
    "    simulated_speed=2\n",
    "    initial_distance=8\n",
    "    duration=60\n",
    "    travel_direction=parameters*-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",
    "    #travel_direction=travel_direction-np.pi/2 \n",
    "    initial_distance_b=np.cos(travel_direction)*initial_distance\n",
    "    delta_cos=np.cumsum(np.repeat(np.cos(travel_direction)*simulated_speed, duration))\n",
    "    distance_cos=initial_distance_b+delta_cos\n",
    "    initial_distance_b=np.sin(travel_direction)*initial_distance\n",
    "    delta_sin=np.cumsum(np.repeat(np.sin(travel_direction)*simulated_speed, duration))\n",
    "    distance_sin=initial_distance_b+delta_sin\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",
    "    # 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.rcParams['font.family'] = 'Helvetica'\n",
    "    # Get the colormap\n",
    "    cmap = plt.get_cmap('viridis')\n",
    "    # Get the total number of trajectories\n",
    "    n = len(a)\n",
    "\n",
    "    #plt.style.use('dark_background') \n",
    "    for i, (key2, grp2) in enumerate(a):\n",
    "        xx = grp2[\"X\"].values\n",
    "        yy = grp2[\"Y\"].values\n",
    "        color = cmap(i/n)\n",
    "        plt.plot(xx, yy, color=color, linewidth=1)\n",
    "\n",
    "\n",
    "        # Calculate angles, radii, etc. (your existing code)\n",
    "\n",
    "    plt.plot(distance_cos, distance_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",
    "    plt.savefig(fig_name)\n",
    "    # Set the aspect ratio to be equal\n",
    "    plt.gca().set_aspect('equal')\n",
    "    plt.show()\n",
    "def plot_travel_distance_set(df_all,analysis_methods,parameter_name,y_axis_lim=[0.1,1000]):\n",
    "    COL = MplColorHelper(colormap_name, 0, 10)\n",
    "    fig, (ax1, ax2,ax3) = plt.subplots(\n",
    "    nrows=1, ncols=3, figsize=(18, 7), tight_layout=True\n",
    ")\n",
    "    colour_code=analysis_methods.get(\"graph_colour_code\")\n",
    "    if parameter_name=='order':\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=\"Change in Travel distance (ratio)\",\n",
    "        xticks=sorted(df_all[0][parameter_name].unique()),\n",
    "        xlabel=parameter_name,\n",
    "    )\n",
    "    ax2.set(\n",
    "        ylabel=\"trial n travel distance\",\n",
    "        xlabel=\"pre-stim interval n travel distance\",\n",
    "    )\n",
    "    ax3.set(\n",
    "        ylabel=\"trial n travel distance\",\n",
    "        xlabel=\"trial n-1 travel distance\",\n",
    "    )\n",
    "    for id in np.arange(len(df_all)):\n",
    "    #for this_dir,this_vr in dir_dict:\n",
    "        df=df_all[id]\n",
    "\n",
    "        #set some thresholds to remove bad tracking \n",
    "        df.loc[(df[\"distTotal\"]<10.0) | (df[\"loss\"]> 0.05), \"distTotal\"] = np.nan\n",
    "        if df.iloc[0][\"VR\"].startswith('VR1'):\n",
    "            this_color=colour_code[0]\n",
    "        elif df.iloc[0][\"VR\"].startswith('VR2'):\n",
    "            this_color=colour_code[1]\n",
    "        elif df.iloc[0][\"VR\"].startswith('VR3'):\n",
    "            this_color=colour_code[2]\n",
    "        else:\n",
    "            this_color=colour_code[3]\n",
    "        #ax1.scatter(df.iloc[1::2][parameter_name], df[1::2][\"distTotal\"]/df[::2][\"distTotal\"],c=this_color)\n",
    "        #ax2.scatter(df[::2][\"distTotal\"],df.iloc[1::2][\"distTotal\"],c=this_color,alpha=df.iloc[1::2]['alpha'])\n",
    "        # below is a quick check when pre-stim interval is 5 min long and isi is 1 min and when there is a post-stim interval\n",
    "        ax1.scatter(df.iloc[3::2][parameter_name], df[3::2][\"distTotal\"]/df[2:-1:2][\"distTotal\"],c=this_color)\n",
    "        ax1.scatter(df.iloc[1][parameter_name], df.iloc[1][\"distTotal\"]/df.iloc[0][\"distTotal\"]*5,c=this_color)\n",
    "        # ax2.scatter(df.iloc[0][\"distTotal\"]/5,df.iloc[1][\"distTotal\"],c=this_color,alpha=df.iloc[1]['alpha']) \n",
    "        # ax2.scatter(df[2:-1:2][\"distTotal\"],df.iloc[3::2][\"distTotal\"],c=this_color,alpha=df.iloc[3::2]['alpha'])\n",
    "        # ax3.scatter(df.iloc[1:-2:2][\"distTotal\"], df[3::2][\"distTotal\"],c=this_color,alpha=df.iloc[3::2]['alpha'])\n",
    "        ##below are some commands to set samples from the same parameters the same colour\n",
    "        ax2.scatter(df.iloc[0][\"distTotal\"]/5,df.iloc[1][\"distTotal\"],c=np.array([COL.get_rgb(int(df.iloc[1]['alpha']*10))]))\n",
    "        ax2.scatter(df[2:-1:2][\"distTotal\"],df.iloc[3::2][\"distTotal\"],c=np.array([COL.get_rgb(int(df.iloc[1]['alpha']*10))]))\n",
    "        ax3.scatter(df.iloc[1:-2:2][\"distTotal\"], df[3::2][\"distTotal\"],c=np.array([COL.get_rgb(int(df.iloc[1]['alpha']*10))]))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.1: 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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the path to your Excel file\n",
    "dir_list = []\n",
    "file_type=\".h5\"\n",
    "using_google_sheet=True\n",
    "sheet_name = \"Unity_MatrexVR\"\n",
    "experiment_name=analysis_methods.get(\"experiment_name\")\n",
    "# if type(thisDataset) == str:\n",
    "#     thisDataset = Path(thisDataset)\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_black\")\n",
    "        #print(animal_of_interest)\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 = itertools.zip_longest(dir_list, vr_no.tolist())\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": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def connect_two_tables(dir_list,test_parameter='order',vr_no=[]):\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",
    "        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",
    "        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",
    "        if test_parameter == 'order':\n",
    "            alpha_converter={0.1: 0.2, 1.0: 0.4, 10.0: 0.6,100000.0:1}\n",
    "        elif test_parameter == 'mu':\n",
    "            #alpha_converter={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",
    "            alpha_converter={0: 0.1,45: 0.4,315: 0.7}\n",
    "        elif test_parameter == 'agent_speed':\n",
    "            alpha_converter={1.0: 0.2,2.0: 0.4, 4.0: 0.6,8.0:1}\n",
    "        elif test_parameter=='three_direction':\n",
    "            test_parameter='mu'\n",
    "            alpha_converter={0: 0.1,45: 0.2,315: 0.3}\n",
    "        print(alpha_converter)\n",
    "        print(df[test_parameter])\n",
    "        df['alpha'] = df[test_parameter].map(alpha_converter)\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": [
    "#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",
    "            simulated_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": [
    "#parameter_name='order'\n",
    "parameter_name='mu'\n",
    "#parameter_name='agent_speed'\n",
    "if 'vr_no' in locals():\n",
    "    df_all,dfxy_all=connect_two_tables(dir_list,parameter_name,vr_no)\n",
    "else:\n",
    "    df_all,dfxy_all=connect_two_tables(dir_dict,parameter_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#agent movement\n",
    "df_agent = pd.read_hdf(r'D:\\MatrexVR_blackbackground_Data\\RunData\\20240904_151537\\VR2_2024-09-04_151654_agent.h5')\n",
    "stim_groups = df_agent.groupby('mu')\n",
    "trajec_lim=500\n",
    "for stim_no, entries in stim_groups:\n",
    "    print(stim_no)\n",
    "    fig, ax = plt.subplots(figsize=(3,3), dpi=300) \n",
    "    plt.rcParams.update(plt.rcParamsDefault)\n",
    "    plt.rcParams.update({'font.size': 8})\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.rcParams['font.family'] = 'Helvetica'\n",
    "    # Get the colormap\n",
    "    cmap = plt.get_cmap('viridis')\n",
    "    # Get the total number of trajectories\n",
    "    #n = len(a)\n",
    "\n",
    "    #plt.style.use('dark_background')\n",
    "    plt.plot(entries['X'], entries['Y'],color='k',linewidth=1)\n",
    "\n",
    "    #plt.plot(xx, yy, color=color, linewidth=1)\n",
    "\n",
    "\n",
    "        # Calculate angles, radii, etc. (your existing code)\n",
    "\n",
    "    #plt.plot(distance_cos, distance_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",
    "    #plt.savefig(fig_name)\n",
    "    # Set the aspect ratio to be equal\n",
    "    plt.gca().set_aspect('equal')\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3: make summary plots of animals' response with customised functions introduced before."
   ]
  },
  {
   "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.1,1000])"
   ]
  },
  {
   "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",
    "parameter_name='mu'\n",
    "y_axis_lim=[0.1,1000]\n",
    "COL = MplColorHelper(colormap_name, 0, 8)\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=='order':\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=\"Change in Travel distance (ratio)\",\n",
    "    xticks=sorted(df_all[0][parameter_name].unique()),\n",
    "    xlabel=parameter_name,\n",
    ")\n",
    "ax2.set_ylim(0,20)\n",
    "ax2.set_xlim(0,20)\n",
    "ax2.set(\n",
    "    ylabel=\"trial n average speed\",\n",
    "    xlabel=\"pre-stim interval n average speed\",\n",
    ")\n",
    "ax3.set_ylim(0,750)\n",
    "ax3.set_xlim(0,750)\n",
    "ax3.set(\n",
    "    ylabel=\"trial n travel distance\",\n",
    "    xlabel=\"trial n-1 travel distance\",\n",
    ")\n",
    "for id in np.arange(len(df_all)):\n",
    "#for this_dir,this_vr in dir_dict:\n",
    "    df=df_all[id]\n",
    "\n",
    "    #set some thresholds to remove bad tracking \n",
    "    df.loc[(df[\"distTotal\"]<10.0) | (df[\"loss\"]> 0.05), \"distTotal\"] = np.nan\n",
    "    if df.iloc[0][\"VR\"].startswith('VR1'):\n",
    "        this_color=colour_code[0]\n",
    "    elif df.iloc[0][\"VR\"].startswith('VR2'):\n",
    "        this_color=colour_code[1]\n",
    "    elif df.iloc[0][\"VR\"].startswith('VR3'):\n",
    "        this_color=colour_code[2]\n",
    "    else:\n",
    "        this_color=colour_code[3]\n",
    "    #ax1.scatter(df.iloc[1::2][parameter_name], df[1::2][\"distTotal\"]/df[::2][\"distTotal\"],c=this_color)\n",
    "    #ax2.scatter(df[::2][\"distTotal\"],df.iloc[1::2][\"distTotal\"],c=this_color,alpha=df.iloc[1::2]['alpha'])\n",
    "    # below is a quick check when pre-stim interval is 5 min long and isi is 1 min and when there is a post-stim interval\n",
    "    alpha=df['alpha']*10\n",
    "    alpha=alpha.astype('int')\n",
    "    ax1.scatter(df.iloc[::2][parameter_name], df.iloc[::2][\"distTotal\"]/df.iloc[::2][\"duration\"],c='k')\n",
    "    ax1.scatter(df.iloc[1::2][parameter_name], df.iloc[1::2][\"distTotal\"]/df.iloc[1::2][\"duration\"],c=COL.get_rgb(alpha[1::2]))\n",
    "    #ax1.scatter(df.iloc[3::2][parameter_name], df[3::2][\"distTotal\"]/df[2:-1:2][\"distTotal\"],c=this_color)\n",
    "    #ax1.scatter(df.iloc[1][parameter_name], df.iloc[1][\"distTotal\"]/df.iloc[0][\"distTotal\"]*5,c=this_color)\n",
    "    # ax2.scatter(df.iloc[0][\"distTotal\"]/5,df.iloc[1][\"distTotal\"],c=this_color,alpha=df.iloc[1]['alpha']) \n",
    "    # ax2.scatter(df[2:-1:2][\"distTotal\"],df.iloc[3::2][\"distTotal\"],c=this_color,alpha=df.iloc[3::2]['alpha'])\n",
    "    # ax3.scatter(df.iloc[1:-2:2][\"distTotal\"], df[3::2][\"distTotal\"],c=this_color,alpha=df.iloc[3::2]['alpha'])\n",
    "    ##below are some commands to set samples from the same parameters the same colour\n",
    "    #ax2.scatter(df.iloc[0][\"distTotal\"]/5,df.iloc[1][\"distTotal\"],c=np.array([COL.get_rgb(int(df.iloc[1]['alpha']*10))]))\n",
    "    #print(df.iloc[::2][\"distTotal\"].shape,df.iloc[1::2][\"distTotal\"].shape)\n",
    "    #print(np.array([COL.get_rgb(int(df.iloc[1::2]['alpha']*10))]))\n",
    "    #test=df_all[0][1::2]['alpha'].values*10\n",
    "\n",
    "    #print(\"test\")\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(alpha[1::2]))\n",
    "    #ax.grid(True)\n",
    "    #scatter.colourbar\n",
    "    ax3.scatter(df.iloc[1:-2:2][\"distTotal\"],df.iloc[3::2][\"distTotal\"],c=COL.get_rgb(alpha[3::2]))\n",
    "\n",
    "    # legend1 = scatter.legend(*scatter.legend_elements(),loc=\"upper right\", title=\"Direction\")\n",
    "    # scatter.add_artist(legend1)\n",
    "    #ax2.set_title('different colour means data from different direction')\n",
    "    #ax3.scatter(df.iloc[1:-2:2][\"distTotal\"], df[3::2][\"distTotal\"],c=np.array([COL.get_rgb(int(df.iloc[1]['alpha']*10))]))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plotting trajectory\n",
    "if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "    stim_or_isi=dfxy_con['initial_distance']\n",
    "elif analysis_methods.get(\"experiment_name\")==\"swarm\":\n",
    "    stim_or_isi=dfxy_con['density']\n",
    "df_stim=dfxy_con.loc[(dfxy_con['VR'].isin(df_con[\"VR\"][good_tracking])) & (stim_or_isi>0)]\n",
    "for key, grp in df_stim.groupby(parameter_name):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercantrajec(grp,key,parameter_name,500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#set condition dfxy_con['density']>0 to select stimulus trial\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\n",
    "df_stim=dfxy_con.loc[(dfxy_con['VR'].isin(df_con[\"VR\"][good_tracking])) & (stim_or_isi>0)]\n",
    "for key, grp in df_stim.groupby(parameter_name):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercantrajec(grp,key,parameter_name,500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "duration=60\n",
    "distance_cos=np.repeat(np.cos(travel_direction)*initial_distance,duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_cos=np.cumsum(np.repeat(np.cos(initial_distance)*simulated_speed, 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distance_cos+delta_distance_cos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#set condition dfxy_con['density']==0 to select ISI\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\n",
    "if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "    stim_or_isi=dfxy_con['initial_distance']\n",
    "elif analysis_methods.get(\"experiment_name\")==\"swarm\":\n",
    "    stim_or_isi=dfxy_con['density']\n",
    "df_isi=dfxy_con.loc[(dfxy_con['VR'].isin(df_con[\"VR\"][good_tracking])) & (stim_or_isi==0)]\n",
    "for key, grp in df_isi.groupby(parameter_name):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercantrajec(grp,key,parameter_name,500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis_methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot mean angle\n",
    "if analysis_methods.get(\"experiment_name\")==\"choice\":\n",
    "    stim_or_isi=df_con['initial_distance']\n",
    "elif analysis_methods.get(\"experiment_name\")==\"swarm\":\n",
    "    stim_or_isi=df_con['density']\n",
    "df_stim=df_con[stim_or_isi>0]\n",
    "for key, grp in df_stim.groupby(parameter_name):\n",
    "    print(f\"{parameter_name}:{key}\")\n",
    "    plot_sercansincos(grp,key,parameter_name)\n",
    "# df_isi=df_con[stim_or_isi==0]\n",
    "# for key, grp in df_isi.groupby(parameter_name):\n",
    "#     print(f\"{parameter_name}:{key}\")\n",
    "#     plot_sercansincos(grp,key,parameter_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Session 2.3.2: plot individual animal's response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_iterator=[]\n",
    "if 'vr_no' in locals():\n",
    "    dir_iterator=zip(dir_list,vr_no)\n",
    "elif 'dir_dict' in locals():\n",
    "    dir_iterator=dir_dict \n",
    "    parameter_name='order'\n",
    "    for this_dir,this_vr in dir_iterator:\n",
    "        this_color=colour_code[this_vr]\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",
    "        df.loc[(df[\"distTotal\"]<5) | (df[\"loss\"]> 0.05), \"distTotal\"] = np.nan\n",
    "        df_stim=df[df['density']>0]\n",
    "        df_isi=df[df['density']==0]\n",
    "        for key, grp in df_stim.groupby('order'):\n",
    "            print(f\"order:{key}\")\n",
    "            plot_sercantrajec(grp,key,parameter_name,2000)\n",
    "            plot_sercansincos(grp,key,parameter_name)\n",
    "else:\n",
    "    print(\"should plot every animal in a for-loop, and a nested-for loop. This section will igonre which VR setup they come from\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "extra session: 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": "Python 3 (ipykernel)",
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}