{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to `RigidBodyDynamics`, we'll be using the `StaticArrays` package, used throughout `RigidBodyDynamics`, which provides stack-allocated, fixed-size arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "using RigidBodyDynamics\n",
    "using StaticArrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a double pendulum `Mechanism`\n",
    "\n",
    "We're going to create a simple `Mechanism` that represents a [double pendulum](https://en.wikipedia.org/wiki/Double_pendulum). The `Mechanism` type can be thought of as an interconnection of rigid bodies and joints.\n",
    "\n",
    "We'll start by creating a 'root' rigid body, representing the fixed world, and using it to create a new `Mechanism`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vertex: world (root)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g = -9.81 # gravitational acceleration in z-direction\n",
    "world = RigidBody{Float64}(\"world\")\n",
    "doublependulum = Mechanism(world; gravity = SVector(0, 0, g))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the `RigidBody` type is parameterized on the 'scalar type', here `Float64`.\n",
    "\n",
    "We'll now add a second body, called 'upper link', to the `Mechanism`. We'll attach it to the world with a revolute joint, with the $y$-axis as the axis of rotation. We'll start by creating a `SpatialInertia`, which stores the inertial properties of the new body:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SpatialInertia expressed in \"upper_link\":\n",
       "mass: 1.0\n",
       "center of mass: Point3D in \"upper_link\": [0.0,0.0,-0.5]\n",
       "moment of inertia:\n",
       "[0.0 0.0 0.0; 0.0 0.333 0.0; 0.0 0.0 0.0]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "axis = SVector(0., 1., 0.) # joint axis\n",
    "I_1 = 0.333 # moment of inertia about joint axis\n",
    "c_1 = -0.5 # center of mass location with respect to joint axis\n",
    "m_1 = 1. # mass\n",
    "frame1 = CartesianFrame3D(\"upper_link\") # the reference frame in which the spatial inertia will be expressed\n",
    "inertia1 = SpatialInertia(frame1, I_1 * axis * axis', SVector(0, 0, c_1), m_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the created `SpatialInertia` is annotated with the frame in which it is expressed (in the form of a `CartesianFrame3D`). This is a common theme among `RigidBodyDynamics` objects. Storing frame information with the data obviates the need for the complicated variable naming conventions that are used in some other libraries to disambiguate the frame in which quantities are expressed. It also enables automated reference frame checks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll now create the second body:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RigidBody: \"upper_link\""
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "upperlink = RigidBody(inertia1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and a new revolute joint called 'shoulder':"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joint \"shoulder\": Revolute joint with axis [0.0,1.0,0.0]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shoulder = Joint(\"shoulder\", Revolute(axis))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating a `Joint` automatically constructs two new `CartesianFrame3D` objects: a frame directly before the joint, and one directly after. To attach the new body to the world by this joint, we'll have to specify where the frame before the joint is located on the parent body (here, the world):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Transform3D from \"before_shoulder\" to \"world\":\n",
       "rotation: 0.0 rad about [1.0,0.0,0.0], translation: [0.0,0.0,0.0]\n"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "beforeShoulderToWorld = Transform3D{Float64}(frame_before(shoulder), default_frame(world)) # creates an identity transform"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can attach the upper link to the world:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vertex: world (root)\n",
       "  Vertex: upper_link, Edge: shoulder"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attach!(doublependulum, world, shoulder, beforeShoulderToWorld, upperlink)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which changes the tree representation of the `Mechanism`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can attach the lower link in similar fashion:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vertex: world (root)\n",
       "  Vertex: upper_link, Edge: shoulder\n",
       "    Vertex: lower_link, Edge: elbow"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l_1 = -1. # length of the upper link\n",
    "I_2 = 0.333 # moment of inertia about joint axis\n",
    "c_2 = -0.5 # center of mass location with respect to joint axis\n",
    "m_2 = 1. # mass\n",
    "inertia2 = SpatialInertia(CartesianFrame3D(\"lower_link\"), I_2 * axis * axis', SVector(0, 0, c_2), m_2)\n",
    "lowerlink = RigidBody(inertia2)\n",
    "elbow = Joint(\"elbow\", Revolute(axis))\n",
    "beforeElbowToAfterShoulder = Transform3D(frame_before(elbow), shoulder.frameAfter, SVector(0, 0, l_1))\n",
    "attach!(doublependulum, upperlink, elbow, beforeElbowToAfterShoulder, lowerlink)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now our double pendulum `Mechanism` is complete."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: instead of defining the `Mechanism` in this way, it is also possible to load in a [URDF](http://wiki.ros.org/urdf) file (an XML file format used in ROS), using the `parse_urdf` function, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vertex: world (root)\n",
       "  Vertex: base_link, Edge: base_link_to_world\n",
       "    Vertex: upper_link, Edge: shoulder\n",
       "      Vertex: lower_link, Edge: elbow"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filename = \"../test/urdf/Acrobot.urdf\"\n",
    "parse_urdf(Float64, filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The state of a `Mechanism`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `Mechanism` stores the joint/rigid body layout, but no state information. State information is separated out into a `MechanismState` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MechanismState{Float64, Float64, Float64}(…)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state = MechanismState(Float64, doublependulum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first set the configurations and velocities of the joints:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "configuration(state, shoulder)[:] = 0.3\n",
    "configuration(state, elbow)[:] = 0.4\n",
    "velocity(state, shoulder)[:] = 1.\n",
    "velocity(state, elbow)[:] = 2.;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Important**: a `MechanismState` contains cache variables that depend on the configurations and velocities of the joints. These need to be invalidated when the configurations and velocities are changed. To do this, call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "setdirty!(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The joint configurations and velocities are stored as `Vector`s (denoted $q$ and $v$ respectively in this package) inside the `MechanismState` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Array{Float64,1}:\n",
       " 0.3\n",
       " 0.4"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q = configuration(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Array{Float64,1}:\n",
       " 1.0\n",
       " 2.0"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v = velocity(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kinematics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to do kinematics. Here's how you transform a point at the origin of the frame after the elbow joint to world frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Point3D in \"world\": [-0.29552,0.0,-0.955336]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transform(state, Point3D(elbow.frameAfter, zeros(SVector{3})), default_frame(world))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Other objects like `Wrench`es, `Twist`s, and `SpatialInertia`s can be transformed in similar fashion."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also ask for the homogeneous transform to world:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Transform3D from \"after_elbow\" to \"world\":\n",
       "rotation: 0.7000000000000001 rad about [0.0,1.0,0.0], translation: [-0.29552,0.0,-0.955336]\n"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transform_to_root(state, elbow.frameAfter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or a relative transform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Transform3D from \"after_elbow\" to \"after_shoulder\":\n",
       "rotation: 0.40000000000000024 rad about [2.65182e-8,1.0,0.0], translation: [0.0,0.0,-1.0]\n"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relative_transform(state, elbow.frameAfter, shoulder.frameAfter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and here's the center of mass of the double pendulum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Point3D in \"world\": [-0.382695,0.0,-0.907713]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "center_of_mass(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `MechanismState` can also be used to compute quantities related to the dynamics of the `Mechanism`. Here we compute the mass matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 Symmetric{Float64,Array{Float64,2}}:\n",
       " 2.58706  0.79353\n",
       " 0.79353  0.333  "
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mass_matrix(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that there is also a zero-allocation version, `mass_matrix!` (the `!` at the end of a method is a Julia convention signifying that the function is 'in-place', i.e. modifies its input data)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can do inverse dynamics as follows (note again that there is a non-allocating version of this method as well):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Array{Float64,1}:\n",
       " 13.5055 \n",
       "  5.94066"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v̇ = [2.; 3.] # the joint acceleration vector, i.e., the time derivative of the joint velocity vector v\n",
    "inverse_dynamics(state, v̇)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's simulate the double pendulum for 5 seconds, starting from the state we defined earlier. For this, we can use the basic `simulate` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ts, qs, vs = simulate(state, 5., Δt = 1e-3);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`simulate` returns a vector of times (`ts`) and associated joint configurations (`qs`) and velocities (`vs`). You can of course plot the trajectories using your favorite plotting package (see e.g. [Plots.jl](https://github.com/JuliaPlots/Plots.jl)). The [RigidBodyTreeInspector](https://github.com/tkoolen/RigidBodyTreeInspector.jl) package can also be used for 3D animation of the double pendulum in action."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A lower level interface for simulation/ODE integration with more options is also available. Consult the documentation for more information."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 0.5.0",
   "language": "julia",
   "name": "julia-0.5"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}