{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show some statistics of the enclosed geometric object with psydat file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "import ipywidgets as widgets\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from ipywidgets import interact\n",
    "from scipy.signal import savgol_filter\n",
    "\n",
    "import vstt\n",
    "from vstt.stats import get_velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import\n",
    "A psydat file can be imported using the psychopy `fromFile` function: \n",
    "If you want to know the detailed content of the data in psydat file, please check the notebook 'raw_data.ipynb'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment = vstt.Experiment(\"example.psydat\")\n",
    "stats = experiment.stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot of results for each trial\n",
    "For example, a scatter plot of the mouse positions for each trial, labelled by the condition, trial number and repetition number:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_to_target(ax, group, colors):\n",
    "    for target_pos, target_radius, positions, color in zip(\n",
    "        group.target_pos, group.target_radius, group.to_target_mouse_positions, colors\n",
    "    ):\n",
    "        ax.plot(positions[:, 0], positions[:, 1], color=color)\n",
    "        ax.add_patch(\n",
    "            plt.Circle(\n",
    "                target_pos,\n",
    "                target_radius,\n",
    "                edgecolor=\"none\",\n",
    "                facecolor=color,\n",
    "                alpha=0.1,\n",
    "            )\n",
    "        )\n",
    "\n",
    "\n",
    "def plot_to_center(ax, group, colors):\n",
    "    for central_target_radius, positions, color in zip(\n",
    "        group.center_radius,\n",
    "        group.to_center_mouse_positions,\n",
    "        colors,\n",
    "    ):\n",
    "        ax.plot(positions[:, 0], positions[:, 1], color=color)\n",
    "        ax.add_patch(\n",
    "            plt.Circle(\n",
    "                [0, 0],\n",
    "                central_target_radius,\n",
    "                edgecolor=\"none\",\n",
    "                facecolor=\"black\",\n",
    "                alpha=0.1,\n",
    "            )\n",
    "        )\n",
    "\n",
    "\n",
    "def concatenate(array1, array2):\n",
    "    return np.concatenate(\n",
    "        (\n",
    "            array1,\n",
    "            array2,\n",
    "        ),\n",
    "        axis=0,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = [\"blue\", \"green\", \"red\", \"cyan\", \"magenta\", \"yellow\", \"black\", \"orange\"]\n",
    "\n",
    "nTrials = len(stats[\"i_trial\"].unique())\n",
    "nReps = len(stats[\"i_rep\"].unique())\n",
    "fig, axs = plt.subplots(nTrials, nReps, figsize=(6, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    ax = axs[trial, rep]\n",
    "    ax.set_title(f\"[Condition {condition_index}] Trial {trial}, Rep {rep}\")\n",
    "    plot_to_target(ax, group, colors)\n",
    "    if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "        plot_to_center(ax, group, colors)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot of results for each target\n",
    "For example, a scatter plot of the mouse positions for each target, labelled by trial number, repetition number, target number and condition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nTrials, nReps * 8, figsize=(6 * 8, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps * 8)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    for positions, color, i in zip(\n",
    "        group.to_target_mouse_positions, colors, range(len(group))\n",
    "    ):\n",
    "        ax = axs[(trial, rep + i)]\n",
    "        ax.set_title(\n",
    "            f\"[Condition {condition_index}] Trial {trial}, Rep {rep}, Target {i}\"\n",
    "        )\n",
    "        ax.set_xlim(-0.5, 0.5)\n",
    "        ax.set_ylim(-0.5, 0.5)\n",
    "        if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "            positions = concatenate(positions, group.to_center_mouse_positions.iloc[i])\n",
    "        ax.plot(positions[:, 0], positions[:, 1], color=color)\n",
    "\n",
    "fig.delaxes(axs[2][6])\n",
    "fig.delaxes(axs[2][7])\n",
    "fig.delaxes(axs[3][6])\n",
    "fig.delaxes(axs[3][7])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot velocity for each trial\n",
    "For example, a scatter plot of the velocity for each target,displayed in a single plot with velocities shown in the time sequence, labelled by trial number, repetition number and condition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nTrials, nReps, figsize=(6, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    ax = axs[trial, rep]\n",
    "    ax.set_title(f\"[Condition {condition_index}] Trial {trial}, Rep {rep} \")\n",
    "    for positions, timestamps, color, i in zip(\n",
    "        group.to_target_mouse_positions,\n",
    "        group.to_target_timestamps,\n",
    "        colors,\n",
    "        range(len(group)),\n",
    "    ):\n",
    "        if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "            positions = concatenate(positions, group.to_center_mouse_positions.iloc[i])\n",
    "            timestamps = concatenate(timestamps, group.to_center_timestamps.iloc[i])\n",
    "        ax.plot(timestamps[:-1], get_velocity(timestamps, positions), color=color)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, a scatter plot of the velocity for each target in separate plot, labelled by trial number, repetition number and condition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nTrials, nReps * 8, figsize=(6 * 8, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps * 8)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    for positions, timestamps, color, i in zip(\n",
    "        group.to_target_mouse_positions,\n",
    "        group.to_target_timestamps,\n",
    "        colors,\n",
    "        range(len(group)),\n",
    "    ):\n",
    "        ax = axs[(trial, rep + i)]\n",
    "        ax.set_title(\n",
    "            f\"[Condition {condition_index}] Trial {trial}, Rep {rep}, Target {i}\"\n",
    "        )\n",
    "        if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "            positions = concatenate(positions, group.to_center_mouse_positions.iloc[i])\n",
    "            timestamps = concatenate(timestamps, group.to_center_timestamps.iloc[i])\n",
    "        ax.plot(timestamps[:-1], get_velocity(timestamps, positions), color=color)\n",
    "\n",
    "\n",
    "fig.delaxes(axs[2][6])\n",
    "fig.delaxes(axs[2][7])\n",
    "fig.delaxes(axs[3][6])\n",
    "fig.delaxes(axs[3][7])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, a scatter plot of the velocity for each target,displayed in a single plot with velocities starting from the same time point 0,labelled by trial number, repetition number and condition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nTrials, nReps, figsize=(6, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    ax = axs[trial, rep]\n",
    "    ax.set_title(f\"[Condition {condition_index}] Trial {trial}, Rep {rep}\")\n",
    "    for positions, timestamps, color, i in zip(\n",
    "        group.to_target_mouse_positions,\n",
    "        group.to_target_timestamps,\n",
    "        colors,\n",
    "        range(len(group)),\n",
    "    ):\n",
    "        if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "            positions = concatenate(positions, group.to_center_mouse_positions.iloc[i])\n",
    "            timestamps = concatenate(timestamps, group.to_center_timestamps.iloc[i])\n",
    "        ax.plot(\n",
    "            timestamps[:-1] - timestamps[0],\n",
    "            get_velocity(timestamps, positions),\n",
    "            color=color,\n",
    "        )\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### apply a Savitzky-Golay filter example\n",
    "Here is an example for illustrating how to apply a [Savitzky-Golay filter](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.savgol_filter.html) to the velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_filter(window_length, polyorder):\n",
    "    \"\"\"\n",
    "    plot the original function and filtered function\n",
    "    :param window_length: The length of the filter window (i.e., the number of coefficients). If mode is ‘interp’, window_length must be less than or equal to the size of x.\n",
    "    :param polyorder: The order of the polynomial used to fit the samples. polyorder must be less than window_length.\n",
    "    \"\"\"\n",
    "    for _, group in stats.groupby([\"i_trial\", \"i_rep\", \"condition_index\"]):\n",
    "        for positions, timestamps in zip(\n",
    "            group.to_target_mouse_positions,\n",
    "            group.to_target_timestamps,\n",
    "        ):\n",
    "            velocity = get_velocity(timestamps, positions)\n",
    "            plt.plot(timestamps[:-1], velocity, linestyle=\"dashed\")\n",
    "            filtered_velocity = savgol_filter(velocity, window_length, polyorder)\n",
    "            plt.plot(timestamps[:-1], filtered_velocity)\n",
    "            break\n",
    "        break\n",
    "    plt.legend([\"original velocity\", \"filtered velocity\"], loc=\"upper left\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "interact(\n",
    "    plot_filter,\n",
    "    window_length=widgets.IntSlider(min=1, max=40, step=1, value=40),\n",
    "    polyorder=widgets.IntSlider(min=1, max=40, step=1, value=8),\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### apply filter to the mouse positions\n",
    "For example, a scatter plot of the movement for each target in separate plots, the filtered movement is displayed in black dashed line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nTrials, nReps * 8, figsize=(6 * 8, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps * 8)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    for positions, color, i in zip(\n",
    "        group.to_target_mouse_positions, colors, range(len(group))\n",
    "    ):\n",
    "        ax = axs[(trial, rep + i)]\n",
    "        ax.set_title(\n",
    "            f\"[Condition {condition_index}] Trial {trial}, Rep {rep}, Target {i}\"\n",
    "        )\n",
    "        ax.set_xlim(-0.5, 0.5)\n",
    "        ax.set_ylim(-0.5, 0.5)\n",
    "        if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "            positions = concatenate(positions, group.to_center_mouse_positions.iloc[i])\n",
    "        ax.plot(positions[:, 0], positions[:, 1], color=color, linestyle=\"dashed\")\n",
    "        window_length = len(positions[:, 0])\n",
    "        polyorder = 8\n",
    "        filtered_y = savgol_filter(positions[:, 1], window_length, polyorder)\n",
    "        ax.plot(positions[:, 0], filtered_y, color=\"black\")\n",
    "\n",
    "\n",
    "fig.delaxes(axs[2][6])\n",
    "fig.delaxes(axs[2][7])\n",
    "fig.delaxes(axs[3][6])\n",
    "fig.delaxes(axs[3][7])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### apply filter to the mouse position first, then plot the velocity\n",
    "For example, a scatter plot of the velocity for each target in separate plots, the filtered velocity is displayed in black dashed line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nTrials, nReps * 8, figsize=(6 * 8, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps * 8)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    for positions, timestamps, color, i in zip(\n",
    "        group.to_target_mouse_positions,\n",
    "        group.to_target_timestamps,\n",
    "        colors,\n",
    "        range(len(group)),\n",
    "    ):\n",
    "        ax = axs[(trial, rep + i)]\n",
    "        ax.set_title(\n",
    "            f\"[Condition {condition_index}] Trial {trial}, Rep {rep}, Target {i}\"\n",
    "        )\n",
    "        if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "            positions = concatenate(positions, group.to_center_mouse_positions.iloc[i])\n",
    "            timestamps = concatenate(timestamps, group.to_center_timestamps.iloc[i])\n",
    "        velocity = get_velocity(timestamps, positions)\n",
    "        ax.plot(timestamps[:-1], velocity, color=color, linestyle=\"dashed\")\n",
    "        filtered_positions = positions.copy()\n",
    "        window_length = len(filtered_positions[:, 1])\n",
    "        polyorder = 8\n",
    "        filtered_positions[:, 1] = savgol_filter(\n",
    "            filtered_positions[:, 1], window_length, polyorder\n",
    "        )\n",
    "        velocity_with_filtered_positions = get_velocity(timestamps, filtered_positions)\n",
    "        ax.plot(\n",
    "            timestamps[:-1],\n",
    "            velocity_with_filtered_positions,\n",
    "            color=\"black\",\n",
    "        )\n",
    "\n",
    "fig.delaxes(axs[2][6])\n",
    "fig.delaxes(axs[2][7])\n",
    "fig.delaxes(axs[3][6])\n",
    "fig.delaxes(axs[3][7])\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### apply filter to the velocity plot to make it smoother\n",
    "For example, a scatter plot of the velocity for each target in separate plots, the filtered velocity is displayed in black dashed line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nTrials, nReps * 8, figsize=(6 * 8, 6 * nTrials * nReps))\n",
    "axs = np.reshape(\n",
    "    axs, (nTrials, nReps * 8)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for (trial, rep, condition_index), group in stats.groupby(\n",
    "    [\"i_trial\", \"i_rep\", \"condition_index\"]\n",
    "):\n",
    "    for positions, timestamps, color, i in zip(\n",
    "        group.to_target_mouse_positions,\n",
    "        group.to_target_timestamps,\n",
    "        colors,\n",
    "        range(len(group)),\n",
    "    ):\n",
    "        ax = axs[(trial, rep + i)]\n",
    "        ax.set_title(\n",
    "            f\"[Condition {condition_index}] Trial {trial}, Rep {rep}, Target {i}\"\n",
    "        )\n",
    "        if not experiment.trial_list[condition_index][\"automove_cursor_to_center\"]:\n",
    "            positions = concatenate(positions, group.to_center_mouse_positions.iloc[i])\n",
    "            timestamps = concatenate(timestamps, group.to_center_timestamps.iloc[i])\n",
    "        velocity = get_velocity(timestamps, positions)\n",
    "        ax.plot(timestamps[:-1], velocity, color=color, linestyle=\"dashed\")\n",
    "        window_length = len(velocity)\n",
    "        polyorder = len(velocity) - 1\n",
    "        filtered_velocity = savgol_filter(velocity, window_length, polyorder)\n",
    "        ax.plot(timestamps[:-1], filtered_velocity, color=\"black\")\n",
    "\n",
    "\n",
    "fig.delaxes(axs[2][6])\n",
    "fig.delaxes(axs[2][7])\n",
    "fig.delaxes(axs[3][6])\n",
    "fig.delaxes(axs[3][7])\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}