{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ffb6b6c8-e58d-4df5-a68a-cb9f191b559c",
   "metadata": {},
   "source": [
    "# Basic math usage\n",
    "\n",
    "This notebook demonstrates how to work with basic math objects using **kdFlex**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "237fca86",
   "metadata": {},
   "source": [
    "<img src=../resources/nb_images/math.png align=center width=600>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5a55f5f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Karana.Math as km\n",
    "import numpy as np\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07d864e1",
   "metadata": {},
   "source": [
    "## Vectors and Matrices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3248bae",
   "metadata": {},
   "source": [
    "**kdFlex** natively supports [`numpy`](https://numpy.org/) arrays at the Python level for vectors and matrices. The Python bindings seamless convert the `numpy` arrays to the C++ `Vec` and `Mat` class. Here's an example of creating a simple three-vector using `numpy`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9412bff8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 2., 3.])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vec3 = np.array([1.0, 2.0, 3.0])\n",
    "vec3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab2746eb",
   "metadata": {},
   "source": [
    "Similarly, here is the $3 \\times 3$ identity matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "18991b6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 0., 0.],\n",
       "       [0., 1., 0.],\n",
       "       [0., 0., 1.]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eye33 = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])\n",
    "eye33"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c611afb5",
   "metadata": {},
   "source": [
    "`numpy` has [rich API](https://numpy.org/learn/), which you are encouraged to explore. For example, the above could also be achieved with: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e11de43d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 0., 0.],\n",
       "       [0., 1., 0.],\n",
       "       [0., 0., 1.]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eye33 = np.eye(3)\n",
    "eye33"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1029bd08",
   "metadata": {},
   "source": [
    "## Spatial vectors\n",
    "\n",
    "**kdFlex** uses *spatial notation* which involves working with 6-dimension _spatial vector_ quantities. For example, spatial velocity is a spatial vector consisting of both the $3$-dimensional angular and linear velocities. A `SpatialVector` is a generic container for these $3 + 3$ dimensional vectors, where $w$ is the angular part and $v$ is the linear part:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aee17dd5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.SpatialVector(\n",
       "\tarray([1., 2., 3.]),\n",
       "\tarray([4., 5., 6.])\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vel = km.SpatialVector(w=[1, 2, 3], v=[4, 5, 6])\n",
    "vel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad142dca",
   "metadata": {},
   "source": [
    "For commonly used quantities, there are specialized, unit-aware `SpatialVector`s. Here we create a `SpatialVelocity` with nonstandard units:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6680ab61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.SpatialVelocity(\n",
       "\t<Quantity([0.01745329 0.03490659 0.05235988], '1 / second')>,\n",
       "\t<Quantity([4000. 5000. 6000.], 'meter / second')>\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from Karana.Math.Kquantities import ureg\n",
    "\n",
    "vel = km.SpatialVelocity([1, 2, 3] * ureg.deg / ureg.s, [4, 5, 6] * ureg.km / ureg.s)\n",
    "vel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c53594d0",
   "metadata": {},
   "source": [
    "## Rotations\n",
    "\n",
    "**kdFlex** primarily represents rotations using unit quaternions, consisting of four components. Unlike three-vector representations such as Euler angles, unit quaternions are free of singularities (often seen as gimbal lock), making them more suitable for general rotations.\n",
    "\n",
    "There are a number of ways to construct a `UnitQuaternion`. For example, this rotates by $\\pi /2 $ about the $y$-axis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "38998103",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.00000000e+00, 0.00000000e+00, 2.22044605e-16])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_vec = np.array([0.0, 0.0, 1.0])\n",
    "y90 = km.UnitQuaternion(math.pi / 2, np.array([0.0, 1.0, 0.0]))\n",
    "y90 * z_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3f76a83",
   "metadata": {},
   "source": [
    "For convenience and brevity, you can typically use native python lists in place of `numpy` arrays in most cases, with the vectors automatically being converted behind the scenes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "695a9f1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.00000000e+00, 0.00000000e+00, 2.22044605e-16])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_vec = [0.0, 0.0, 1.0]\n",
    "y90 = km.UnitQuaternion(math.pi / 2, [0.0, 1.0, 0.0])\n",
    "y90 * z_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae1e6ab9",
   "metadata": {},
   "source": [
    "Here's an example of converting $ZXZ$ Euler angles to a `UnitQuaternion`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "edf1c685",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.UnitQuaternion(1.1102230246251565e-16, 0.7071067811865475, 0, 0.7071067811865476)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "angles = km.EulerAngles(math.pi / 2, math.pi / 2, -math.pi / 2, km.EulerSystem.ZXZ)\n",
    "quat = km.UnitQuaternion(angles)\n",
    "quat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bdc2f64",
   "metadata": {},
   "source": [
    "Converting a `UnitQuaternion` to a direction cosine matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7fa501ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.RotationMatrix(\n",
       "array([[ 2.22044605e-16,  0.00000000e+00,  1.00000000e+00],\n",
       "       [ 0.00000000e+00,  1.00000000e+00,  0.00000000e+00],\n",
       "       [-1.00000000e+00,  0.00000000e+00,  2.22044605e-16]])\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quat = km.UnitQuaternion(math.pi / 2, [0.0, 1.0, 0.0])\n",
    "quat.toRotationMatrix()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0208e41a",
   "metadata": {},
   "source": [
    "## Homogeneous Transforms\n",
    "\n",
    "In **kdFlex**, a homogeneous transform, or `HomTran`, consists of a translation together with a rotation. Here we build the transform which translates all components by $100$ and rotates the $y$-axis to the $z$-axis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cd5a8490",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([101.,  97., 102.])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transform = km.HomTran(\n",
    "    vec=[100, 100, 100],\n",
    "    q=km.UnitQuaternion([0, 1, 0], [0, 0, 1]),  # rotates to send y-component to z-axis\n",
    ")\n",
    "\n",
    "transform * [1, 2, 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdebb3b7",
   "metadata": {},
   "source": [
    "Applying the inverse transform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c2bba569",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-99., -97.,  98.])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transform.inverse() * [1, 2, 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70a8623b",
   "metadata": {},
   "source": [
    "Transforms can also be composed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "02b0d4e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.HomTran(\n",
       "\tKarana.Math.UnitQuaternion(0, 0, 0, 1),\n",
       "\tarray([0.00000000e+00, 4.26325641e-14, 0.00000000e+00])\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transform * transform.inverse()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21fa7c82",
   "metadata": {},
   "source": [
    "## Transforming spatial vectors\n",
    "\n",
    "`HomTran` also provides a rich set of SOA operations. For example, we can use the `phiStar` operator to transform a `SpatialVelocity` to another frame on a rigid body:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7acf2f09-802f-4c76-aabd-14f910c7af6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.SpatialVelocity(\n",
       "\t<Quantity([1. 2. 3.], '1 / second')>,\n",
       "\t<Quantity([4. 5. 6.], 'meter / second')>\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "T_ab = km.HomTran(\n",
    "    vec=[100, 100, 100],\n",
    "    q=km.UnitQuaternion([0, 1, 0], [0, 0, 1]),  # rotates to send y-component to z-axis\n",
    ")\n",
    "vel_a = km.SpatialVelocity([1, 2, 3], [4, 5, 6])\n",
    "vel_a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2df4cc91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.SpatialVector(\n",
       "\tarray([ 1.,  3., -2.]),\n",
       "\tarray([ -96.,  -94., -205.])\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vel_b = T_ab.phiStar(vel_a)\n",
    "vel_b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d91a8efc-1b4f-4d3f-aa66-bcae6bd29fa6",
   "metadata": {},
   "source": [
    "The `phi` operator can also be used to apply `parallel axis transformation` to `spatial inertias` as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "bb300cb7-3e67-42b6-a49c-f673023c06f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Karana.Math.SpatialInertia(\n",
       "<Quantity(0.7, 'kilogram')>,\n",
       "<Quantity([0.2 0.3 0.4], 'meter')>,\n",
       "<Quantity([[5. 0. 0.]\n",
       " [0. 1. 0.]\n",
       " [0. 0. 1.]], 'kilogram * meter ** 2')>\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i = np.eye(3)\n",
    "i[0, 0] = 5.0\n",
    "M_a = km.SpatialInertia(0.7, np.array([0.2, 0.3, 0.4]), i)\n",
    "M_a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a2b2e099-bf71-4092-a8c0-a0774faf1d77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply parallel axis theorem to shift M_a from 'a' to 'b'\n",
    "M_b = M_a.parallelAxis(T_ab).matrix()\n",
    "\n",
    "# use the phi() operator to do this transformation\n",
    "M_b_alt = T_ab.phiMatrix() @ M_a.matrix() @ T_ab.phiMatrix().transpose()\n",
    "\n",
    "# verify that the two paths give the same result\n",
    "assert np.allclose(M_b - M_b_alt, 0, atol=1e-11)"
   ]
  }
 ],
 "metadata": {
  "docs": {
   "tags": [
    "math", 
    "beginner"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}