{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Gravity alignment code\n\n@author: theja\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt \nfrom scipy.spatial.transform import Rotation\nimport numpy as np \nnp.random.seed(82319)\n\n\nt = np.linspace(0,3,75)\nxyz = np.zeros((t.size,3))\nh = 1.5\nv = 10\ntheta = np.pi/3\ng = 9.8\nxyz[:,2] = h + v*t*np.sin(theta) - (g*t**2)/2\nxyz[:,1] =  v*t*np.cos(theta)\n# add noise to the estimates \nxyz += np.random.normal(0,0.2,xyz.size).reshape(xyz.shape)\n# also rotate all the coordinates a bit\n\ndef rotate_point(point):\n    rotmat = Rotation.from_euler('zxy', [10,-15,-70],degrees=True).as_matrix()\n    return np.matmul(rotmat, point)\n\nxyz_rot = np.apply_along_axis(rotate_point, 1, xyz)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fps = 25\ndef smooth_2nddeg(xyz):\n    t = np.arange(xyz.shape[0])\n    xyz_fit = [np.polyfit(t,xyz[:,each], 2) for each in range(xyz.shape[1])]\n    # now calculate the smoothed positions\n    xyz_smooth = [np.polyval(each, t) for each in xyz_fit]\n    xyz_sm = np.column_stack([each for each in xyz_smooth])\n    return xyz_sm\n\ndef calc_acc(xyz_smooth, fps):\n    v = np.apply_along_axis(np.diff, 0, xyz_smooth)\n    acc = np.apply_along_axis(np.diff, 0, v)/(1.0/(fps**2))\n    return acc\n\ndef smooth_and_acc(xyz, fps=25):\n    smooth_xyz = smooth_2nddeg(xyz)\n    acc = calc_acc(smooth_xyz, fps)\n    return acc\n\ndef calc_norm(X):\n    return np.linalg.norm(X)\n\ndef row_calc_norm(X):\n    return np.apply_along_axis(calc_norm, 1, X)\n    \n\n\nown_sm = smooth_and_acc(xyz_rot)\norig_sm = smooth_and_acc(xyz)\n\nmean_acc = np.mean(own_sm,0)\nmean_acc/np.linalg.norm(mean_acc)\n\n# if __name__ == '__main__':\n#     plt.figure()\n#     plt.plot(xyz_rot)"
      ]
    }
  ],
  "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.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}