{
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {
            "display_name": "Python",
            "language": "python",
            "name": "python"
        },
        "language_info": {
            "name": "python"
        }
    },
    "cells": [
        {
            "id": "7abee766",
            "cell_type": "markdown",
            "source": "<div class='alert alert-warning'>\n\nSciPy's interactive examples with Jupyterlite are experimental and may not always work as expected. Execution of cells containing imports may result in large downloads (up to 60MB of content for the first import from SciPy). Load times when importing from SciPy may take roughly 10-20 seconds. If you notice any problems, feel free to open an [issue](https://github.com/scipy/scipy/issues/new/choose).\n\n</div>",
            "metadata": {}
        },
        {
            "id": "870b5984",
            "cell_type": "markdown",
            "source": "A `RigidTransform` instance can be initialized in any of the above formats\nand converted to any of the others. The underlying object is independent of\nthe representation used for initialization.\n\n**Notation Conventions and Composition**\n\nThe notation here largely follows the convention defined in [5]. When we\nname transforms, we read the subscripts from right to left. So ``tf_A_B``\nrepresents a transform A <- B and can be interpreted as:\n\n- the coordinates and orientation of B relative to A\n- the transformation of points from B to A\n- the pose of B described in A's coordinate system\n\n.. parsed-literal\n```\n:class: highlight-none\n\ntf_A_B\n^ ^\n| |\n| --- from B\n|\n----- to A\n\n```\nWhen composing transforms, the order is important. Transforms are not\ncommutative, so in general ``tf_A_B * tf_B_C`` is not the same as\n``tf_B_C * tf_A_B``. Transforms are composed and applied to vectors\nright-to-left. So ``(tf_A_B * tf_B_C).apply(p_C)`` is the same as\n``tf_A_B.apply(tf_B_C.apply(p_C))``.\n\nWhen composed, transforms should be ordered such that the multiplication\noperator is surrounded by a single frame, so the frame \"cancels out\" and\nthe outside frames are left. In the example below, B cancels out and the\noutside frames A and C are left. Or to put it another way, A <- C is the\nsame as A <- B <- C.\n\n.. parsed-literal\n```\n:class: highlight-none\n\n----------- B cancels out\n|      |\nv      v\ntf_A_C = tf_A_B * tf_B_C\n^          ^\n|          |\n------------ to A, from C are left\n\n```\nWhen we notate vectors, we write the subscript of the frame that the vector\nis defined in. So ``p_B`` means the point ``p`` defined in frame B. To\ntransform this point from frame B to coordinates in frame A, we apply the\ntransform ``tf_A_B`` to the vector, lining things up such that the notated\nB frames are next to each other and \"cancel out\".\n\n.. parsed-literal\n```\n:class: highlight-none\n\n------------ B cancels out\n|         |\nv         v\np_A = tf_A_B.apply(p_B)\n^\n|\n-------------- A is left\n\n```\n**Visualization**\n",
            "metadata": {}
        },
        {
            "id": "b5db416e",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "from scipy.spatial.transform import RigidTransform as Tf\nfrom scipy.spatial.transform import Rotation as R\nimport numpy as np",
            "outputs": []
        },
        {
            "id": "4339481b",
            "cell_type": "markdown",
            "source": "The following function can be used to plot transforms with Matplotlib\nby showing how they transform the standard x, y, z coordinate axes:\n",
            "metadata": {}
        },
        {
            "id": "692b60e0",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "import matplotlib.pyplot as plt\ncolors = (\"#FF6666\", \"#005533\", \"#1199EE\")  # Colorblind-safe RGB\ndef plot_transformed_axes(ax, tf, name=None, scale=1):\n    r = tf.rotation\n    t = tf.translation\n    loc = np.array([t, t])\n    for i, (axis, c) in enumerate(zip((ax.xaxis, ax.yaxis, ax.zaxis),\n                                      colors)):\n        axlabel = axis.axis_name\n        axis.set_label_text(axlabel)\n        axis.label.set_color(c)\n        axis.line.set_color(c)\n        axis.set_tick_params(colors=c)\n        line = np.zeros((2, 3))\n        line[1, i] = scale\n        line_rot = r.apply(line)\n        line_plot = line_rot + loc\n        ax.plot(line_plot[:, 0], line_plot[:, 1], line_plot[:, 2], c)\n        text_loc = line[1]*1.2\n        text_loc_rot = r.apply(text_loc)\n        text_plot = text_loc_rot + t\n        ax.text(*text_plot, axlabel.upper(), color=c,\n                va=\"center\", ha=\"center\")\n    ax.text(*tf.translation, name, color=\"k\", va=\"center\", ha=\"center\",\n            bbox={\"fc\": \"w\", \"alpha\": 0.8, \"boxstyle\": \"circle\"})",
            "outputs": []
        },
        {
            "id": "f5e2aa8f",
            "cell_type": "markdown",
            "source": "**Defining Frames**\n\nLet's work through an example.\n\nFirst, define the \"world frame\" A, also called the \"base frame\".\nAll frames are the identity transform from their own perspective.\n",
            "metadata": {}
        },
        {
            "id": "c5105247",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "tf_A = Tf.identity()",
            "outputs": []
        },
        {
            "id": "24b58df9",
            "cell_type": "markdown",
            "source": "We will visualize a new frame B in A's coordinate system. So we need to\ndefine the transform that converts coordinates from frame B to frame A\n(A <- B).\n\nPhysically, let's imagine constructing B from A by:\n\n1) Rotating A by +90 degrees around its x-axis.\n2) Translating the rotated frame 2 units in A's -x direction.\n\nFrom A's perspective, B is at [-2, 0, 0] and rotated +90 degrees about the\nx-axis, which is exactly the transform A <- B.\n",
            "metadata": {}
        },
        {
            "id": "29bab9c4",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "t_A_B = np.array([-2, 0, 0])\nr_A_B = R.from_euler('xyz', [90, 0, 0], degrees=True)\ntf_A_B = Tf.from_components(t_A_B, r_A_B)",
            "outputs": []
        },
        {
            "id": "964f0153",
            "cell_type": "markdown",
            "source": "Let's plot these frames.\n",
            "metadata": {}
        },
        {
            "id": "a47ec354",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "fig, ax = plt.subplots(subplot_kw={\"projection\": \"3d\"})\nplot_transformed_axes(ax, tf_A, name=\"tfA\")     # A plotted in A\nplot_transformed_axes(ax, tf_A_B, name=\"tfAB\")  # B plotted in A\nax.set_title(\"A, B frames with respect to A\")\nax.set_aspect(\"equal\")\nax.figure.set_size_inches(6, 5)\nplt.show()",
            "outputs": []
        },
        {
            "id": "40a08f89",
            "cell_type": "markdown",
            "source": "Now let's visualize a new frame C in B's coordinate system.\nLet's imagine constructing C from B by:\n\n1) Translating B by 2 units in its +z direction.\n2) Rotating B by +30 degrees around its z-axis.\n",
            "metadata": {}
        },
        {
            "id": "c0527da6",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "t_B_C = np.array([0, 0, 2])\nr_B_C = R.from_euler('xyz', [0, 0, 30], degrees=True)\ntf_B_C = Tf.from_components(t_B_C, r_B_C)",
            "outputs": []
        },
        {
            "id": "8b3f23a8",
            "cell_type": "markdown",
            "source": "In order to plot these frames from a consistent perspective, we need to\ncalculate the transform between A and C. Note that we do not make this\ntransform directly, but instead compose intermediate transforms that let us\nget from C to A:\n",
            "metadata": {}
        },
        {
            "id": "ac4b9ad2",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "tf_A_C = tf_A_B * tf_B_C  # A <- B <- C",
            "outputs": []
        },
        {
            "id": "6db8ba83",
            "cell_type": "markdown",
            "source": "Now we can plot these three frames from A's perspective.\n",
            "metadata": {}
        },
        {
            "id": "469f45bb",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "fig, ax = plt.subplots(subplot_kw={\"projection\": \"3d\"})\nplot_transformed_axes(ax, tf_A, name=\"tfA\")     # A plotted in A\nplot_transformed_axes(ax, tf_A_B, name=\"tfAB\")  # B plotted in A\nplot_transformed_axes(ax, tf_A_C, name=\"tfAC\")  # C plotted in A\nax.set_title(\"A, B, C frames with respect to A\")\nax.set_aspect(\"equal\")\nax.figure.set_size_inches(6, 5)\nplt.show()",
            "outputs": []
        },
        {
            "id": "d12a43e7",
            "cell_type": "markdown",
            "source": "**Transforming Vectors**\n\nLet's transform a vector from A, to B and C. To do this, we will first\ninvert the transforms we already have from B and C, to A.\n",
            "metadata": {}
        },
        {
            "id": "c0a42650",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "tf_B_A = tf_A_B.inv()  # B <- A\ntf_C_A = tf_A_C.inv()  # C <- A",
            "outputs": []
        },
        {
            "id": "fe832c0f",
            "cell_type": "markdown",
            "source": "Now we can define a point in A and use the above transforms to get its\ncoordinates in B and C:\n",
            "metadata": {}
        },
        {
            "id": "d129c2b9",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "p1_A = np.array([1, 0, 0])  # +1 in x_A direction\np1_B = tf_B_A.apply(p1_A)\np1_C = tf_C_A.apply(p1_A)\nprint(p1_A)  # Original point 1 in A",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "[1 0 0]"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "5cabf462",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "print(p1_B)  # Point 1 in B",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "[3. 0. 0.]"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "dc380077",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "print(p1_C)  # Point 1 in C",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "[ 2.59807621 -1.5       -2.        ]"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "e49c9e20",
            "cell_type": "markdown",
            "source": "We can also do the reverse. We define a point in C and transform it to A:\n",
            "metadata": {}
        },
        {
            "id": "bd01c1a0",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "p2_C = np.array([0, 1, 0])  # +1 in y_C direction\np2_A = tf_A_C.apply(p2_C)\nprint(p2_C)  # Original point 2 in C",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "[0 1 0]"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "0a28c2c3",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "print(p2_A)  # Point 2 in A",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "[-2.5       -2.         0.8660254]"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "4c38451e",
            "cell_type": "markdown",
            "source": "Plot the frames with respect to A again, but also plot these two points:\n",
            "metadata": {}
        },
        {
            "id": "b105051e",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "fig, ax = plt.subplots(subplot_kw={\"projection\": \"3d\"})\nplot_transformed_axes(ax, tf_A, name=\"tfA\")     # A plotted in A\nplot_transformed_axes(ax, tf_A_B, name=\"tfAB\")  # B plotted in A\nplot_transformed_axes(ax, tf_A_C, name=\"tfAC\")  # C plotted in A\nax.scatter(p1_A[0], p1_A[1], p1_A[2], color=colors[0])  # +1 x_A\nax.scatter(p2_A[0], p2_A[1], p2_A[2], color=colors[1])  # +1 y_C\nax.set_title(\"A, B, C frames and points with respect to A\")\nax.set_aspect(\"equal\")\nax.figure.set_size_inches(6, 5)\nplt.show()",
            "outputs": []
        },
        {
            "id": "454758fe",
            "cell_type": "markdown",
            "source": "**Switching Base Frames**\n\nUp to this point, we have been visualizing frames from A's perspective.\nLet's use the transforms we defined to visualize the frames from C's\nperspective.\n\nNow C is the \"base frame\" or \"world frame\". All frames are the identity\ntransform from their own perspective.\n",
            "metadata": {}
        },
        {
            "id": "395d4fbd",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "tf_C = Tf.identity()",
            "outputs": []
        },
        {
            "id": "2c847d4c",
            "cell_type": "markdown",
            "source": "We've already defined the transform C <- A, and can obtain C <- B by\ninverting the existing transform B <- C.\n",
            "metadata": {}
        },
        {
            "id": "a8e7d6d6",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "tf_C_B = tf_B_C.inv()  # C <- B",
            "outputs": []
        },
        {
            "id": "168eeab4",
            "cell_type": "markdown",
            "source": "This lets us plot everything from C's perspective:\n",
            "metadata": {}
        },
        {
            "id": "3486fe72",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "fig, ax = plt.subplots(subplot_kw={\"projection\": \"3d\"})\nplot_transformed_axes(ax, tf_C, name=\"tfC\")     # C plotted in C\nplot_transformed_axes(ax, tf_C_B, name=\"tfCB\")  # B plotted in C\nplot_transformed_axes(ax, tf_C_A, name=\"tfCA\")  # A plotted in C\nax.scatter(p1_C[0], p1_C[1], p1_C[2], color=colors[0])\nax.scatter(p2_C[0], p2_C[1], p2_C[2], color=colors[1])\nax.set_title(\"A, B, C frames and points with respect to C\")\nax.set_aspect(\"equal\")\nax.figure.set_size_inches(6, 5)\nplt.show()",
            "outputs": []
        }
    ]
}