{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linear and angular velocity\n",
    "\n",
    "We have seen how to look up and transform positions and orientations across reference frames in the previous section. Working with velocities adds some complexity that will be explained in this section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "Note\n",
    "    \n",
    "The following examples require the `matplotlib` library.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import rigid_body_motion as rbm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (6, 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up\n",
    "\n",
    "Like in the previous section, we first set up the world frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.register_frame(\"world\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For simplicity, all reference frames will share the same timestamps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_timestamps = 5\n",
    "ts = np.linspace(0, 5, n_timestamps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a body moving 5 meters in the x direction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body_world = np.zeros((n_timestamps, 3))\n",
    "p_body_world[:, 0] = np.linspace(0, 5, n_timestamps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, this body will rotate 90° around the z axis. We use helper functions from the `quaternion` package for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from quaternion import as_float_array, from_euler_angles\n",
    "\n",
    "o_body_world = as_float_array(\n",
    "    from_euler_angles(np.linspace(0, np.pi / 4, n_timestamps), 0, 0)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can attach a reference frame to this body:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.register_frame(\n",
    "    \"body\",\n",
    "    translation=p_body_world,\n",
    "    rotation=o_body_world,\n",
    "    timestamps=ts,\n",
    "    parent=\"world\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now define a second moving body whose motion we describe wrt the frame of the first body. It is located at 1 meter in the y direction and moves 1 meter in the negative x direction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body2_body = np.zeros((n_timestamps, 3))\n",
    "p_body2_body[:, 0] = -np.linspace(0, 1, n_timestamps)\n",
    "p_body2_body[:, 1] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This body also rotates, but this time around the y axis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "o_body2_body = as_float_array(\n",
    "    from_euler_angles(0, np.linspace(0, np.pi / 4, n_timestamps), 0)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can register a frame attached to the second body as a child frame of the first body frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.register_frame(\n",
    "    \"body2\",\n",
    "    translation=p_body2_body,\n",
    "    rotation=o_body2_body,\n",
    "    timestamps=ts,\n",
    "    parent=\"body\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the position and orientation of both bodies wrt the world frame. We use the [lookup_pose()](_generated/rigid_body_motion.lookup_pose.rst) method to obtain the position of the second body in the world frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body2_world, o_body2_world = rbm.lookup_pose(\"body2\", \"world\")\n",
    "\n",
    "fig = plt.figure(figsize=(10, 5))\n",
    "\n",
    "ax = fig.add_subplot(121, projection=\"3d\")\n",
    "rbm.plot.points(p_body_world, ax=ax, fmt=\"yo-\")\n",
    "rbm.plot.quaternions(o_body_world, base=p_body_world, ax=ax)\n",
    "\n",
    "ax = fig.add_subplot(122, projection=\"3d\", sharex=ax, sharey=ax, sharez=ax)\n",
    "rbm.plot.points(p_body2_world, ax=ax, fmt=\"co-\")\n",
    "rbm.plot.quaternions(o_body2_world, base=p_body2_world, ax=ax)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear velocity\n",
    "\n",
    "The linear velocity of a frame wrt another frame can be calculated with the [lookup_linear_velocity()](_generated/rigid_body_motion.lookup_linear_velocity.rst) method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_body2_body = rbm.lookup_linear_velocity(\"body2\", \"body\")\n",
    "v_body2_body[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the velocity of the second body wrt the first body $v_{B2/B}$ is 0.2 m/s in the negative x direction. Next, we will see how we can obtain the velocity of the second body wrt the world frame $W$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the reference frame\n",
    "\n",
    "When transforming linear velocities across frames, we need to use the so-called \"three term velocity formula\". In this case, we have the velocity of the _moving frame_ $B2$ wrt the _reference frame_ $B$. We can transform the current reference frame $B$ to the new frame $W$ as follows:\n",
    "\n",
    "$$v_{B2/W} = \\underbrace{v_{B2/B}}_{\\text{Input}} + \\underbrace{v_{B/W} + \\omega_{B/W} \\times t_{B2/B}}_{\\text{Lookup}}$$\n",
    "\n",
    "In addition to the velocities between the reference frames $B$ and $W$, this formula also requires the translation between the moving frame $B2$ and the original reference frame $B$. This is why we also need to specify the `moving_frame` argument when using [transform_linear_velocity()](_generated/rigid_body_motion.transform_linear_velocity.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_body2_world = rbm.transform_linear_velocity(\n",
    "    v_body2_body, outof=\"body\", into=\"world\", moving_frame=\"body2\", timestamps=ts,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, we can also use [lookup_linear_velocity()](_generated/rigid_body_motion.lookup_linear_velocity.rst) to lookup the position of $B2$ wrt $W$ and differentiate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_body2_world_lookup = rbm.lookup_linear_velocity(\"body2\", \"world\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following short helper function can be used to compare the two methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_velocities(transform, lookup, timestamps=None, mode=\"linear\"):\n",
    "    \"\"\" Compare velocities from transform and lookup. \"\"\"\n",
    "    fig, axarr = plt.subplots(2, 1, sharex=True, sharey=True)\n",
    "\n",
    "    ylabel = f\"{mode.capitalize()} velocity ({'rad/s' if mode == 'angular' else 'm/s'})\"\n",
    "\n",
    "    axarr[0].plot(timestamps, transform)\n",
    "    axarr[0].set_ylabel(ylabel)\n",
    "    axarr[0].set_title(\"Transform\")\n",
    "\n",
    "    axarr[1].plot(timestamps, lookup)\n",
    "    axarr[1].set_xlabel(\"Time (s)\")\n",
    "    axarr[1].set_ylabel(ylabel)\n",
    "    axarr[1].set_title(\"Lookup\")\n",
    "    axarr[1].legend([\"x\", \"y\", \"z\"])\n",
    "\n",
    "    fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_velocities(v_body2_world, v_body2_world_lookup, ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see a slight discrepancy due to numerical inconsistencies in the derivative calculation. However, these errors are reduced at higher sampling rates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the moving frame\n",
    "\n",
    "In a different scenario, we might be given the velocity of the first body wrt the world frame and want to transform the _moving frame_ from $B$ to $B2$ to compute the velocity of the second body wrt $W$. The same formula applies, although this time the input is $v_{B/W}$:\n",
    "\n",
    "$$v_{B2/W} = \\underbrace{v_{B/W}}_{\\text{Input}} + \\underbrace{v_{B2/B} + \\omega_{B/W} \\times t_{B2/B}}_{\\text{Lookup}}$$\n",
    "\n",
    "When using [transform_linear_velocity()](_generated/rigid_body_motion.transform_linear_velocity.rst) we need to be careful that the velocity is represented in the coordinates of the frame we want to transform. Therefore, $v_{B/W}$ has to be represented in $B$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_body_world = rbm.lookup_linear_velocity(\"body\", \"world\", represent_in=\"body\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use [transform_linear_velocity()](_generated/rigid_body_motion.transform_linear_velocity.rst) with `what=\"moving_frame\"` to transform from $B$ to $B2$. Since the method needs to look up $\\omega_{B/W}$, this time we need to provide the `reference_frame` of the velocity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_body2_world = rbm.transform_linear_velocity(\n",
    "    v_body_world,\n",
    "    outof=\"body\",\n",
    "    into=\"body2\",\n",
    "    reference_frame=\"world\",\n",
    "    what=\"moving_frame\",\n",
    "    timestamps=ts,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare this result against the lookup method again. Note that the transformation also changes the representation frame to the new frame, which is why the resulting velocity is different compared to the first case:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_body2_world_lookup = rbm.lookup_linear_velocity(\"body2\", \"world\", represent_in=\"body2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_velocities(v_body2_world, v_body2_world_lookup, ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Angular velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Angular velocities can be looked up with [lookup_angular_velocity()](_generated/rigid_body_motion.lookup_angular_velocity.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_body2_body = rbm.lookup_angular_velocity(\"body2\", \"body\")\n",
    "w_body2_body[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the reference frame\n",
    "\n",
    "Transforming the reference frame of angular velocity is similar to the case of linear velocity, although the formula is a lot simpler:\n",
    "\n",
    "$$\\omega_{B2/W} = \\underbrace{\\omega_{B2/B}}_{\\text{Input}} + \\underbrace{\\omega_{B/W}}_{\\text{Lookup}}$$\n",
    "\n",
    "Because of this, [transform_angular_velocity()](_generated/rigid_body_motion.transform_angular_velocity.rst) also has a simpler interface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_body2_world = rbm.transform_angular_velocity(\n",
    "    w_body2_body, outof=\"body\", into=\"world\", timestamps=ts\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comparing the transform to the lookup shows no differences in the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_body2_world_lookup = rbm.lookup_angular_velocity(\"body2\", \"world\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_velocities(w_body2_world, w_body2_world_lookup, ts, mode=\"angular\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforming the moving frame\n",
    "\n",
    "As before, we can also transform the moving frame:\n",
    "\n",
    "$$\\omega_{B2/W} = \\underbrace{\\omega_{B/W}}_{\\text{Input}} + \\underbrace{\\omega_{B2/B}}_{\\text{Lookup}}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_body_world = rbm.lookup_angular_velocity(\"body\", \"world\", represent_in=\"body\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_body2_world = rbm.transform_angular_velocity(\n",
    "    w_body_world, outof=\"body\", into=\"body2\", what=\"moving_frame\", timestamps=ts,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w_body2_world_lookup = rbm.lookup_angular_velocity(\"body2\", \"world\", represent_in=\"body2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_velocities(w_body2_world, w_body2_world_lookup, ts, mode=\"angular\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:rbm]",
   "language": "python",
   "name": "conda-env-rbm-py"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}