{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geometry Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an introductory walkthrough of the symbolic [geometry package](../api/symforce.geo.html) in symforce. Symforce natively contains the following geometry objects: complex numbers, quaternions, dual quaternions, matrices, 2D and 3D rotations, and 2D and 3D poses (translation + rotation).\n",
    "\n",
    "In this tutorial we will demonstrate a few of the ways these geometry objects can be constructed and composed. Here we pay special attention to 3D rotations, poses, and vectors, as they are typically the most commonly used geometric objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup\n",
    "import symforce\n",
    "\n",
    "symforce.set_symbolic_api(\"sympy\")\n",
    "symforce.set_log_level(\"warning\")\n",
    "\n",
    "import symforce.symbolic as sf\n",
    "from symforce import ops\n",
    "from symforce.notebook_util import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rotations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rotations can be defined from and converted to a number of different representations as shown below, but always use a quaternion as the underlying representation. We use the notation world_R_body to represent a rotation that rotates a point written in the body frame into the world frame.\n",
    "\n",
    "Note that in some cases a small epsilon can be provided to prevent numerical errors (e.g. division by zero) when converting to/from rotation representations. Furthermore, converting between certain representations can require additional symbolic expressions to guard against degenerate cases. For example, a rotation constructed from a rotation matrix results in more complexity than when constructed using an axis-angle representation as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Identity definition\n",
    "display(sf.Rot3())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Symbolic definition\n",
    "display(sf.Rot3.symbolic(\"R\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From a quaternion of the form qw + qv0 + qv1 + qv2\n",
    "qv = sf.V3.symbolic(\"qv\")\n",
    "qw = sf.Symbol(\"qw\")\n",
    "quat = sf.Quaternion(xyz=qv, w=qw)\n",
    "display(quat)\n",
    "# we can use this quaternion to initialize a sf.Rot3:\n",
    "R_from_quat = sf.Rot3(quat)\n",
    "display(R_from_quat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To/From rotation matrix\n",
    "\n",
    "# Rotate about x-axis\n",
    "theta = sf.Symbol(\"theta\")\n",
    "R_mat = sf.Matrix(\n",
    "    [\n",
    "        [1, 0, 0],\n",
    "        [0, sf.cos(theta), -sf.sin(theta)],\n",
    "        [0, sf.sin(theta), sf.cos(theta)],\n",
    "    ]\n",
    ")\n",
    "R = sf.Rot3.from_rotation_matrix(R_mat)\n",
    "\n",
    "display(R_mat)\n",
    "display(R)  # Note the additional expressions required to avoid numerical errors\n",
    "display(R.to_rotation_matrix())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To/From Euler angles\n",
    "R = sf.Rot3.from_yaw_pitch_roll(0, 0, theta)  # Yaw rotation only\n",
    "ypr = R.to_yaw_pitch_roll()\n",
    "\n",
    "display(R)\n",
    "display(ops.StorageOps.simplify(list(ypr)))  # Simplify YPR expression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From axis-angle representation\n",
    "\n",
    "# Rotate about x-axis\n",
    "R = sf.Rot3.from_angle_axis(angle=theta, axis=sf.Vector3(1, 0, 0))\n",
    "\n",
    "display(R)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we can construct rotations, we can use them to rotate vectors as one would expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rotation defining orientation of body frame wrt world frame\n",
    "world_R_body = sf.Rot3.symbolic(\"R\")\n",
    "\n",
    "# Point written in body frame\n",
    "body_t_point = sf.Vector3.symbolic(\"p\")\n",
    "\n",
    "# Point written in world frame\n",
    "world_t_point = world_R_body * body_t_point\n",
    "\n",
    "display(world_t_point)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Chaining rotations and inverting rotations works as one would expect as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "body_R_cam = sf.Rot3.symbolic(\"R_cam\")\n",
    "world_R_cam = world_R_body * body_R_cam\n",
    "\n",
    "# Rotation inverse = negate vector part of quaternion\n",
    "cam_R_body = body_R_cam.inverse()\n",
    "display(body_R_cam)\n",
    "display(cam_R_body)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also easily substitute numerical values into symbolic expressions using geo objects themselves. This makes it very convenient to substitute numeric values into large symbolic expressions constructed using many different geo objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "world_R_body_numeric = sf.Rot3.from_yaw_pitch_roll(0.1, -2.3, 0.7)\n",
    "display(world_t_point.subs(world_R_body, world_R_body_numeric))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Poses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Poses are defined as a rotation plus a translation, and are constructed as such. We use the notation world_T_body to represent a pose that transforms from the body frame to the world frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Symbolic construction\n",
    "world_T_body = sf.Pose3.symbolic(\"T\")\n",
    "display(world_T_body)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construction from a rotation and translation\n",
    "\n",
    "# Orientation of body frame wrt world frame\n",
    "world_R_body = sf.Rot3.symbolic(\"R\")\n",
    "\n",
    "# Position of body frame wrt world frame written in the world frame\n",
    "world_t_body = sf.Vector3.symbolic(\"t\")\n",
    "\n",
    "world_T_body = sf.Pose3(R=world_R_body, t=world_t_body)\n",
    "display(world_T_body)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to rotations, we can compose poses with poses, compose poses with 3D points, and invert poses as one would expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compose pose with a pose\n",
    "body_T_cam = sf.Pose3.symbolic(\"T_cam\")\n",
    "world_T_cam = world_T_body * body_T_cam\n",
    "\n",
    "# Compose pose with a point\n",
    "body_t_point = sf.Vector3.symbolic(\"p\")  # Position in body frame\n",
    "# Equivalent to: world_R_body * body_t_point + world_t_body\n",
    "world_t_point = world_T_body * body_t_point\n",
    "display(world_t_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Invert a pose\n",
    "body_T_world = world_T_body.inverse()\n",
    "display(world_T_body)\n",
    "display(body_T_world)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vectors and Matrices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vectors and matrices are all represented using subclasses of sf.Matrix class, and can be constructed in several different ways as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Matrix construction. The statements below all create the same 2x3 matrix object\n",
    "\n",
    "# Construction from 2D list\n",
    "m1 = sf.Matrix([[1, 2, 3], [4, 5, 6]])\n",
    "\n",
    "# Construction using specified size + data\n",
    "m2 = sf.Matrix(2, 3, [1, 2, 3, 4, 5, 6])\n",
    "\n",
    "# sf.MatrixNM creates a matrix with shape NxM (defined by default for 6x6\n",
    "# matrices and smaller)\n",
    "m3 = sf.Matrix23(1, 2, 3, 4, 5, 6)\n",
    "m4 = sf.Matrix23([1, 2, 3, 4, 5, 6])\n",
    "\n",
    "# Construction using aliases\n",
    "m5 = sf.M([[1, 2, 3], [4, 5, 6]])\n",
    "m6 = sf.M(2, 3, [1, 2, 3, 4, 5, 6])\n",
    "m7 = sf.M23(1, 2, 3, 4, 5, 6)\n",
    "m8 = sf.M23([1, 2, 3, 4, 5, 6])\n",
    "\n",
    "# Construction from block matrices of appropriate dimensions\n",
    "m9 = sf.Matrix23.block_matrix([[sf.M13([1, 2, 3])], [sf.M13([3, 4, 5])]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vector constructors. The statements below all create the same 3x1 vector object\n",
    "\n",
    "# Construction from 2D list\n",
    "v1 = sf.Matrix([[1], [2], [3]])\n",
    "\n",
    "# Construction from 1D list. We assume a 1D list represents a column vector.\n",
    "v2 = sf.Matrix([1, 2, 3])\n",
    "\n",
    "# Construction using aliases (defined by default for 9x1 vectors and smaller)\n",
    "v3 = sf.Matrix31(1, 2, 3)\n",
    "v4 = sf.M31(1, 2, 3)\n",
    "v5 = sf.Vector3(1, 2, 3)\n",
    "v6 = sf.V3(1, 2, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use a few typical matrix constructors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Matrix of zeros\n",
    "z1 = sf.Matrix23.zero()\n",
    "z2 = sf.Matrix.zeros(2, 3)\n",
    "\n",
    "# Matrix of ones\n",
    "o1 = sf.Matrix23.one()\n",
    "o2 = sf.Matrix.ones(2, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the Matrix class itself does not contain group or lie group methods, to prevent confusion\n",
    "between the identity matrix and inverse matrix, and the group operations under addition.  The group\n",
    "operations are implemented separately for matrices under addition, and are accessed through\n",
    "ops.GroupOps and ops.LieGroupOps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zero_matrix = sf.Matrix33.zero()\n",
    "identity_matrix = sf.Matrix33.eye()\n",
    "\n",
    "# We could also write:\n",
    "zero_matrix = ops.GroupOps.identity(sf.Matrix33)\n",
    "\n",
    "display(zero_matrix)\n",
    "display(identity_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, of course, matrix math works as one would expect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Matrix multiplication\n",
    "m23 = sf.M23.symbolic(\"lhs\")\n",
    "m31 = sf.V3.symbolic(\"rhs\")\n",
    "display(m23 * m31)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vector operations\n",
    "norm = m31.norm()\n",
    "squared_norm = m31.squared_norm()\n",
    "unit_vec = m31.normalized()\n",
    "display(unit_vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m33 = 5 * sf.Matrix33.eye()  # Element-wise multiplication with scalar\n",
    "display(m33.inv())  # Matrix inverse"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the most powerful operations we can use matrices for is to compute jacobians with respect to other geo objects. By default we compute jacobians with respect to the tangent space of the given object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "R0 = sf.Rot3.symbolic(\"R0\")\n",
    "R1 = sf.Rot3.symbolic(\"R1\")\n",
    "residual = sf.M(R0.local_coordinates(R1))\n",
    "display(residual)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jacobian = residual.jacobian(R1)\n",
    "# The jacobian is quite a complex symbolic expression, so we don't display it for\n",
    "# convenience.\n",
    "# The shape is equal to (dimension of residual) x (dimension of tangent space)\n",
    "display(jacobian.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General properties of geo objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Storage operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All geometric types implement the \"Storage\" interface. This means that they can:\n",
    "\n",
    "1. Be serialized into a list of scalar expressions (`.to_storage()`)\n",
    "2. Be reconstructed from a list of scalar expressions (`.from_storage()`)\n",
    "3. Use common symbolic operations (symbolic construction, substitution, simplification, etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Serialization to scalar list\n",
    "rot = sf.Rot3()\n",
    "elements = rot.to_storage()\n",
    "assert len(elements) == rot.storage_dim()\n",
    "display(elements)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construction from scalar list\n",
    "rot2 = sf.Rot3.from_storage(elements)\n",
    "assert rot == rot2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Symbolic operations\n",
    "rot_sym = sf.Rot3.symbolic(\"rot_sym\")\n",
    "rot_num = rot_sym.subs(rot_sym, rot)\n",
    "\n",
    "display(rot_sym)\n",
    "display(rot_num)\n",
    "display(rot_num.simplify())  # Simplify internal symbolic expressions\n",
    "display(rot_num.evalf())  # Numerical evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Group operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All geometric types also implement the \"Group\" interface, meaning that geometric objects:\n",
    "\n",
    "1. Can be composed with objects of the same type to produce an object of the same type (`.compose()`)\n",
    "2. Have an identity element (`.identity()`)\n",
    "3. Can be inverted (`.inverse()`)\n",
    "4. Can be created to represent the relation between two other objects of the same type (`.between()`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct two random rotations\n",
    "R1 = sf.Rot3.random()\n",
    "R2 = sf.Rot3.random()\n",
    "\n",
    "# Composition\n",
    "display(R1.compose(R2))  # For rotations this is the same as R1 * R2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Identity\n",
    "R_identity = sf.Rot3.identity()\n",
    "display(R1)\n",
    "display(R_identity * R1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inverse\n",
    "R1_inv = R1.inverse()\n",
    "display(R_identity)\n",
    "display(R1_inv * R1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Between\n",
    "R_delta = R1.between(R2)\n",
    "display(R1 * R_delta)\n",
    "display(R2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lie Group operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rotations, poses, and matrices all implement the \"LieGroup\" interface, meaning that they each have a\n",
    "tangent space. There are many great references on Lie groups out there already, so instead of\n",
    "introducing them here, we recommend checking out\n",
    "[Frank Dellaert's](https://github.com/borglab/gtsam/blob/develop/doc/LieGroups.pdf),\n",
    "[Ethan Eade's](https://www.ethaneade.com/lie.pdf), or\n",
    "[JL Blanco's](https://ingmec.ual.es/~jlblanco/papers/jlblanco2010geometry3D_techrep.pdf) tutorials.\n",
    "In SymForce, objects which are a Lie Group can:\n",
    "\n",
    "1. Be used to compute the tangent space vector about the identity element (`.to_tangent()`)\n",
    "2. Be constructed from a tangent space vector about the identity element (`.from_tangent()`)\n",
    "3. Be perturbed by a tangent space vector about the given element (`.retract()`)\n",
    "4. Be used to compute the tangent space perturbation needed to obtain another given element (`.local_coordinates()`)\n",
    "5. Be used to compute a jacobian describing the relation between the underlying data of the object (e.g. a quaternion for a rotation) and the tangent space vector about the given element (`.storage_D_tangent()`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To/From tangent space vector about identity element\n",
    "R1 = sf.Rot3.random()\n",
    "tangent_vec = R1.to_tangent()\n",
    "R1_recovered = sf.Rot3.from_tangent(tangent_vec)\n",
    "\n",
    "assert len(tangent_vec) == R1.tangent_dim()\n",
    "display(R1)\n",
    "display(R1_recovered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tangent space perturbations\n",
    "\n",
    "# Perturb R1 by the given vector in the tangent space around R1\n",
    "R2 = R1.retract([0.1, 2.3, -0.5])\n",
    "\n",
    "# Compute the tangent vector pointing from R1 to R2, in the tangent space\n",
    "# around R1\n",
    "recovered_tangent_vec = R1.local_coordinates(R2)\n",
    "\n",
    "display(recovered_tangent_vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Jacobian of storage w.r.t tangent space perturbation\n",
    "\n",
    "# We chain storage_D_tangent together with jacobians of larger symbolic\n",
    "# expressions taken with respect to the symbolic elements of the object (e.g. a\n",
    "# quaternion for rotations) to compute the jacobian wrt the tanget space about\n",
    "# the element.\n",
    "# I.e. residual_D_tangent = residual_D_storage * storage_D_tangent\n",
    "\n",
    "jacobian = R1.storage_D_tangent()\n",
    "assert jacobian.shape == (R1.storage_dim(), R1.tangent_dim())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more details on Storage/Group/LieGroup operations, see the [Ops tutorial](../tutorials/ops_tutorial.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Poses and SE(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Poses in robotics are often represented using the [Special Euclidean group](https://en.wikipedia.org/wiki/Euclidean_group),\n",
    "SE(2) and SE(3) for 2D and 3D respectively. Additionally, the manifold is often defined so that the\n",
    "`from_tangent` function is the matrix exponential of a linear combination of generators.\n",
    "\n",
    "This has the advantage of simpler theoretical treatment in some cases, but has disadvantages in\n",
    "terms of the number of operations required to implement commonly used functions like `retract`.\n",
    "For these reasons, instead of using the SE(n) `from_tangent` and `to_tangent` functions, ours\n",
    "decouple the perturbations to rotation and translation.  This means our `Pose2` and `Pose3` objects\n",
    "represent neither `SE(n)` nor `SO(n) x R^n`; they compose like objects on `SE(n)` but have tangent\n",
    "space operations analogous to `SO(n) x R^n`.\n",
    "\n",
    "For most users, this should all happen under the hood.  The most common observable difference is\n",
    "that `retract(g, v) != compose(g, from_tangent(v))`, so simplifications done by hand that rely on\n",
    "this fact must be treated with care.  We do also provide a `Pose3_SE3` class for users who need this\n",
    "functionality in symbolic expressions - this class has no runtime equivalent and is unsupported by\n",
    "much of SymForce, including the Optimizer."
   ]
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.10.6"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython2",
  "version": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}