RotationWithNumba.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-27T12:58:32.842364Z",
     "start_time": "2018-06-27T12:58:32.269446Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import quaternion\n",
    "import numba as nb\n",
    "\n",
    "@nb.njit\n",
    "def sv_plus_rxv(q, v, w):\n",
    "    w[0] = q[0] * v[0] + q[2]*v[2] - q[3]*v[1]\n",
    "    w[1] = q[0] * v[1] + q[3]*v[0] - q[1]*v[2]\n",
    "    w[2] = q[0] * v[2] + q[1]*v[1] - q[2]*v[0]\n",
    "    return\n",
    "\n",
    "@nb.njit\n",
    "def v_plus_2rxv_over_m(q, v, w, two_over_m):\n",
    "    v[0] = v[0] + two_over_m * (q[2]*w[2] - q[3]*w[1])\n",
    "    v[1] = v[1] + two_over_m * (q[3]*w[0] - q[1]*w[2])\n",
    "    v[2] = v[2] + two_over_m * (q[1]*w[1] - q[2]*w[0])\n",
    "    return\n",
    "\n",
    "@nb.njit\n",
    "def rotate(q, v, w):\n",
    "    m = q[1]*q[1]+q[2]*q[2]+q[3]*q[3]\n",
    "    sv_plus_rxv(q, v, w)\n",
    "    v_plus_2rxv_over_m(q, v, w, 2.0/m)\n",
    "    return\n",
    "\n",
    "@nb.njit\n",
    "def rotate_and_update(Q, v, decay_constant, relaxation_constant):\n",
    "    w = np.empty((3,))  # temporary working memory; allocate once to save time\n",
    "    for i in range(v.shape[0]):\n",
    "        rotate(Q[i], v[i], w)\n",
    "        v[i, 0] *= decay_constant  # scale x component\n",
    "        v[i, 1] *= decay_constant  # scale y component separately for speed\n",
    "        v[i, 2] += (1 - v[i, 2]) * relaxation_constant  # add relaxation for z component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-27T12:58:33.397460Z",
     "start_time": "2018-06-27T12:58:32.844538Z"
    }
   },
   "outputs": [],
   "source": [
    "Q = np.random.rand(3969, 4)\n",
    "v = np.random.rand(3969, 3)\n",
    "decay_constant = 0.99  # Just making this up\n",
    "relaxation_constant = 0.01  # Just making this up\n",
    "rotate_and_update(Q, v, decay_constant, relaxation_constant)  # Run it once to compile before timing it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-27T12:58:43.994242Z",
     "start_time": "2018-06-27T12:58:33.399484Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "131 µs ± 2.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit Q = np.random.rand(3969, 4); v = np.random.rand(3969, 3)\n",
    "rotate_and_update(Q, v, decay_constant, relaxation_constant)  # operates in place"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-27T12:58:44.084346Z",
     "start_time": "2018-06-27T12:58:43.996183Z"
    }
   },
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def update_v_rot(v_rot):\n",
    "    for i in range(v_rot.shape[0]):\n",
    "        v_rot[i, 0] *= decay_constant  # scale x component\n",
    "        v_rot[i, 1] *= decay_constant  # scale y component separately for speed\n",
    "        v_rot[i, 2] += (1 - v_rot[i, 2]) * relaxation_constant  # add relaxation for z component\n",
    "\n",
    "update_v_rot(np.random.rand(2, 3))  # Run it once to compile before timing it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-27T12:58:48.188923Z",
     "start_time": "2018-06-27T12:58:44.086332Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "497 µs ± 7.21 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit Q = quaternion.as_quat_array(np.random.rand(3969, 4)); v = np.random.rand(3969, 3)\n",
    "v = (quaternion.as_rotation_matrix(Q) @ v[..., np.newaxis])[..., 0]\n",
    "update_v_rot(v)  # operates in place"
   ]
  },
  {
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}