firmware_emulator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code for the coursework for the Virtual and Augmented Reality submodule in the Contemporary Computer Science IV module\n",
    "Libraries used: python (3.7.2), pandas, matplotlib, numpy\n",
    "\n",
    "\n",
    "### To run this code, please follow these instructions:\n",
    "The code works best when executed in order (top to bottom), one cell at a time. The matplotlib outputs are \n",
    "given in pop-out windows, but the 3D (animated) plots seem to replace one-another and the previous outputs\n",
    "when they are produced.\n",
    "\n",
    "Thus, to see all output correctly, the cells for problems 5 and 6 should be run one at a time. The previous cells can be run together (in order) in order to save a little time. After these have been run, the output-producing cells for problems 5 and 6 can be run and re-run as desired"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import figure\n",
    "from matplotlib import animation\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "# Create aliases for the headings of the IMU data for easier processing\n",
    "time   = 'time'\n",
    "gyro_X = ' gyroscope.X'\n",
    "gyro_Y = ' gyroscope.Y'\n",
    "gyro_Z = ' gyroscope.Z'\n",
    "acc_X  = ' accelerometer.X'\n",
    "acc_Y  = ' accelerometer.Y'\n",
    "acc_Z  = ' accelerometer.Z'\n",
    "mag_X  = ' magnetometer.X'\n",
    "mag_Y  = ' magnetometer.Y'\n",
    "mag_Z  = ' magnetometer.Z '\n",
    "\n",
    "# NOTE:\n",
    "# X is 'depth' (forward/backward) (i.e. the axis of roll)\n",
    "# Y is 'across' (left/right) (i.e. the axis of pitch)\n",
    "# Z is 'up' (up/down) (i.e. the axis of yaw)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 1\n",
    "\n",
    "# Read the IMU data from the .csv file\n",
    "IMU_Data = pd.read_csv(\"IMUData.csv\")\n",
    "# Create raw acc reading values to be used for Problem 6\n",
    "IMU_Data['raw_accelerometer.X'] = IMU_Data[acc_X]\n",
    "IMU_Data['raw_accelerometer.Y'] = IMU_Data[acc_Y]\n",
    "IMU_Data['raw_accelerometer.Z'] = IMU_Data[acc_Z]\n",
    "\n",
    "# Convert rotational rate from deg/s to rad/s\n",
    "IMU_Data[gyro_X] = IMU_Data[gyro_X].apply(np.radians)\n",
    "IMU_Data[gyro_Y] = IMU_Data[gyro_Y].apply(np.radians)\n",
    "IMU_Data[gyro_Z] = IMU_Data[gyro_Z].apply(np.radians)\n",
    "\n",
    "# Normalize the accelerometer and magnetometer values\n",
    "def normalize(x, y, z):\n",
    "    \"\"\"Normalizes the magnitude of the given [X, Y, Z] vector\n",
    "    \n",
    "    \\in x: the x-axis value\n",
    "    \\in y: the y-axis value\n",
    "    \\in z: the z-axis value\n",
    "    \\out x_norm: the normalised x-axis value\n",
    "    \\out y_norm: the normalised y-axis value\n",
    "    \\out z_norm: the normalised z-axis value\n",
    "    \"\"\"\n",
    "    magnitude = np.sqrt(x**2 + y**2 + z**2)\n",
    "    \n",
    "    # Only normalize if magnitude > 0\n",
    "    if magnitude:\n",
    "        return x/magnitude, y/magnitude, z/magnitude\n",
    "    \n",
    "    # If magnitude == 0, return the vector as-is\n",
    "    return x, y, z\n",
    "\n",
    "for index, row in IMU_Data.iterrows():\n",
    "    row[acc_X], row[acc_Y], row[acc_Z] = normalize(row[acc_X], row[acc_Y], row[acc_Z])\n",
    "    row[mag_X], row[mag_Y], row[mag_Z] = normalize(row[mag_X], row[mag_Y], row[mag_Z])\n",
    "\n",
    "# Methods for converting Euler angles (radians) <---> quaternions\n",
    "\n",
    "def euler2quaternion(eulers):\n",
    "    \"\"\"Converts the given euler angles into a quaternion format\n",
    "    \n",
    "    \\in eulers: the input euler angles, of format (x, y, z), equivalently (roll, pitch, yaw)\n",
    "    \\out quat: the equivalent quaternion, of format (w, x, y, z)\n",
    "    \"\"\"\n",
    "    # extract the euler angles from eulers\n",
    "    roll, pitch, yaw = eulers\n",
    "    \n",
    "    # Compute the sinusoidal values needed\n",
    "    cos_y = np.cos(yaw/2)\n",
    "    sin_y = np.sin(yaw/2)\n",
    "    cos_p = np.cos(pitch/2)\n",
    "    sin_p = np.sin(pitch/2)\n",
    "    cos_r = np.cos(roll/2)\n",
    "    sin_r = np.sin(roll/2)\n",
    "    \n",
    "    # Compute the quaternion values\n",
    "    quat_w = cos_y * cos_p * cos_r + sin_y * sin_p * sin_r\n",
    "    quat_x = cos_y * cos_p * sin_r - sin_y * sin_p * cos_r\n",
    "    quat_y = sin_y * cos_p * sin_r + cos_y * sin_p * cos_r\n",
    "    quat_z = sin_y * cos_p * cos_r - cos_y * sin_p * sin_r\n",
    "    \n",
    "    return (quat_w, quat_x, quat_y, quat_z)\n",
    "\n",
    "def quaternion2euler(quat):\n",
    "    \"\"\"Converts the given quaternion into euler angle format\n",
    "    \n",
    "    \\in quat: the input quaternion, of format (w, x, y, z)\n",
    "    \\out euler: the equivalent euler angles, of format (x, y, z), equivalently (roll, pitch, yaw)\n",
    "    \"\"\"\n",
    "    # extract the quaternion values from quat\n",
    "    w, x, y, z = quat\n",
    "    \n",
    "    # Compute the roll\n",
    "    roll  = np.arctan2(2 * (w*x + y*z), 1 - 2 * (x**2 + y**2))\n",
    "    \n",
    "    # Compute the pitch\n",
    "    asin_val = 2 * (w*y - z*x)\n",
    "    # Check if this value is out of the range of arcsin (-1, 1)\n",
    "    if abs(asin_val) >= 1:\n",
    "        # Use pi/2 as the value as the value is out of the range of arcsin\n",
    "        sign = 1 if asin_val >= 0 else -1\n",
    "        pitch = np.pi/2 * sign\n",
    "    else:\n",
    "        pitch = np.arcsin(asin_val)\n",
    "    \n",
    "    # Compute the yaw\n",
    "    yaw = np.arctan2(2 * (w*z + x*y), 1 - 2 * (y**2 + z**2))\n",
    "    \n",
    "    return (roll, pitch, yaw)\n",
    "    \n",
    "\n",
    "# Method for converting quaternion to its conjugate\n",
    "def quaternion_conjugate(quat):\n",
    "    \"\"\"Convert the given quaternion to its conjugate\n",
    "    \n",
    "    \\in quat: the input quaternion, of format (w, x, y, z)\n",
    "    \\out conj_quat: the conjugate of quat, of format (w, x, y, z)\n",
    "    \"\"\"\n",
    "    return (quat[0], -quat[1], -quat[2], -quat[3])\n",
    "\n",
    "# Method for calculating quaternion product of two quaternions\n",
    "def quaternion_product(quat1, quat2):\n",
    "    \"\"\"Calculate the product of the two quaternions\n",
    "    \n",
    "    \\in quat1: the first quaternion, of format (w, x, y, z)\n",
    "    \\in quat2: the second quaternion, of format (w, x, y, z)\n",
    "    \\out quat: the product of quat1 and quat2, of format (w, x, y, z)\n",
    "    \"\"\"\n",
    "    # Extract the quaternion values\n",
    "    w1, x1, y1, z1 = quat1\n",
    "    w2, x2, y2, z2 = quat2\n",
    "    \n",
    "    # Construct the quaternion product\n",
    "    w = w1*w2 - x1*x2 - y1*y2 - z1*z2\n",
    "    x = w1*x2 + x1*w2 + y1*z2 - z1*y2\n",
    "    y = w1*y2 - x1*z2 + y1*w2 + z1*x2\n",
    "    z = w1*z2 + x1*y2 - y1*x2 + z1*w2\n",
    "    \n",
    "    return (w, x, y, z)\n",
    "\n",
    "# General (non-required) utility functions\n",
    "def quaternion_rotation(q, p):\n",
    "    \"\"\"Rotates the point p by the quaternion q\n",
    "    \n",
    "    \\in q: the quaternion of the rotation\n",
    "    \\in p: the point to be rotated by q\n",
    "    \\out result: the point produced by the rotation\"\"\"\n",
    "    if len(p) == 3:\n",
    "        p = (0, p[0], p[1], p[2])\n",
    "    result = quaternion_product(p, quaternion_conjugate(q))\n",
    "    result = quaternion_product(q, result)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Problem 2\n",
    "\n",
    "# Create the DataFrame used to store the head orientations\n",
    "head_orientation = pd.DataFrame(columns=['time', 'quaternion', 'W', 'X', 'Y', 'Z']) # List of quaternions corresponding to orientation at each time stamp\n",
    "head_orientation['time'] = IMU_Data['time'][:]\n",
    "\n",
    "def simple_dead_reckoning(prev_ho_row, current_IMU_Data_row):\n",
    "    \"\"\"Performs simple dead reckoning to track the current head orientation by integrating\n",
    "    the gyroscope readings.\n",
    "    \n",
    "    \\in prev_ho_row: the previous head_orientation row\n",
    "    \\in current_IMU_Data_row: the row of data in the IMU_Data frame corresponding to the current time step\n",
    "    \\out current_orientation: the estimated current head orientation\n",
    "    \"\"\"\n",
    "    # Extract the data we need from the current_IMU_Data_row (row in IMU_Data frame that corresponds to the current time step)\n",
    "    current_time = current_IMU_Data_row[time]\n",
    "    gyroX        = current_IMU_Data_row[gyro_X]\n",
    "    gyroY        = current_IMU_Data_row[gyro_Y]\n",
    "    gyroZ        = current_IMU_Data_row[gyro_Z]\n",
    "    \n",
    "    # Extract the data we need from the prev_ho_row (previous row in the head_orientation_tilt_correction frame)\n",
    "    last_orientation = prev_ho_row['quaternion']\n",
    "    last_time = prev_ho_row[time]\n",
    "    \n",
    "    # Calculate the magnitude of the gyro readings\n",
    "    l = np.sqrt(gyroX ** 2 + gyroY ** 2 + gyroZ ** 2)\n",
    "    \n",
    "    # Normalize the gyro readings (so they can be used as axes in axis-angle representation)\n",
    "    gyroX /= l\n",
    "    gyroY /= l\n",
    "    gyroZ /= l\n",
    "    \n",
    "    # Calculate the angle of the rotation\n",
    "    theta = l*(current_time - last_time)\n",
    "    \n",
    "    # Create the quaternion from this axis-angle representation\n",
    "    change_quat = (np.cos(theta/2), gyroX*np.sin(theta/2), gyroY*np.sin(theta/2), gyroZ*np.sin(theta/2))\n",
    "\n",
    "    # Produce the current orientation by combining (multiplying) the last orientation and the change quaternion\n",
    "    current_orientation = quaternion_product(last_orientation, change_quat)\n",
    "    \n",
    "    return current_orientation\n",
    "\n",
    "# Initial orientation\n",
    "head_orientation.at[0, 'quaternion'] = (1, 0, 0, 0)\n",
    "for index in range(len(head_orientation)):\n",
    "    # Skip the first one as this has the value of the identity quaternion\n",
    "    if index != 0:\n",
    "        prev_row = head_orientation.iloc[index - 1]\n",
    "        imu_row = IMU_Data.iloc[index]\n",
    "        # Estimate the current head orientation\n",
    "        head_orientation.at[index, 'quaternion'] = simple_dead_reckoning(prev_row, imu_row)\n",
    "    # Set the separate W, X, Y, Z values of this row\n",
    "    quat = head_orientation.iloc[index]['quaternion']\n",
    "    head_orientation.at[index, 'W'] = quat[0]\n",
    "    head_orientation.at[index, 'X'] = quat[1]\n",
    "    head_orientation.at[index, 'Y'] = quat[2]\n",
    "    head_orientation.at[index, 'Z'] = quat[3]\n",
    "\n",
    "# Convert the quaternion orientations to Euler angles for 2D plotting in Problem 5\n",
    "head_orientation_euler = pd.DataFrame(columns=['time', 'X', 'Y', 'Z'])\n",
    "head_orientation_euler['time'] = IMU_Data['time'][:]\n",
    "for index in range(len(head_orientation)):\n",
    "    quat = head_orientation.iloc[index]['quaternion']\n",
    "    x, y, z = quaternion2euler(quat)\n",
    "    head_orientation_euler.at[index, 'X'] = x\n",
    "    head_orientation_euler.at[index, 'Y'] = y\n",
    "    head_orientation_euler.at[index, 'Z'] = z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 3\n",
    "\n",
    "# Create the DataFrame used to store the head orientations\n",
    "head_orientation_tilt_correction = pd.DataFrame(columns=['time', 'quaternion', 'W', 'X', 'Y', 'Z']) # List of quaternions corresponding to orientation at each time stamp\n",
    "head_orientation_tilt_correction['time'] = IMU_Data['time'][:]\n",
    "\n",
    "def tilt_correction(prev_ho_row, current_IMU_Data_row):\n",
    "    \"\"\"Estimates the head orientation with additionally applying gravity-based tilt correction\n",
    "    \n",
    "    \\in prev_ho_row: the row of data (time and quaternion) corresponding to the last time step\n",
    "    \\in current_IMU_Data_row: the row of data in the IMU_Data frame that corresponds to the current time step\n",
    "    \\out tilt_corrected_orientation: the estimated head orientation, with gravity-based tilt correction\n",
    "    \"\"\"\n",
    "    # Extract the data we need from current_IMU_Data_row\n",
    "    accX = current_IMU_Data_row[acc_X]\n",
    "    accY = current_IMU_Data_row[acc_Y]\n",
    "    accZ = current_IMU_Data_row[acc_Z]\n",
    "        \n",
    "    # Get the result from simple dead reckoning\n",
    "    simple_orientation = simple_dead_reckoning(prev_ho_row, current_IMU_Data_row)\n",
    "    \n",
    "    # Perform tilt correction\n",
    "    \n",
    "    # Transform acceleration measurements into the global frame\n",
    "    acc_quat = (0, accX, accY, accZ) # convert the acc vector into a quaternion as per LaValle paper\n",
    "    a_global = quaternion_rotation(simple_orientation, acc_quat)\n",
    "    \n",
    "    # Calculate the tilt axis - this is orthogonal to the XY plane (as Z is 'up'), so we take it as (-a_y, a_x, 0)\n",
    "    tilt_axis = (-a_global[2], a_global[1], 0)\n",
    "    \n",
    "    # Find the angle phi between the up vector (0, 0, 1) and the accelerometer vector (a_x, a_y, a_z)\n",
    "    # phi = arccos(dot_product(a, up)/(magnitude(a)*magnitude(up))). However, magnitude of both a and up are 1, as a is normalized\n",
    "    # and up is (0, 0, 1). Also, since up is (0, 0, 1) the dot product of a and up is simply a_z. Thus:\n",
    "    phi = np.arccos(a_global[3])\n",
    "    \n",
    "    # Use the complementary filter to fuse the two signals\n",
    "    \n",
    "    # Construct the tilt quaternion\n",
    "    alpha = 0.03\n",
    "    theta = -alpha * phi\n",
    "\n",
    "    tilt_quat = (np.cos(theta/2), tilt_axis[0]*np.sin(theta/2), tilt_axis[1]*np.sin(theta/2), tilt_axis[2]*np.sin(theta/2))\n",
    "    \n",
    "    tilt_corrected_orientation = quaternion_product(tilt_quat, simple_orientation)\n",
    "    \n",
    "    return tilt_corrected_orientation\n",
    "\n",
    "# Initial orientation\n",
    "head_orientation_tilt_correction.at[0, 'quaternion'] = (1, 0, 0, 0)\n",
    "for index in range(len(head_orientation_tilt_correction)):\n",
    "    # Skip the first one as this has the value of the identity quaternion\n",
    "    if index != 0:\n",
    "        prev_row = head_orientation_tilt_correction.iloc[index - 1]\n",
    "        imu_row = IMU_Data.iloc[index]\n",
    "        # Estimate the current head orientation, with tilt correction\n",
    "        head_orientation_tilt_correction.at[index, 'quaternion'] = tilt_correction(prev_row, imu_row)\n",
    "    # Set the separate W, X, Y, Z values of this row\n",
    "    quat = head_orientation_tilt_correction.iloc[index]['quaternion']\n",
    "    head_orientation_tilt_correction.at[index, 'W'] = quat[0]\n",
    "    head_orientation_tilt_correction.at[index, 'X'] = quat[1]\n",
    "    head_orientation_tilt_correction.at[index, 'Y'] = quat[2]\n",
    "    head_orientation_tilt_correction.at[index, 'Z'] = quat[3]\n",
    "\n",
    "# Convert the quaternion orientations to Euler angles for 2D plotting in Problem 5\n",
    "head_orientation_tilt_correction_euler = pd.DataFrame(columns=['time', 'X', 'Y', 'Z'])\n",
    "head_orientation_tilt_correction_euler['time'] = IMU_Data['time'][:]\n",
    "for index in range(len(head_orientation_tilt_correction)):\n",
    "    quat = head_orientation_tilt_correction.iloc[index]['quaternion']\n",
    "    x, y, z = quaternion2euler(quat)\n",
    "    head_orientation_tilt_correction_euler.at[index, 'X'] = x\n",
    "    head_orientation_tilt_correction_euler.at[index, 'Y'] = y\n",
    "    head_orientation_tilt_correction_euler.at[index, 'Z'] = z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 4\n",
    "\n",
    "# Create the DataFrame used to store the head orientations\n",
    "head_orientation_yaw_correction = pd.DataFrame(columns=['time', 'quaternion', 'W', 'X', 'Y', 'Z']) # List of quaternions corresponding to orientation at each time stamp\n",
    "head_orientation_yaw_correction['time'] = IMU_Data['time'][:]\n",
    "\n",
    "def yaw_correction(prev_ho_row, current_IMU_Data_row):\n",
    "    \"\"\"Estimates the head orientation with additionally applying gravity-based tilt correction\n",
    "    \n",
    "    \\in prev_ho_row: the row of data (time and quaternion) corresponding to the last time step\n",
    "    \\in current_IMU_Data_row: the row of data in the IMU_Data frame that corresponds to the current time step\n",
    "    \\out tilt_corrected_orientation: the estimated head orientation, with gravity-based tilt correction\n",
    "    \"\"\"\n",
    "    # Extract the data we need from current_IMU_Data_row\n",
    "    magX = current_IMU_Data_row[mag_X]\n",
    "    magY = current_IMU_Data_row[mag_Y]\n",
    "    magZ = current_IMU_Data_row[mag_Z]\n",
    "    \n",
    "    # Get the result from tilt correction\n",
    "    tilt_corrected_orientation = tilt_correction(prev_ho_row, current_IMU_Data_row)\n",
    "    \n",
    "    # Perform yaw correction\n",
    "    \n",
    "    # Get reference readings\n",
    "    mag_ref = (0, IMU_Data.iloc[0][mag_X], IMU_Data.iloc[0][mag_Y], IMU_Data.iloc[0][mag_Z])\n",
    "    quat_ref = head_orientation_yaw_correction.iloc[0]['quaternion']\n",
    "    \n",
    "    mag_current = (0, magX, magY, magZ)\n",
    "    \n",
    "    # Bring the current readings into the estimated global frame\n",
    "    q = tilt_corrected_orientation\n",
    "    q_inv = quaternion_conjugate(tilt_corrected_orientation)\n",
    "    mag_current = quaternion_product(q, quaternion_product(mag_current, q_inv))\n",
    "    \n",
    "    # Find the angular difference between mag_current and mag_ref, when projected into the XY plane\n",
    "    theta_current = np.arctan2(mag_current[1], mag_current[2])\n",
    "    theta_ref = np.arctan2(mag_ref[1], mag_ref[2])\n",
    "    \n",
    "    # Apply the complementary filter to fix the yaw drift\n",
    "    alpha = 0.0002\n",
    "    # Construct the correction quaternion\n",
    "    theta = -alpha * (theta_current - theta_ref)\n",
    "    correction_quat = (np.cos(theta/2), 0, 0, np.sin(theta/2))\n",
    "    \n",
    "    yaw_corrected_orientation = quaternion_product(correction_quat, tilt_corrected_orientation) #if (theta_current - theta_ref) > \n",
    "    \n",
    "    return yaw_corrected_orientation\n",
    "\n",
    "# Initial orientation\n",
    "head_orientation_yaw_correction.at[0, 'quaternion'] = (1, 0, 0, 0)\n",
    "for index in range(len(head_orientation_yaw_correction)):\n",
    "    # Skip the first one as this has the value of the identity quaternion\n",
    "    if index != 0:\n",
    "        prev_row = head_orientation_yaw_correction.iloc[index - 1]\n",
    "        imu_row = IMU_Data.iloc[index]\n",
    "        # Estimate the current head orientation, with tilt correction\n",
    "        head_orientation_yaw_correction.at[index, 'quaternion'] = yaw_correction(prev_row, imu_row)\n",
    "    # Set the separate W, X, Y, Z values of this row\n",
    "    quat = head_orientation_yaw_correction.iloc[index]['quaternion']\n",
    "    head_orientation_yaw_correction.at[index, 'W'] = quat[0]\n",
    "    head_orientation_yaw_correction.at[index, 'X'] = quat[1]\n",
    "    head_orientation_yaw_correction.at[index, 'Y'] = quat[2]\n",
    "    head_orientation_yaw_correction.at[index, 'Z'] = quat[3]\n",
    "\n",
    "# Convert the quaternion orientations to Euler angles for 2D plotting in Problem 5\n",
    "head_orientation_yaw_correction_euler = pd.DataFrame(columns=['time', 'X', 'Y', 'Z'])\n",
    "head_orientation_yaw_correction_euler['time'] = IMU_Data['time'][:]\n",
    "\n",
    "for index in range(len(head_orientation_yaw_correction)):\n",
    "    quat = head_orientation_yaw_correction.iloc[index]['quaternion']\n",
    "    x, y, z = quaternion2euler(quat)\n",
    "    head_orientation_yaw_correction_euler.at[index, 'X'] = x\n",
    "    head_orientation_yaw_correction_euler.at[index, 'Y'] = y\n",
    "    head_orientation_yaw_correction_euler.at[index, 'Z'] = z\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 5 - Gyroscope Readings\n",
    "\n",
    "# Make plots be produced in pop-out windows and be interactive\n",
    "%matplotlib qt\n",
    "# Tri-axial angular rate in deg/s\n",
    "fig = figure(figsize=(24, 8))\n",
    "fig.canvas.set_window_title('Gyroscope Readings')\n",
    "\n",
    "plt.plot(IMU_Data[time], np.degrees(IMU_Data[gyro_X]), label='X')\n",
    "plt.plot(IMU_Data[time], np.degrees(IMU_Data[gyro_Y]), label='Y')\n",
    "plt.plot(IMU_Data[time], np.degrees(IMU_Data[gyro_Z]), label='Z')\n",
    "\n",
    "plt.title(\"Tri-axial angular rate (Gyroscope)\")\n",
    "plt.xlabel(\"Time (s)\")\n",
    "plt.ylabel(\"Angular rate ($deg/s$)\")\n",
    "plt.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 5 - Accelerometer Readings\n",
    "\n",
    "# Make plots be produced in pop-out windows and be interactive\n",
    "%matplotlib qt\n",
    "# Normalized tri-axial acceleration in g (m/s^2)\n",
    "fig = figure(figsize=(24, 8))\n",
    "fig.canvas.set_window_title('Normalized Accelerometer Readings')\n",
    "plt.plot(IMU_Data[time], IMU_Data[acc_X], label='X')\n",
    "plt.plot(IMU_Data[time], IMU_Data[acc_Y], label='Y')\n",
    "plt.plot(IMU_Data[time], IMU_Data[acc_Z], label='Z')\n",
    "\n",
    "plt.title(\"Normalized tri-axial acceleration (Accelerometer)\")\n",
    "plt.xlabel(\"Time (s)\")\n",
    "plt.ylabel(\"Acceleration in $g$ ($m/s^2$)\")\n",
    "plt.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 5 - Magnetometer Readings\n",
    "\n",
    "# Make plots be produced in pop-out windows and be interactive\n",
    "%matplotlib qt\n",
    "# Normalized tri-axial magnetic flux in Gauss (G)\n",
    "fig = figure(figsize=(24, 8))\n",
    "fig.canvas.set_window_title('Normalized Magnetometer Readings')\n",
    "plt.plot(IMU_Data[time], IMU_Data[mag_X], label='X')\n",
    "plt.plot(IMU_Data[time], IMU_Data[mag_Y], label='Y')\n",
    "plt.plot(IMU_Data[time], IMU_Data[mag_Z], label='Z')\n",
    "\n",
    "plt.title(\"Normalized tri-axial magnetic flux (Magnetometer)\")\n",
    "plt.xlabel(\"Time (s)\")\n",
    "plt.ylabel(\"Magnetic flux in Gauss ($G$)\")\n",
    "plt.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 5 - Euler Angles of Head Orientation\n",
    "\n",
    "# Tri-axial Euler angles in degrees for gyroscope integration, gyroscope + accelerometer, gyroscope + accelerometer + magnetometer\n",
    "\n",
    "fig = figure(figsize=(24, 12))\n",
    "fig.canvas.set_window_title('Tri-axial Euler angles (deg) for 3 types of orientation tracking')\n",
    "plt.tight_layout()\n",
    "plt.subplot(3, 1, 1)\n",
    "# Gyroscope\n",
    "plt.plot(head_orientation_euler[time], np.degrees(np.float32(head_orientation_euler['X'])), label='X')\n",
    "plt.plot(head_orientation_euler[time], np.degrees(np.float32(head_orientation_euler['Y'])), label='Y')\n",
    "plt.plot(head_orientation_euler[time], np.degrees(np.float32(head_orientation_euler['Z'])), label='Z')\n",
    "plt.axhline(0, color='black', alpha=0.3)\n",
    "plt.title(\"Gyroscope Integration\")\n",
    "plt.xlabel(\"Time (s)\")\n",
    "plt.ylabel(\"Euler angle (deg)\")\n",
    "plt.legend(loc='upper right')\n",
    "\n",
    "plt.subplot(3, 1, 2)\n",
    "# Gyroscope + Accelerometer\n",
    "plt.plot(head_orientation_tilt_correction_euler[time], np.degrees(np.float32(head_orientation_tilt_correction_euler['X'])), label='X')\n",
    "plt.plot(head_orientation_tilt_correction_euler[time], np.degrees(np.float32(head_orientation_tilt_correction_euler['Y'])), label='Y')\n",
    "plt.plot(head_orientation_tilt_correction_euler[time], np.degrees(np.float32(head_orientation_tilt_correction_euler['Z'])), label='Z')\n",
    "plt.axhline(0, color='black', alpha=0.3)\n",
    "plt.title(\"Gyroscope + Accelerometer\")\n",
    "plt.xlabel(\"Time (s)\")\n",
    "plt.ylabel(\"Euler angle (deg)\")\n",
    "plt.legend(loc='upper right')\n",
    "\n",
    "plt.subplot(3, 1, 3)\n",
    "# Gyroscope + Accelerometer + Magnetometer\n",
    "plt.plot(head_orientation_yaw_correction_euler[time], np.degrees(np.float32(head_orientation_yaw_correction_euler['X'])), label='X')\n",
    "plt.plot(head_orientation_yaw_correction_euler[time], np.degrees(np.float32(head_orientation_yaw_correction_euler['Y'])), label='Y')\n",
    "plt.plot(head_orientation_yaw_correction_euler[time], np.degrees(np.float32(head_orientation_yaw_correction_euler['Z'])), label='Z')\n",
    "plt.axhline(0, color='black', alpha=0.3)\n",
    "plt.title(\"Gyroscope + Accelerometer + Magnetometer\")\n",
    "plt.xlabel(\"Time (s)\")\n",
    "plt.ylabel(\"Euler angle (deg)\")\n",
    "plt.legend(loc='upper right')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 5 - 3D plots full speed\n",
    "# 3D plots of the 3 main axis of the IMU for gyroscope integration, gyroscope + acceleration, gyroscope + acceleration + magnetometer\n",
    "\n",
    "# Make plots pop-out and be interactive\n",
    "%matplotlib qt\n",
    "\n",
    "N = 6959 # Number of frames\n",
    "interval = 4 # Time between frames in ms - 1000/256 as time steps are at 256Hz\n",
    "resolution = 25 # How often frames are generated (i.e. resolution=4 means every 4th data point is plotted)\n",
    "# The resolution variable is used to ensure the animation is full speed/half speed as required. With the given values it runs at the correct speeds on my computer, but if yours is faster/slower the actual speeds will likely differ\n",
    "\n",
    "N //= resolution # Set the total number of frames as adjusted by resolution\n",
    "interval *= resolution # Set the time between frames as adjusted by resolution\n",
    "scale = 2 # Scale of the axis limits\n",
    "\n",
    "axisX, axisY, axisZ = [0, 0, 0], [0, 0, 0], [0, 0, 0]\n",
    "\n",
    "def animate_axes(num, ax1, ax2, ax3):\n",
    "    \"\"\"Generate the new quivers\"\"\"\n",
    "    global axisX, axisY, axisZ\n",
    "    # Check that this isn't the first frame\n",
    "    if axisX[0] != 0:\n",
    "        for i in axisX, axisY, axisZ:\n",
    "            for j in i:\n",
    "                j.remove()\n",
    "    \n",
    "    num *= resolution\n",
    "    \n",
    "    x_proj_vec = (0, 1, 0, 0)\n",
    "    y_proj_vec = (0, 0, 1, 0)\n",
    "    z_proj_vec = (0, 0, 0, 1)\n",
    "    # Apply the current quaternion (indexed by num) to these vectors\n",
    "    quaternion = head_orientation.iloc[num]['quaternion']\n",
    "    tilt_corrected_quaternion = head_orientation_tilt_correction.iloc[num]['quaternion']\n",
    "    yaw_corrected_quaternion = head_orientation_yaw_correction.iloc[num]['quaternion']\n",
    "    \n",
    "    axes = [ax1, ax2, ax3]\n",
    "    quats = [quaternion, tilt_corrected_quaternion, yaw_corrected_quaternion]\n",
    "    \n",
    "    for i in range(3):\n",
    "        ax = axes[i]\n",
    "        quat = quats[i]\n",
    "\n",
    "        x_vec = quaternion_rotation(quat, x_proj_vec)\n",
    "        y_vec = quaternion_rotation(quat, y_proj_vec)\n",
    "        z_vec = quaternion_rotation(quat, z_proj_vec)\n",
    "\n",
    "        axisX[i] = ax.quiver(0, 0, 0, *x_vec[1:], color='Blue', label='X')\n",
    "        axisY[i] = ax.quiver(0, 0, 0, *y_vec[1:], color='Orange', label='Y')\n",
    "        axisZ[i] = ax.quiver(0, 0, 0, *z_vec[1:], color='Green', label='Z')\n",
    "        \n",
    "        ax.legend(loc='center left')\n",
    "\n",
    "fig = plt.figure(figsize=(24, 8))\n",
    "fig.canvas.set_window_title('Orientation Tracking - Full speed')\n",
    "# Gyroscope integration\n",
    "ax1 = fig.add_subplot(131, projection='3d')\n",
    "ax1.set_title(\"Gyroscope integration\")\n",
    "ax1.set_xlim3d([-1.0*scale, 1.0*scale]); ax1.set_xlabel('X')\n",
    "ax1.set_ylim3d([-1.0*scale, 1.0*scale]); ax1.set_ylabel('Y')\n",
    "ax1.set_zlim3d([-1.0*scale, 1.0*scale]); ax1.set_zlabel('Z')\n",
    "\n",
    "# Gyroscope + Accelerometer\n",
    "ax2 = fig.add_subplot(132, projection='3d')\n",
    "ax2.set_title(\"Gyroscope + Accelerometer\")\n",
    "ax2.set_xlim3d([-1.0*scale, 1.0*scale]); ax2.set_xlabel('X')\n",
    "ax2.set_ylim3d([-1.0*scale, 1.0*scale]); ax2.set_ylabel('Y')\n",
    "ax2.set_zlim3d([-1.0*scale, 1.0*scale]); ax2.set_zlabel('Z')\n",
    "\n",
    "# Gyroscope + Accelerometer + Magnetometer\n",
    "ax3 = fig.add_subplot(133, projection='3d')\n",
    "ax3.set_title(\"Gyroscope + Accelerometer + Magnetometer\")\n",
    "ax3.set_xlim3d([-1.0*scale, 1.0*scale]); ax3.set_xlabel('X')\n",
    "ax3.set_ylim3d([-1.0*scale, 1.0*scale]); ax3.set_ylabel('Y')\n",
    "ax3.set_zlim3d([-1.0*scale, 1.0*scale]); ax3.set_zlabel('Z')\n",
    "\n",
    "ani = animation.FuncAnimation(fig, animate_axes, N, fargs=(ax1, ax2, ax3), interval=interval, blit=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 5 - 3D plots half speed\n",
    "# 3D plots of the 3 main axis of the IMU for gyroscope integration, gyroscope + acceleration, gyroscope + acceleration + magnetometer\n",
    "\n",
    "# Make plots pop-out and be interactive\n",
    "%matplotlib qt\n",
    "\n",
    "# Redefine all variables used above so this cell can be called several times in a row\n",
    "N = 6959 # Number of frames\n",
    "interval = 4 # Time between frames in ms - 1000/256 as time steps are at 256Hz\n",
    "resolution = 12 # How often frames are generated (i.e. resolution=4 means every 4th data point is plotted)\n",
    "# The resolution variable is used to ensure the animation is full speed/half speed as required. With the given values it runs at the correct speeds on my computer, but if yours is faster/slower the actual speeds will likely differ\n",
    "\n",
    "N //= resolution # Set the total number of frames as adjusted by resolution\n",
    "interval *= resolution # Set the time between frames as adjusted by resolution\n",
    "scale = 2 # Scale of the axis limits\n",
    "\n",
    "axisX, axisY, axisZ = [0, 0, 0], [0, 0, 0], [0, 0, 0]\n",
    "\n",
    "def animate_axes(num, ax1, ax2, ax3):\n",
    "    \"\"\"Generate the new quivers\"\"\"\n",
    "    global axisX, axisY, axisZ\n",
    "    # Check that this isn't the first frame\n",
    "    if axisX[0] != 0:\n",
    "        for i in axisX, axisY, axisZ:\n",
    "            for j in i:\n",
    "                j.remove()\n",
    "        \n",
    "    num *= resolution\n",
    "    \n",
    "    x_proj_vec = (0, 1, 0, 0)\n",
    "    y_proj_vec = (0, 0, 1, 0)\n",
    "    z_proj_vec = (0, 0, 0, 1)\n",
    "    # Apply the current quaternion (indexed by num) to these vectors\n",
    "    quaternion = head_orientation.iloc[num]['quaternion']\n",
    "    tilt_corrected_quaternion = head_orientation_tilt_correction.iloc[num]['quaternion']\n",
    "    yaw_corrected_quaternion = head_orientation_yaw_correction.iloc[num]['quaternion']\n",
    "    \n",
    "    axes = [ax1, ax2, ax3]\n",
    "    quats = [quaternion, tilt_corrected_quaternion, yaw_corrected_quaternion]\n",
    "    \n",
    "    for i in range(3):\n",
    "        ax = axes[i]\n",
    "        quat = quats[i]\n",
    "\n",
    "        x_vec = quaternion_rotation(quat, x_proj_vec)\n",
    "        y_vec = quaternion_rotation(quat, y_proj_vec)\n",
    "        z_vec = quaternion_rotation(quat, z_proj_vec)\n",
    "\n",
    "        axisX[i] = ax.quiver(0, 0, 0, *x_vec[1:], color='Blue', label='X')\n",
    "        axisY[i] = ax.quiver(0, 0, 0, *y_vec[1:], color='Orange', label='Y')\n",
    "        axisZ[i] = ax.quiver(0, 0, 0, *z_vec[1:], color='Green', label='Z')\n",
    "        \n",
    "        ax.legend(loc='center left')\n",
    "\n",
    "fig = plt.figure(figsize=(24, 8))\n",
    "fig.canvas.set_window_title('Orientation Tracking - Half speed')\n",
    "# Gyroscope integration\n",
    "ax1 = fig.add_subplot(131, projection='3d')\n",
    "ax1.set_title(\"Gyroscope integration\")\n",
    "ax1.set_xlim3d([-1.0*scale, 1.0*scale]); ax1.set_xlabel('X')\n",
    "ax1.set_ylim3d([-1.0*scale, 1.0*scale]); ax1.set_ylabel('Y')\n",
    "ax1.set_zlim3d([-1.0*scale, 1.0*scale]); ax1.set_zlabel('Z')\n",
    "\n",
    "# Gyroscope + Accelerometer\n",
    "ax2 = fig.add_subplot(132, projection='3d')\n",
    "ax2.set_title(\"Gyroscope + Accelerometer\")\n",
    "ax2.set_xlim3d([-1.0*scale, 1.0*scale]); ax2.set_xlabel('X')\n",
    "ax2.set_ylim3d([-1.0*scale, 1.0*scale]); ax2.set_ylabel('Y')\n",
    "ax2.set_zlim3d([-1.0*scale, 1.0*scale]); ax2.set_zlabel('Z')\n",
    "\n",
    "# Gyroscope + Accelerometer + Magnetometer\n",
    "ax3 = fig.add_subplot(133, projection='3d')\n",
    "ax3.set_title(\"Gyroscope + Accelerometer + Magnetometer\")\n",
    "ax3.set_xlim3d([-1.0*scale, 1.0*scale]); ax3.set_xlabel('X')\n",
    "ax3.set_ylim3d([-1.0*scale, 1.0*scale]); ax3.set_ylabel('Y')\n",
    "ax3.set_zlim3d([-1.0*scale, 1.0*scale]); ax3.set_zlabel('Z')\n",
    "\n",
    "ani = animation.FuncAnimation(fig, animate_axes, N, fargs=(ax1, ax2, ax3), interval=interval, blit=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 6\n",
    "\n",
    "# Gyroscope Integration\n",
    "head_position = pd.DataFrame(columns=['time', 'posGlobal', 'quatBody', 'velocity', 'quaternion', 'W', 'X', 'Y', 'Z'])\n",
    "# Copy data from the simple head orientation estimate\n",
    "for col in ['time', 'quaternion', 'W', 'X', 'Y', 'Z']:\n",
    "    head_position[col] = head_orientation[col][:]\n",
    "\n",
    "# Gyroscope + Accelerometer\n",
    "head_position_tilt = pd.DataFrame(columns=['time', 'posGlobal', 'quatBody', 'velocity', 'quaternion', 'W', 'X', 'Y', 'Z'])\n",
    "# Copy data from the tilt corrected head orientation estimate\n",
    "for col in ['time', 'quaternion', 'W', 'X', 'Y', 'Z']:\n",
    "    head_position_tilt[col] = head_orientation_tilt_correction[col][:]\n",
    "\n",
    "# Gyroscope + Accelerometer + Magnetometer\n",
    "head_position_yaw = pd.DataFrame(columns=['time', 'posGlobal', 'quatBody', 'velocity', 'quaternion', 'W', 'X', 'Y', 'Z'])\n",
    "# Copy data from the yaw corrected head orientation estimate\n",
    "for col in ['time', 'quaternion', 'W', 'X', 'Y', 'Z']:\n",
    "    head_position_yaw[col] = head_orientation_yaw_correction[col][:]\n",
    "\n",
    "l_n = 0.15 # Distance from height of pupils to the base of the neck (0.15m)\n",
    "l_t = 0.8 # Distance from the base of the neck to the waist (0.8m)\n",
    "\n",
    "def positional_tracking(prev_pos_row, current_pos_row, current_IMU_Data_row):\n",
    "    # Extract the raw accelerometer readings from the IMU_Data row\n",
    "    # Raw accelerometer readings are used to allow a_l to more easily be extracted from a by subtracting 1 from a_l (as accelerometer measures in terms of g)\n",
    "    accX = current_IMU_Data_row['raw_accelerometer.X']\n",
    "    accY = current_IMU_Data_row['raw_accelerometer.Y']\n",
    "    accZ = current_IMU_Data_row['raw_accelerometer.Z']\n",
    "    \n",
    "    # Extract data from prev_pos_row\n",
    "    q = current_pos_row['quaternion'] # head orientation\n",
    "    prev_vel = prev_pos_row['velocity']\n",
    "    q1_prev = prev_pos_row['quatBody'] # quaternion describing the body orientation\n",
    "    \n",
    "    dt = current_IMU_Data_row['time'] - prev_pos_row['time'] # Time since last timestep\n",
    "    \n",
    "    # Convert accelerometer readings to global frame\n",
    "    acc = (0, accX, accY, accZ)\n",
    "    \n",
    "    a_g = quaternion_rotation(q, acc) # Acceleration in global frame\n",
    "\n",
    "    # Subtract 1 (gravity) from Z acceleration to get (estimated) linear acceleration\n",
    "    a_g = (a_g[1], a_g[2], a_g[3] - 1)\n",
    "\n",
    "    # Double integrate global accelerometer readings to produce body position\n",
    "    \n",
    "    # 1. Do this by converting linear acceleration into change in angular velocity as omega = (acc_linear * t)/l_t, where l_t is the length of the body\n",
    "    #    (this makes the assumption that the linear acceleration is caused by a rotation of the body, not of the neck)\n",
    "    # 2. Then, add this to the previous angular velocity to get the current angular velocity\n",
    "    # 3. Then, apply simple dead reckoning to get an estimate for orientation of the body\n",
    "    # 4. Lastly, compose this estimate with the estimate for the head position produced by the simple kinematic head model.\n",
    "    # This produces a estimated global position based on double integration of acceleration and a kinematic model of the body (and head).\n",
    "    \n",
    "    # 1. Linear acc to angular velocity conversion\n",
    "    omega = list(map(lambda acc: acc*dt/l_t, a_g)) # omega = angular velocity of body\n",
    "    \n",
    "    # 2. Add the previous velocity\n",
    "    omega = [omega[0] + prev_vel[0], omega[1] + prev_vel[1], omega[2] + prev_vel[2]]\n",
    "    \n",
    "    # 3. Apply simple dead reckoning filter\n",
    "    \n",
    "    # Normalize the angular velocity\n",
    "    mag_omega = np.sqrt(omega[0]**2 + omega[1]**2 + omega[2]**2)\n",
    "    omega_normed = list(map(lambda ang_vel: ang_vel/mag_omega, omega))\n",
    "    \n",
    "    # Calculate the angle of the rotation\n",
    "    theta = mag_omega*dt\n",
    "    \n",
    "    # Create the change quaternion from this axis-angle representation\n",
    "    d_q = (np.cos(theta/2), omega_normed[0]*np.sin(theta/2), omega_normed[1]*np.sin(theta/2), omega_normed[2]*np.sin(theta/2))\n",
    "\n",
    "    # Produce the current torso orientation by combining (multiplying) the last orientation and the change quaternion\n",
    "    q1 = quaternion_product(q1_prev, d_q)\n",
    "    \n",
    "    # Rotate (0, 0, l_t) by the torso orientation quaternion\n",
    "    r1 = quaternion_rotation(q1, (0, 0, 0, l_t))[1:] # torso position\n",
    "    \n",
    "    # Get the estimate for head position from simple kinematic head model\n",
    "    # Apply kinematic head model to produce head position\n",
    "    r = (0, 0, 0, l_n) # head position\n",
    "    r = quaternion_rotation(q, r)[1:] # head position after simple kinematic head model\n",
    "    \n",
    "    # Compose body and head position to produce the final position\n",
    "    p = np.add(r1, r) # global position\n",
    "    \n",
    "    # Return all required values\n",
    "    return p, q1, omega\n",
    "\n",
    "# Initial values - position = (0, 0, ln), body position = (0, 0, 0), velocity = (0, 0, 0)\n",
    "for table in [head_position, head_position_tilt, head_position_yaw]:\n",
    "    table.at[0, 'posGlobal'] = (0, 0, l_n+l_t)\n",
    "    table.at[0, 'quatBody'] = (1, 0, 0, 0)\n",
    "    table.at[0, 'velocity'] = (0, 0, 0)\n",
    "\n",
    "for index in range(len(head_position)):\n",
    "    # Skip the first one as this has the value of the identity quaternion\n",
    "    if index != 0:\n",
    "        prev_row = head_position.iloc[index - 1]\n",
    "        cur_row = head_position.iloc[index]\n",
    "        imu_row = IMU_Data.iloc[index]\n",
    "        # Estimate the current head position, with tilt correction\n",
    "        pos_global, posBody, velocity = positional_tracking(prev_row, cur_row, imu_row)\n",
    "        head_position.at[index, 'posGlobal'] = pos_global\n",
    "        head_position.at[index, 'quatBody'] = posBody\n",
    "        head_position.at[index, 'velocity'] = velocity\n",
    "\n",
    "# Gyroscope + Accelerometer\n",
    "for index in range(len(head_position_tilt)):\n",
    "    # Skip the first one as this has the value of the identity quaternion\n",
    "    if index != 0:\n",
    "        prev_row = head_position_tilt.iloc[index - 1]\n",
    "        cur_row = head_position_tilt.iloc[index]\n",
    "        imu_row = IMU_Data.iloc[index]\n",
    "        # Estimate the current head position, with tilt correction\n",
    "        pos_global, posBody, velocity = positional_tracking(prev_row, cur_row, imu_row)\n",
    "        head_position_tilt.at[index, 'posGlobal'] = pos_global\n",
    "        head_position_tilt.at[index, 'quatBody'] = posBody\n",
    "        head_position_tilt.at[index, 'velocity'] = velocity\n",
    "\n",
    "# Gyroscope + Accelerometer + Magnetometer\n",
    "for index in range(len(head_position_yaw)):\n",
    "    # Skip the first one as this has the value of the identity quaternion\n",
    "    if index != 0:\n",
    "        prev_row = head_position_yaw.iloc[index - 1]\n",
    "        cur_row =  head_position_yaw.iloc[index]\n",
    "        imu_row = IMU_Data.iloc[index]\n",
    "        # Estimate the current head position, with tilt correction\n",
    "        pos_global, posBody, velocity = positional_tracking(prev_row, cur_row, imu_row)\n",
    "        head_position_yaw.at[index, 'posGlobal'] = pos_global\n",
    "        head_position_yaw.at[index, 'quatBody'] = posBody\n",
    "        head_position_yaw.at[index, 'velocity'] = velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem 6 - 3D plots (full speed)\n",
    "\n",
    "# Make plots pop-out and be interactive\n",
    "%matplotlib qt\n",
    "\n",
    "# 3D plots of the 3 main axis rotations + translations of the IMU for gyroscope integration, gyroscope + acceleration, gyroscope + acceleration + magnetometer\n",
    "\n",
    "N = 6959 # Number of frames\n",
    "interval = 4 # Time between frames in ms - 1000/256 as time steps are at 256Hz\n",
    "resolution = 25 # How often frames are generated (i.e. resolution=4 means every 4th data point is plotted)\n",
    "# The resolution variable is used to ensure the animation is full speed/half speed as required. With the given values it runs at the correct speeds on my computer, but if yours is faster/slower the actual speeds will likely differ\n",
    "\n",
    "# Set the total number of frames as adjusted by resolution\n",
    "N //= resolution\n",
    "interval *= resolution\n",
    "scale = 2 # Scale of the axis limits\n",
    "\n",
    "axisX, axisY, axisZ = [0, 0, 0], [0, 0, 0], [0, 0, 0]\n",
    "\n",
    "def animate_axes(num, ax1, ax2, ax3):\n",
    "    \"\"\"Generate the new quivers\"\"\"\n",
    "    global axisX, axisY, axisZ\n",
    "    # Check that this isn't the first frame\n",
    "    if axisX[0] != 0:\n",
    "        for i in axisX, axisY, axisZ:\n",
    "            for j in i:\n",
    "                j.remove()\n",
    "        \n",
    "    num *= resolution\n",
    "    \n",
    "    x_proj_vec = (0, 1, 0, 0)\n",
    "    y_proj_vec = (0, 0, 1, 0)\n",
    "    z_proj_vec = (0, 0, 0, 1)\n",
    "    # Apply the current quaternion (indexed by num) to these vectors\n",
    "    quaternion = head_orientation.iloc[num]['quaternion']\n",
    "    tilt_corrected_quaternion = head_orientation_tilt_correction.iloc[num]['quaternion']\n",
    "    yaw_corrected_quaternion = head_orientation_yaw_correction.iloc[num]['quaternion']\n",
    "    pos_simple = head_position.iloc[num]['posGlobal']\n",
    "    pos_tilt   = head_position_tilt.iloc[num]['posGlobal']\n",
    "    pos_yaw    = head_position_yaw.iloc[num]['posGlobal']\n",
    "    \n",
    "    axes = [ax1, ax2, ax3]\n",
    "    quats = [quaternion, tilt_corrected_quaternion, yaw_corrected_quaternion]\n",
    "    poss = [pos_simple, pos_tilt, pos_yaw]\n",
    "    \n",
    "    for i in range(3):\n",
    "        ax = axes[i]\n",
    "        quat = quats[i]\n",
    "        pos = poss[i]\n",
    "        \n",
    "        x_vec = quaternion_rotation(quat, x_proj_vec)\n",
    "        y_vec = quaternion_rotation(quat, y_proj_vec)\n",
    "        z_vec = quaternion_rotation(quat, z_proj_vec)\n",
    "\n",
    "        # Add the position tracking\n",
    "        axisX[i] = ax.quiver(*pos, *x_vec[1:], color='Blue', label='X')\n",
    "        axisY[i] = ax.quiver(*pos, *y_vec[1:], color='Orange', label='Y')\n",
    "        axisZ[i] = ax.quiver(*pos, *z_vec[1:], color='Green', label='Z')\n",
    "        \n",
    "        ax.legend(loc='center left')\n",
    "\n",
    "fig = plt.figure(figsize=(24, 8))\n",
    "fig.canvas.set_window_title('Positional Tracking')\n",
    "# Gyroscope integration\n",
    "ax1 = fig.add_subplot(131, projection='3d')\n",
    "ax1.set_title(\"Gyroscope integration\")\n",
    "ax1.set_xlim3d([-1.0*scale, 1.0*scale]); ax1.set_xlabel('X')\n",
    "ax1.set_ylim3d([-1.0*scale, 1.0*scale]); ax1.set_ylabel('Y')\n",
    "ax1.set_zlim3d([-1.0*scale, 1.0*scale]); ax1.set_zlabel('Z')\n",
    "\n",
    "# Gyroscope + Accelerometer\n",
    "ax2 = fig.add_subplot(132, projection='3d')\n",
    "ax2.set_title(\"Gyroscope + Accelerometer\")\n",
    "ax2.set_xlim3d([-1.0*scale, 1.0*scale]); ax2.set_xlabel('X')\n",
    "ax2.set_ylim3d([-1.0*scale, 1.0*scale]); ax2.set_ylabel('Y')\n",
    "ax2.set_zlim3d([-1.0*scale, 1.0*scale]); ax2.set_zlabel('Z')\n",
    "\n",
    "# Gyroscope + Accelerometer + Magnetometer\n",
    "ax3 = fig.add_subplot(133, projection='3d')\n",
    "ax3.set_title(\"Gyroscope + Accelerometer + Magnetometer\")\n",
    "ax3.set_xlim3d([-1.0*scale, 1.0*scale]); ax3.set_xlabel('X')\n",
    "ax3.set_ylim3d([-1.0*scale, 1.0*scale]); ax3.set_ylabel('Y')\n",
    "ax3.set_zlim3d([-1.0*scale, 1.0*scale]); ax3.set_zlabel('Z')\n",
    "\n",
    "ani = animation.FuncAnimation(fig, animate_axes, N, fargs=(ax1, ax2, ax3), interval=interval, blit=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}