{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "diverse-gallery",
   "metadata": {},
   "source": [
    "# Estimating transforms from data\n",
    "\n",
    "It is often necessary to estimate transformations between rigid bodies that are not explicitly known. This happens for example when the motion of the same rigid body is measured by different tracking systems that represent their data in different world frames."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "international-ceramic",
   "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": "demonstrated-charleston",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rigid_body_motion as rbm\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (10, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "placed-holly",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.register_frame(\"world\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "satisfactory-politics",
   "metadata": {},
   "source": [
    "## Shortest arc rotation\n",
    "\n",
    "Let's assume we have two vectors $v_1$ and $v_2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brief-representative",
   "metadata": {},
   "outputs": [],
   "source": [
    "v1 = (1, 0, 0)\n",
    "v2 = (np.sqrt(2) / 2, np.sqrt(2) / 2, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "disabled-petroleum",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(\"world\", ax=ax)\n",
    "rbm.plot.vectors(v1, ax=ax, color=\"y\")\n",
    "rbm.plot.vectors(v2, ax=ax, color=\"c\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "identical-montreal",
   "metadata": {},
   "source": [
    "The quaternion $r$ that rotates $v_1$ in the same direction as $v_2$, i.e., that satisfies:\n",
    "\n",
    "$$\\frac{v_2}{\\|v_2\\|} = \\frac{\\operatorname{rot}(r, v_1)}{\\|v_1\\|}$$\n",
    "\n",
    "can be computed with the [shortest_arc_rotation()](_generated/rigid_body_motion.shortest_arc_rotation.rst) method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "romance-phrase",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.shortest_arc_rotation(v1, v2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "arctic-rental",
   "metadata": {},
   "source": [
    "The method also works with arrays of vectors. Let's first construct an array of progressive rotations around the yaw axis with the [from_euler_angles()](_generated/rigid_body_motion.from_euler_angles.rst) method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "otherwise-russell",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = rbm.from_euler_angles(yaw=np.linspace(0, np.pi / 8, 10))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "split-sodium",
   "metadata": {},
   "source": [
    "Now we can rotate $v_2$ with $r$. Because we rotate a single vector with multiple quaternions we have to specify `one_to_one=False`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "certain-armstrong",
   "metadata": {},
   "outputs": [],
   "source": [
    "v2_arr = rbm.rotate_vectors(r, v2, one_to_one=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sunset-accused",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(\"world\", ax=ax)\n",
    "rbm.plot.vectors(v1, ax=ax, color=\"y\")\n",
    "rbm.plot.vectors(v2_arr, ax=ax, color=\"c\", alpha=0.3)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "opponent-transition",
   "metadata": {},
   "source": [
    "[shortest_arc_rotation()](_generated/rigid_body_motion.shortest_arc_rotation.rst) now returns an array of quaternions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "placed-circumstances",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.shortest_arc_rotation(v1, v2_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "suspended-irish",
   "metadata": {},
   "source": [
    "## Best fit rotation\n",
    "\n",
    "In a different scenario, we might have two vectors that are offset by a fixed rotation and are rotating in space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "italic-comparison",
   "metadata": {},
   "outputs": [],
   "source": [
    "v1_arr = rbm.rotate_vectors(r, v1, one_to_one=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "express-investment",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(\"world\", ax=ax)\n",
    "rbm.plot.vectors(v1_arr, ax=ax, color=\"y\", alpha=0.3)\n",
    "rbm.plot.vectors(v2_arr, ax=ax, color=\"c\", alpha=0.3)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "metallic-offset",
   "metadata": {},
   "source": [
    "The rotation between the vectors can be found with a least-squares minimization:\n",
    "\n",
    "$$\\min_r \\left\\| v_2 - \\operatorname{rot}(r, v_1) \\right\\|$$\n",
    "\n",
    "This is implemented in the [best_fit_rotation()](_generated/rigid_body_motion.best_fit_rotation.rst) method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "italian-thousand",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.best_fit_rotation(v1_arr, v2_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pressing-indianapolis",
   "metadata": {},
   "source": [
    "## Best fit transform\n",
    "\n",
    "In yet another case, we might have two arrays of points (e.g. point clouds) with a fixed transform (rotation *and* translation) between them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "interior-values",
   "metadata": {},
   "outputs": [],
   "source": [
    "p1_arr = 0.1 * np.random.randn(100, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pursuant-excuse",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.array((1, 1, 0))\n",
    "r = rbm.from_euler_angles(yaw=np.pi / 4)\n",
    "p2_arr = rbm.rotate_vectors(r, p1_arr, one_to_one=False) + t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "willing-sixth",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "rbm.plot.reference_frame(\"world\", ax=ax)\n",
    "rbm.plot.points(p1_arr, ax=ax, fmt=\"yo\")\n",
    "rbm.plot.points(p2_arr, ax=ax, fmt=\"co\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "confused-color",
   "metadata": {},
   "source": [
    "To estimate this transform, we can minimize:\n",
    "\n",
    "$$\\min_r \\left\\| p_2 - (\\operatorname{rot}(r, p_1) + t) \\right\\|$$\n",
    "\n",
    "This algorithm (also called point set registration) is implemented in the [best_fit_transform()](_generated/rigid_body_motion.best_fit_transform.rst) method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "announced-warning",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.best_fit_transform(p1_arr, p2_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "understanding-dayton",
   "metadata": {},
   "source": [
    "## Iterative closest point\n",
    "\n",
    "The above algorithm only works for known correspondences between points $p_1$ and $p_2$ (i.e., each point in `p1_arr` corresponds to the same index in `p2_arr`). This is not always the case - in fact, something like a point cloud from different laser scans of the same object might yield sets of completely different points. An approximate transform can still be found with the iterative closest point (ICP) algorithm. We can simulate the case of unknown correspondences by randomly permuting the second array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "atomic-architect",
   "metadata": {},
   "outputs": [],
   "source": [
    "rbm.iterative_closest_point(p1_arr, np.random.permutation(p2_arr))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "royal-removal",
   "metadata": {},
   "source": [
    "Note that there is a discrepancy in the estimated transform compared to the best fit transform. ICP usually yields better results with a larger number of points that have more spatial structure."
   ]
  }
 ],
 "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
}