{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import sympy as sym"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# link masses\n",
    "M1 = 3\n",
    "M2 = 2\n",
    "M3 = 1\n",
    "\n",
    "# link lengths\n",
    "LX = 0.1\n",
    "LY = 0.1\n",
    "L2 = 1\n",
    "L3 = 0.5\n",
    "\n",
    "# moments of inertia (base link doesn't rotate so don't need its inertia)\n",
    "I2 = M2 * L2 ** 2 / 12\n",
    "I3 = M3 * L3 ** 2 / 12\n",
    "\n",
    "# inertia matrices\n",
    "G1 = np.diag(np.array([M1, M1, 0]))\n",
    "G2 = np.diag(np.array([M2, M2, I2]))\n",
    "G3 = np.diag(np.array([M3, M3, I3]))\n",
    "\n",
    "# gravity\n",
    "G = 9.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ManualDynamics:\n",
    "    @classmethod\n",
    "    def mass_matrix(cls, q):\n",
    "        x1, θ2, θ3 = q\n",
    "        θ23 = θ2 + θ3\n",
    "\n",
    "        m11 = M1 + M2 + M3\n",
    "        m12 = -(0.5 * M2 + M3) * L2 * np.sin(θ2) - 0.5 * M3 * L3 * np.sin(θ23)\n",
    "        m13 = -0.5 * M3 * L3 * np.sin(θ23)\n",
    "\n",
    "        m22 = (\n",
    "            (0.25 * M2 + M3) * L2 ** 2\n",
    "            + 0.25 * M3 * L3 ** 2\n",
    "            + M3 * L2 * L3 * np.cos(θ3)\n",
    "            + I2\n",
    "            + I3\n",
    "        )\n",
    "        m23 = 0.5 * M3 * L3 * (0.5 * L3 + L2 * np.cos(θ3)) + I3\n",
    "\n",
    "        m33 = 0.25 * M3 * L3 ** 2 + I3\n",
    "\n",
    "        M = np.array([[m11, m12, m13], [m12, m22, m23], [m13, m23, m33]])\n",
    "        return M\n",
    "\n",
    "    @classmethod\n",
    "    def christoffel_matrix(cls, q):\n",
    "        x1, θ2, θ3 = q\n",
    "        θ23 = θ2 + θ3\n",
    "\n",
    "        # Partial derivatives of mass matrix\n",
    "        dMdx1 = np.zeros((3, 3))\n",
    "\n",
    "        dMdθ2_12 = (\n",
    "            -0.5 * M2 * L2 * np.cos(θ2) - M3 * L2 * np.cos(θ2) - 0.5 * M3 * L3 * np.cos(θ23)\n",
    "        )\n",
    "        dMdθ2_13 = -0.5 * M3 * L3 * np.cos(θ23)\n",
    "        dMdθ2 = np.array([[0, dMdθ2_12, dMdθ2_13], [dMdθ2_12, 0, 0], [dMdθ2_13, 0, 0]])\n",
    "\n",
    "        dMdθ3_12 = -0.5 * M3 * L3 * np.cos(θ23)\n",
    "        dMdθ3_13 = -0.5 * M3 * L3 * np.cos(θ23)\n",
    "        dMdθ3_22 = -M3 * L2 * L3 * np.sin(θ3)\n",
    "        dMdθ3_23 = -0.5 * M3 * L2 * L3 * np.sin(θ3)\n",
    "        dMdθ3 = np.array([\n",
    "            [0, dMdθ3_12, dMdθ3_13],\n",
    "            [dMdθ3_12, dMdθ3_22, dMdθ3_23],\n",
    "            [dMdθ3_13, dMdθ3_23, 0],\n",
    "        ])\n",
    "\n",
    "        dMdq = np.zeros((3, 3, 3))\n",
    "        dMdq[0, :, :] = dMdx1\n",
    "        dMdq[1, :, :] = dMdθ2\n",
    "        dMdq[2, :, :] = dMdθ3\n",
    "\n",
    "        # This is equivalent to a set of for loops that fill out Γ such that\n",
    "        # Γ[i, j, k] = dMdq[k, j, i] - 0.5*dMdq[i, j, k]\n",
    "        Γ = dMdq.T - 0.5 * dMdq\n",
    "        return Γ\n",
    "\n",
    "    @classmethod\n",
    "    def coriolis_matrix(cls, q, dq):\n",
    "        Γ = cls.christoffel_matrix(q)\n",
    "        # note numpy dot broadcasting rules: this is a sum product of dq with \n",
    "        # the second last dimension of Γ\n",
    "        return np.dot(dq, Γ)\n",
    "\n",
    "    @classmethod\n",
    "    def gravity_vector(cls, q):\n",
    "        x1, θ2, θ3 = q\n",
    "        θ23 = θ2 + θ3\n",
    "        return np.array([\n",
    "            0,\n",
    "            (0.5 * M2 + M3) * G * L2 * np.cos(θ2) + 0.5 * M3 * L3 * G * np.cos(θ23),\n",
    "            0.5 * M3 * L3 * G * np.cos(θ23),\n",
    "        ])\n",
    "\n",
    "    @classmethod\n",
    "    def force(cls, q, dq, ddq):\n",
    "        M = cls.mass_matrix(q)\n",
    "        C = cls.coriolis_matrix(q, dq)\n",
    "        g = cls.gravity_vector(q)\n",
    "\n",
    "        return M @ ddq + C @ dq + g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "# We indicate that the functions should be just-in-time compiled\n",
    "# with the @partial(jax.jit, static_argnums=(0,)) lines\n",
    "\n",
    "class AutoDiffDynamics:\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def link1_pose(cls, q):\n",
    "        return jnp.array([q[0], 0, 0])\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def link2_pose(cls, q):\n",
    "        P1 = cls.link1_pose(q)\n",
    "        return P1 + jnp.array([\n",
    "            LX + 0.5 * L2 * jnp.cos(q[1]),\n",
    "            LY + 0.5 * L2 * jnp.sin(q[1]), q[1]]\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def link3_pose(cls, q):\n",
    "        P2 = cls.link2_pose(q)\n",
    "        return P2 + jnp.array([\n",
    "            0.5 * L2 * jnp.cos(q[1]) + 0.5 * L3 * jnp.cos(q[1] + q[2]),\n",
    "            0.5 * L2 * jnp.sin(q[1]) + 0.5 * L3 * jnp.sin(q[1] + q[2]),\n",
    "            q[2],\n",
    "        ])\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def mass_matrix(cls, q):\n",
    "        # Jacobians\n",
    "        J1 = jax.jacfwd(cls.link1_pose)(q)\n",
    "        J2 = jax.jacfwd(cls.link2_pose)(q)\n",
    "        J3 = jax.jacfwd(cls.link3_pose)(q)\n",
    "\n",
    "        return J1.T @ G1 @ J1 + J2.T @ G2 @ J2 + J3.T @ G3 @ J3\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def christoffel_matrix(cls, q):\n",
    "        # Here dMdq is transposed with respect to the dMdq's in the manual and\n",
    "        # symbolic implementations. Thus, the expression for Γ is also\n",
    "        # transposed, giving the same end result.\n",
    "        dMdq = jax.jacfwd(cls.mass_matrix)(q)\n",
    "        Γ = dMdq - 0.5 * dMdq.T\n",
    "        return Γ\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def coriolis_matrix(cls, q, dq):\n",
    "        Γ = cls.christoffel_matrix(q)\n",
    "        return jnp.dot(dq, Γ)\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def potential_energy(cls, q):\n",
    "        P1 = cls.link1_pose(q)\n",
    "        P2 = cls.link2_pose(q)\n",
    "        P3 = cls.link3_pose(q)\n",
    "        return G * (M1 * P1[1] + M2 * P2[1] + M3 * P3[1])\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def gravity_vector(cls, q):\n",
    "        return jax.jacfwd(cls.potential_energy)(q)\n",
    "\n",
    "    @classmethod\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def force(cls, q, dq, ddq):\n",
    "        M = cls.mass_matrix(q)\n",
    "        C = cls.coriolis_matrix(q, dq)\n",
    "        g = cls.gravity_vector(q)\n",
    "\n",
    "        return M @ ddq + C @ dq + g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SymbolicDynamics:\n",
    "    # joint configuration and velocity are functions of time\n",
    "    t = sym.symbols(\"t\")\n",
    "    q = sym.Array(\n",
    "        [sym.Function(\"x1\")(t), sym.Function(\"θ2\")(t), sym.Function(\"θ3\")(t)]\n",
    "    )\n",
    "    dq = q.diff(t)\n",
    "\n",
    "    # link poses\n",
    "    P1 = sym.Matrix([q[0], 0, 0])\n",
    "    P2 = P1 + sym.Matrix([\n",
    "        LX + 0.5 * L2 * sym.cos(q[1]),\n",
    "        LY + 0.5 * L2 * sym.sin(q[1]),\n",
    "        q[1],\n",
    "    ])\n",
    "    P3 = P2 + sym.Matrix([\n",
    "        0.5 * L2 * sym.cos(q[1]) + 0.5 * L3 * sym.cos(q[1] + q[2]),\n",
    "        0.5 * L2 * sym.sin(q[1]) + 0.5 * L3 * sym.sin(q[1] + q[2]),\n",
    "        q[2],\n",
    "    ])\n",
    "\n",
    "    # link Jacobians\n",
    "    J1 = P1.jacobian(q)\n",
    "    J2 = P2.jacobian(q)\n",
    "    J3 = P3.jacobian(q)\n",
    "\n",
    "    # mass matrix\n",
    "    M = J1.transpose() * G1 * J1 + J2.transpose() * G2 * J2 + J3.transpose() * G3 * J3\n",
    "\n",
    "    # Christoffel symbols and Coriolis matrix\n",
    "    dMdq = M.diff(q)\n",
    "    Γ = sym.permutedims(dMdq, (2, 1, 0)) - 0.5 * dMdq\n",
    "    C = sym.tensorcontraction(sym.tensorproduct(dq, Γ), (0, 2)).tomatrix()\n",
    "\n",
    "    # gravity vector\n",
    "    V = G * (M1 * P1[1] + M2 * P2[1] + M3 * P3[1])\n",
    "    g = V.diff(q)\n",
    "\n",
    "    # compile functions to numerical code\n",
    "    mass_matrix = sym.lambdify([q], M)\n",
    "    christoffel_matrix = sym.lambdify([q], Γ)\n",
    "    coriolis_matrix = sym.lambdify([q, dq], C)\n",
    "    gravity_vector = sym.lambdify([q], g)\n",
    "\n",
    "    @classmethod\n",
    "    def force(cls, q, dq, ddq):\n",
    "        M = cls.mass_matrix(q)\n",
    "        C = cls.coriolis_matrix(q, dq)\n",
    "        g = cls.gravity_vector(q)\n",
    "\n",
    "        return M @ ddq + C @ dq + g\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test a random configuration to check all implementations give the same results\n",
    "\n",
    "np.random.seed(0)\n",
    "\n",
    "# random generate configuration, velocities, and accelerations\n",
    "q = np.random.random(3) * [2, 2 * np.pi, 2 * np.pi] - [1, np.pi, np.pi]\n",
    "dq = np.random.random(3) * 2 - 1\n",
    "ddq = np.random.random(3) * 2 - 1\n",
    "\n",
    "# make sure M, Γ, C, g are the same\n",
    "M = ManualDynamics.mass_matrix(q)\n",
    "assert np.allclose(M, AutoDiffDynamics.mass_matrix(q))\n",
    "assert np.allclose(M, SymbolicDynamics.mass_matrix(q))\n",
    "\n",
    "Γ = ManualDynamics.christoffel_matrix(q)\n",
    "assert np.allclose(Γ, AutoDiffDynamics.christoffel_matrix(q))\n",
    "assert np.allclose(Γ, SymbolicDynamics.christoffel_matrix(q))\n",
    "\n",
    "C = ManualDynamics.coriolis_matrix(q, dq)\n",
    "assert np.allclose(C, AutoDiffDynamics.coriolis_matrix(q, dq))\n",
    "assert np.allclose(C, SymbolicDynamics.coriolis_matrix(q, dq))\n",
    "\n",
    "g = ManualDynamics.gravity_vector(q)\n",
    "assert np.allclose(g, AutoDiffDynamics.gravity_vector(q))\n",
    "assert np.allclose(g, SymbolicDynamics.gravity_vector(q))\n",
    "\n",
    "# generalized force vector from each method\n",
    "print(ManualDynamics.force(q, dq, ddq))\n",
    "print(AutoDiffDynamics.force(q, dq, ddq))\n",
    "print(SymbolicDynamics.force(q, dq, ddq))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next we'll do some benchmarking on each method with random motions\n",
    "# Be sure to run the above cell before this one, to give jax a chance to JIT\n",
    "# compile on its first run (which is much slower than all subsequent runs)\n",
    "\n",
    "def random_motion():\n",
    "    q = np.random.random(3) * [2, 2 * np.pi, 2 * np.pi] - [1, np.pi, np.pi]\n",
    "    dq = np.random.random(3) * 2 - 1\n",
    "    ddq = np.random.random(3) * 2 - 1\n",
    "    return q, dq, ddq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit q, dq, ddq = random_motion()\n",
    "ManualDynamics.force(q, dq, ddq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit q, dq, ddq = random_motion()\n",
    "AutoDiffDynamics.force(q, dq, ddq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit q, dq, ddq = random_motion()\n",
    "SymbolicDynamics.force(q, dq, ddq)"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}