{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "brutal-distribution",
   "metadata": {},
   "source": [
    "# Reference frames\n",
    "\n",
    "rigid_body_motion provides a flexible high-performance framework for working offline with motion data. The core of this framework is a mechanism for constructing trees of both static and dynamic reference frames that supports automatic lookup and application of transformations across the tree."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pacific-amplifier",
   "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,
   "id": "genuine-change",
   "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",
   "id": "hungarian-tribute",
   "metadata": {},
   "source": [
    "## Static frames\n",
    "\n",
    "We will begin by defining a world reference frame using the [ReferenceFrame](_generated/rigid_body_motion.ReferenceFrame.rst) class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bearing-freeware",
   "metadata": {},
   "outputs": [],
   "source": [
    "rf_world = rbm.ReferenceFrame(\"world\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acknowledged-knowing",
   "metadata": {},
   "source": [
    "Now we can add a second reference frame as a child of the world frame. This frame is translated by 5 meters in the x direction and rotated 90° around the z axis. Note that rotations are represented as [unit quaternions](https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation) by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "laden-leisure",
   "metadata": {},
   "outputs": [],
   "source": [
    "rf_observer = rbm.ReferenceFrame(\n",
    "    \"observer\",\n",
    "    parent=rf_world,\n",
    "    translation=(5, 0, 0),\n",
    "    rotation=(np.sqrt(2) / 2, 0, 0, np.sqrt(2) / 2),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "floating-night",
   "metadata": {},
   "source": [
    "We can show the reference frame tree with the [render_tree](_generated/rigid_body_motion.render_tree.rst) function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "residential-auckland",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.render_tree(rf_world)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aggressive-rubber",
   "metadata": {},
   "source": [
    "It is also possible to show a 3d plot of static reference frames with [plot.reference_frame()](_generated/rigid_body_motion.plot.reference_frame.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fatal-opinion",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(rf_world, ax=ax)\n",
    "rbm.plot.reference_frame(rf_observer, rf_world, ax=ax)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "whole-table",
   "metadata": {},
   "source": [
    "To facilitate referring to previously defined frames, the library has a registry where frames can be stored by name with [ReferenceFrame.register()](_generated/rigid_body_motion.ReferenceFrame.register.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "protective-division",
   "metadata": {},
   "outputs": [],
   "source": [
    "rf_world.register()\n",
    "rf_observer.register()\n",
    "rbm.registry"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "according-northern",
   "metadata": {},
   "source": [
    "## Transformations\n",
    "\n",
    "Now that we've set up a basic tree, we can use it to transform motion between reference frames. We use the [lookup_transform()](_generated/rigid_body_motion.lookup_transform.rst) method to obtain the transformation from the world to the observer frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "requested-wilderness",
   "metadata": {},
   "outputs": [],
   "source": [
    "t, r = rbm.lookup_transform(outof=\"world\", into=\"observer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "offensive-disco",
   "metadata": {},
   "source": [
    "This transformation consists of a translation $t$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cleared-cookie",
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "strange-birmingham",
   "metadata": {},
   "source": [
    "and a rotation $r$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "third-revision",
   "metadata": {},
   "outputs": [],
   "source": [
    "r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "positive-township",
   "metadata": {},
   "source": [
    "### Position\n",
    "\n",
    "rigid_body_motion uses the convention that a transformation is a rotation followed by a translation. Here, when applying the transformation to a point $p$ expressed with respect to (wrt) the world frame $W$ it yields the point wrt the observer frame $O$:\n",
    "\n",
    "$$p_O = \\operatorname{rot}\\left(r, p_W\\right) + t$$\n",
    "\n",
    "The $\\operatorname{rot}()$ function denotes the [rotation of a vector by a quaternion](https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation#Using_quaternions_as_rotations)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "isolated-engineer",
   "metadata": {},
   "source": [
    "Let's assume we have a rigid body located at 2 meters in the x direction from the origin of the world reference frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "colonial-emerald",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body_world = np.array((2, 0, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prepared-house",
   "metadata": {},
   "source": [
    "We can add the body position to the plot with [plot.points()](_generated/rigid_body_motion.plot.points.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceramic-objective",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(rf_world, ax=ax)\n",
    "rbm.plot.reference_frame(rf_observer, rf_world, ax=ax)\n",
    "rbm.plot.points(p_body_world, ax=ax, fmt=\"yo\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "white-basket",
   "metadata": {},
   "source": [
    "We can use the above formula to transform the position of the body into the observer frame. The [rotate_vectors()](_generated/rigid_body_motion.rotate_vectors.rst) method implements the rotation of a vector by a quaternion:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "understood-jacksonville",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body_observer = rbm.rotate_vectors(r, p_body_world) + t\n",
    "p_body_observer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sustainable-ballet",
   "metadata": {},
   "source": [
    "As expected, the resulting position of the body is 3 meters from the observer in the y direction. For convenience, the [transform_points()](_generated/rigid_body_motion.transform_points.rst) method performs all of the above steps:\n",
    "\n",
    "1. Lookup of the frames by name in the registry (if applicable)\n",
    "2. Computing the transformation from the source to the target frame\n",
    "3. Applying the transformation to the point(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fatty-baking",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body_observer = rbm.transform_points(p_body_world, outof=\"world\", into=\"observer\")\n",
    "p_body_observer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "varied-shock",
   "metadata": {},
   "source": [
    "### Orientation\n",
    "\n",
    "Orientations expressed in quaternions are transformed by quaternion multiplication:\n",
    "\n",
    "$$o_O = r \\cdot o_W $$\n",
    "\n",
    "This multiplication is implemented in the [qmul()](_generated/rigid_body_motion.qmul.rst) function to which you can pass an arbtrary number of quaternions to multiply. Assuming the body is oriented in the same direction as the world frame, transforming the orientation into the observer frame results in a rotation around the yaw axis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "juvenile-assumption",
   "metadata": {},
   "outputs": [],
   "source": [
    "o_body_world = np.array((1, 0, 0, 0))\n",
    "rbm.qmul(r, o_body_world)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "traditional-talent",
   "metadata": {},
   "source": [
    "We can add the orientation to the plot with [plot.quaternions()](_generated/rigid_body_motion.plot.quaternions.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "single-exhibit",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(rf_world, ax=ax)\n",
    "rbm.plot.reference_frame(rf_observer, rf_world, ax=ax)\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",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "premier-stable",
   "metadata": {},
   "source": [
    "Again, for convenience, the [transform_quaternions()](_generated/rigid_body_motion.transform_quaternions.rst) function can be used in the same way as [transform_points()](_generated/rigid_body_motion.transform_points.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fluid-spine",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.transform_quaternions(o_body_world, outof=\"world\", into=\"observer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "analyzed-fabric",
   "metadata": {},
   "source": [
    "### Vectors\n",
    "\n",
    "Let's assume the body moves in the x direction with a velocity of 1 m/s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "northern-thunder",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_body_world = np.array((1, 0, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "smoking-polls",
   "metadata": {},
   "source": [
    "We can add the velocity to the plot with [plot.vectors()](_generated/rigid_body_motion.plot.vectors.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "noble-secretariat",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(rf_world, ax=ax)\n",
    "rbm.plot.reference_frame(rf_observer, rf_world, ax=ax)\n",
    "rbm.plot.points(p_body_world, ax=ax, fmt=\"yo\")\n",
    "rbm.plot.vectors(v_body_world, base=p_body_world, ax=ax, color=\"y\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "legitimate-nigeria",
   "metadata": {},
   "source": [
    "From the point of view of the observer, the body moves with the same speed, but in the negative y direction. Therefore, we need to apply a coordinate transformation to represent the velocity vector in the observer frame:\n",
    "\n",
    "$$ v_O = \\operatorname{rot}\\left(r, v_W\\right) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "leading-plaintiff",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.rotate_vectors(r, v_body_world)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "consistent-embassy",
   "metadata": {},
   "source": [
    "Like before, the [transform_vectors()](_generated/rigid_body_motion.transform_vectors.rst) function can be used in the same way as [transform_points()](_generated/rigid_body_motion.transform_points.rst):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "creative-rebecca",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.transform_vectors(v_body_world, outof=\"world\", into=\"observer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "serial-burns",
   "metadata": {},
   "source": [
    "## Moving frames\n",
    "\n",
    "Now, let's assume that the body moves from the origin of the world frame to the origin of the observer frame in 5 steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "nominated-harbor",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body_world = np.zeros((5, 3))\n",
    "p_body_world[:, 0] = np.linspace(0, 5, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "industrial-nevada",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(rf_world, ax=ax)\n",
    "rbm.plot.reference_frame(rf_observer, rf_world, ax=ax)\n",
    "rbm.plot.points(p_body_world, ax=ax, fmt=\"yo-\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "desperate-concert",
   "metadata": {},
   "source": [
    "We will now attach a reference frame to the moving body to explain the handling of moving reference frames. For this, we need to associate the positions of the body with corresponding timestamps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "oriented-standard",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_body = np.arange(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ancient-terrorism",
   "metadata": {},
   "source": [
    "Let's construct the moving body frame and add it to the registry. We will use the [register_frame()](_generated/rigid_body_motion.register_frame.rst) convenience method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "described-sally",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.register_frame(\"body\", translation=p_body_world, timestamps=ts_body, parent=\"world\")\n",
    "rbm.render_tree(\"world\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "automatic-petite",
   "metadata": {},
   "source": [
    "If we transform a static point from the world into the body frame its position will change over time, which is why [transform_points()](_generated/rigid_body_motion.transform_points.rst) will return an array of points even though we pass only a single point:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "impossible-broadcasting",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.transform_points((2, 0, 0), outof=\"world\", into=\"body\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "forty-frank",
   "metadata": {},
   "source": [
    "One of the central features of the reference frame mechanism is its ability to consolidate arrays of timestamped motion even when the timestamps don't match. To illustrate this, let's create a second body moving in the y direction in world coordinates whose timestamps are offset by 0.5 seconds compared to the first body:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "regional-output",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body2_world = np.zeros((5, 3))\n",
    "p_body2_world[:, 1] = np.linspace(0, 2, 5)\n",
    "ts_body2 = ts_body - 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "controversial-drama",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(rf_world, ax=ax)\n",
    "rbm.plot.reference_frame(rf_observer, rf_world, ax=ax)\n",
    "rbm.plot.points(p_body_world, ax=ax, fmt=\"yo-\")\n",
    "rbm.plot.points(p_body2_world, ax=ax, fmt=\"co-\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "under-father",
   "metadata": {},
   "source": [
    "Transforming the position of the second body into the frame of the first body still works, despite the timestamp mismatch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "photographic-northeast",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_body2_body, ts_body2_body = rbm.transform_points(\n",
    "    p_body2_world,\n",
    "    outof=\"world\",\n",
    "    into=\"body\",\n",
    "    timestamps=ts_body2,\n",
    "    return_timestamps=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "published-practice",
   "metadata": {},
   "source": [
    "This is because behind the scenes, [transform_points()](_generated/rigid_body_motion.transform_points.rst) matches the timestamps of the array to transform with those of the transformation across the tree by\n",
    "\n",
    "1. computing the range of timestamps for which the transformation is defined,\n",
    "2. intersecting that range with the range of timestamps to be transformed and\n",
    "3. interpolating the resulting transformation across the tree to match the timestamps of the array.\n",
    "\n",
    "Note that we specified `return_timestamps=True` to obtain the timestamps of the transformed array as they are different from the original timestamps. Let's plot the position of both bodies wrt the world frame as well as the position of the second body wrt the first body to see how the timestamp matching works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "modular-charger",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axarr = plt.subplots(3, 1, sharex=True, sharey=True)\n",
    "\n",
    "axarr[0].plot(ts_body, p_body_world, \"*-\")\n",
    "axarr[0].set_ylabel(\"Position (m)\")\n",
    "axarr[0].set_title(\"First body wrt world frame\")\n",
    "axarr[0].grid(\"on\")\n",
    "\n",
    "axarr[1].plot(ts_body2, p_body2_world, \"*-\")\n",
    "axarr[1].set_ylabel(\"Position (m)\")\n",
    "axarr[1].set_title(\"Second body wrt world frame\")\n",
    "axarr[1].grid(\"on\")\n",
    "\n",
    "axarr[2].plot(ts_body2_body, p_body2_body, \"*-\")\n",
    "axarr[2].set_xlabel(\"Time (s)\")\n",
    "axarr[2].set_ylabel(\"Position (m)\")\n",
    "axarr[2].set_title(\"Second body wrt first body frame\")\n",
    "axarr[2].grid(\"on\")\n",
    "axarr[2].legend([\"x\", \"y\", \"z\"], loc=\"upper left\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "awful-village",
   "metadata": {},
   "source": [
    "As you can see, the resulting timestamps are the same as those of the second body; however, the first sample has been dropped because the transformation is not defined there.\n"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}