{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial Rotations\n",
    "\n",
    "This tutorial demonstrates how to use quaternions for 3D rotations and visualize the results with a spherical cap. We'll use `jax.scipy.spatial.transform.Rotation` to create initial rotations and convert them to FastQuat quaternions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from jax.scipy.spatial.transform import Rotation\n",
    "\n",
    "from fastquat import Quaternion\n",
    "\n",
    "# Configure matplotlib for better rendering\n",
    "plt.rcParams['figure.dpi'] = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Initial Rotation Around X-axis Using JAX Scipy\n",
    "\n",
    "We'll create a 30° rotation around the X-axis using `jax.scipy.spatial.transform.Rotation.from_euler` and convert it to a FastQuat quaternion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create initial rotation using scipy's Rotation around X-axis\n",
    "rotation_main = Rotation.from_euler('x', 30.0, degrees=True).as_matrix()\n",
    "\n",
    "# Convert to FastQuat quaternion\n",
    "q_main = Quaternion.from_rotation_matrix(rotation_main)\n",
    "\n",
    "print(f'Main rotation quaternion (30° around X): {q_main}')\n",
    "print(f'Main quaternion norm: {abs(q_main):.6f}')\n",
    "print(f'Rotation matrix shape: {rotation_main.shape}')\n",
    "\n",
    "k_vector = jnp.array([0, 0, 1])\n",
    "rotated_k_vector = q_main.rotate_vector(k_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Display\n",
    "\n",
    "Create a 3D visualization showing the unit sphere, the original and rotated k vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_visualization(original_vector, rotated_vector):\n",
    "    \"\"\"Create comprehensive visualization of quaternion rotation and spherical cap.\"\"\"\n",
    "\n",
    "    fig = plt.figure(figsize=(15, 6))\n",
    "\n",
    "    # Generate unit sphere for visualization\n",
    "    u = np.linspace(0, 2 * np.pi, 50)\n",
    "    v = np.linspace(0, np.pi, 50)\n",
    "    sphere_x = np.outer(np.cos(u), np.sin(v))\n",
    "    sphere_y = np.outer(np.sin(u), np.sin(v))\n",
    "    sphere_z = np.outer(np.ones(np.size(u)), np.cos(v))\n",
    "\n",
    "    # Left plot: Vector rotation demonstration\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "    ax.plot_surface(sphere_x, sphere_y, sphere_z, alpha=0.2, color='lightblue')\n",
    "\n",
    "    # Original k vector (blue)\n",
    "    ax.quiver(\n",
    "        0,\n",
    "        0,\n",
    "        0,\n",
    "        *original_vector,\n",
    "        color='blue',\n",
    "        arrow_length_ratio=0.1,\n",
    "        linewidth=4,\n",
    "        label='Original k',\n",
    "    )\n",
    "\n",
    "    # Rotated k vector (red)\n",
    "    ax.quiver(\n",
    "        0,\n",
    "        0,\n",
    "        0,\n",
    "        *rotated_vector,\n",
    "        color='red',\n",
    "        arrow_length_ratio=0.1,\n",
    "        linewidth=4,\n",
    "        label='Rotated k (30° around X)',\n",
    "    )\n",
    "\n",
    "    # Show rotation arc\n",
    "    theta_arc = np.linspace(0, np.pi / 6, 20)  # 20 degrees in radians\n",
    "    arc_x = np.zeros_like(theta_arc)\n",
    "    arc_y = -np.sin(theta_arc)\n",
    "    arc_z = np.cos(theta_arc)\n",
    "    ax.plot(arc_x, arc_y, arc_z, 'g--', linewidth=2, alpha=0.7, label='Rotation Arc')\n",
    "\n",
    "    ax.set_xlabel('X')\n",
    "    ax.set_ylabel('Y')\n",
    "    ax.set_zlabel('Z')\n",
    "    ax.set_title('Quaternion Rotation of Unit Vector k\\n(30° around X-axis)')\n",
    "    ax.legend()\n",
    "    ax.set_box_aspect([1, 1, 1])\n",
    "    #    plt.tight_layout()\n",
    "    # plt.savefig('spherical_cap_visualization.png', dpi=150, bbox_inches='tight')\n",
    "    # plt.show()\n",
    "    return fig\n",
    "\n",
    "\n",
    "fig = create_visualization(k_vector, rotated_k_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create Small Rotation for Spherical Cap\n",
    "\n",
    "We'll create a 1° rotation around the Z-axis to generate the spherical cap through quaternion multiplication and we'll create the spherical cap by repeatedly applying the small rotation to demonstrate quaternion multiplication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create small rotation for spherical cap of 1 degree rotation around the Z-axis\n",
    "small_rotation = Rotation.from_euler('z', 1, degrees=True)\n",
    "small_matrix = small_rotation.as_matrix()\n",
    "q_small = Quaternion.from_rotation_matrix(small_matrix)\n",
    "\n",
    "print(f'Small rotation quaternion (1° around Z): {q_small}')\n",
    "print(f'Small quaternion norm: {abs(q_small):.6f}')\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def apply_small_rotation(quaternion):\n",
    "    quaternion = q_small * quaternion\n",
    "    vector = quaternion.rotate_vector(k_vector)\n",
    "    return quaternion, vector\n",
    "\n",
    "\n",
    "def create_spherical_cap():\n",
    "    \"\"\"Create a spherical cap using quaternion multiplication.\"\"\"\n",
    "\n",
    "    # Create cap by applying small rotations iteratively\n",
    "    points = []\n",
    "\n",
    "    # Generate 360 points by repeatedly applying 1° rotation\n",
    "    q_accumulated = q_main\n",
    "    for i in range(360):\n",
    "        q_accumulated, rotated_point = apply_small_rotation(q_accumulated)\n",
    "        points.append(rotated_point)\n",
    "    return jnp.array(points)\n",
    "\n",
    "\n",
    "# Generate the spherical cap\n",
    "cap_points = create_spherical_cap()\n",
    "\n",
    "print(f'Original k vector: {k_vector}')\n",
    "print(f'Rotated k vector: {rotated_k_vector}')\n",
    "print(f'Generated {len(cap_points)} points for spherical cap')\n",
    "print(f'Cap points shape: {cap_points.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Display\n",
    "\n",
    "Create a 3D visualization showing the unit sphere, the original and rotated k vectors, and the spherical cap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = create_visualization(k_vector, rotated_k_vector)\n",
    "ax = plt.gca()\n",
    "\n",
    "# Plot the spherical cap as connected points\n",
    "ax.plot(*cap_points.T, 'ro-', markersize=3, linewidth=1, alpha=0.8, label='Spherical Cap')\n",
    "ax.legend()\n",
    "ax.set_box_aspect([1, 1, 1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mathematical Background\n",
    "\n",
    "This example demonstrates several key concepts:\n",
    "\n",
    "### Quaternion Representation\n",
    "The rotation of a vector **v** by a unit quaternion **q** is given by:\n",
    "$$\\mathbf{v}' = q \\mathbf{v} q^{-1}$$\n",
    "\n",
    "### Quaternion Multiplication\n",
    "Quaternion multiplication represents composition of rotations:\n",
    "$$q_3 = q_2 \\cdot q_1$$\n",
    "means \"first apply rotation $q_1$, then apply rotation $q_2$\".\n",
    "\n",
    "### Spherical Cap Generation\n",
    "By iteratively applying small rotations:\n",
    "$$q_n = q_{\\text{small}} \\cdot q_{n-1}$$\n",
    "\n",
    "we generate points that form a spherical cap around the main rotation axis.\n",
    "\n",
    "## Key Insights\n",
    "\n",
    "1. **Conversion Fidelity**: Converting from Euler angles → rotation matrix → quaternion preserves the rotation exactly\n",
    "2. **Quaternion Multiplication**: Represents smooth composition of rotations\n",
    "3. **Normalization**: All operations preserve quaternion normalization automatically\n",
    "4. **Geometric Consistency**: The spherical cap forms a perfect circle on the unit sphere\n",
    "\n",
    "This makes FastQuat ideal for applications requiring precise 3D rotations, such as:\n",
    "- Robotics and mechanical engineering\n",
    "- Computer graphics and animation\n",
    "- Aerospace and navigation systems\n",
    "- Scientific visualization"
   ]
  }
 ],
 "metadata": {
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}