{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geometric Algebra and Its Applications in Deep Learning\n",
    "\n",
    "## Outline\n",
    "\n",
    "1. [Introduction to Geometric Algebra](#Introduction)\n",
    "\n",
    "2. [Geometric Algebra in Deep Learning](#key-papers-leveraging-clifford-algebra-in-deep-learning)\n",
    "   \n",
    "\n",
    "3. [Camera Pose Prediction using GA-based NN](#camera-pose-predicition-using-a-neural-network-based-on-geometric-algebra)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "Geometric Algebra is a rich, unifying language that extends traditional algebraic systems. It blends the simplicity of vector algebra with the depth of complex numbers and quaternions, offering a coordinate-free approach to geometry. GA not only simplifies calculations and clarifies geometric intuition, but it also reveals underlying symmetries in physical laws and provides a robust framework for applications in fields ranging from physics and engineering to computer graphics and deep learning. To see how this works in practice, we can start by looking at some basic geometric primitives and then explore the different products that GA uses to combine and interpret them.\n",
    "\n",
    "We begin with the idea of vectors, which have both magnitude (length) and direction in space. In a 2D plane, it is common to work with the unit vectors\n",
    "$$\n",
    "\\hat{x} = \\begin{pmatrix}1 \\\\ 0\\end{pmatrix}, \\quad \n",
    "\\hat{y} = \\begin{pmatrix}0 \\\\ 1\\end{pmatrix},\n",
    "$$\n",
    "so that any 2D vector $\\vec{v}$ can be written as\n",
    "$$\n",
    "\\vec{v} = v_x \\hat{x} + v_y \\hat{y}.\n",
    "$$\n",
    "Moving to 3D adds $\\hat{z}$:\n",
    "$$\n",
    "\\hat{x} = \\begin{pmatrix}1 \\\\ 0 \\\\ 0\\end{pmatrix}, \\quad \n",
    "\\hat{y} = \\begin{pmatrix}0 \\\\ 1 \\\\ 0\\end{pmatrix}, \\quad\n",
    "\\hat{z} = \\begin{pmatrix}0 \\\\ 0 \\\\ 1\\end{pmatrix}.\n",
    "$$\n",
    "These simple elements are central to GA because they form the “building blocks” of more complex geometric objects. Meanwhile, scalars such as $a \\in \\mathbb{R}$ are just numbers, while vectors like $\\vec{v} = (v_1, v_2, v_3)$ capture direction. When two vectors in 2D define a parallelogram, their oriented area is expressed by $\\vec{a} \\wedge \\vec{b}$, sometimes called a directed area. GA pays attention not only to magnitudes (like lengths or areas) but also to orientation (how those elements are “pointing”).\n",
    "\n",
    "To make sense of these areas, GA introduces bivectors, which represent the size and handedness of a plane. In 2D, the wedge product $\\vec{a} \\wedge \\vec{b}$ becomes a scalar multiple of the unit bivector $\\hat{x}\\hat{y}$, whose magnitude is given by $|a_1 b_2 - a_2 b_1|$ and whose sign encodes orientation. In 3D, if we add a third vector $\\vec{c}$, then\n",
    "$$\n",
    "\\vec{a} \\wedge \\vec{b} \\wedge \\vec{c}\n",
    "$$\n",
    "defines a trivector, corresponding to the oriented volume of the parallelepiped described by $\\vec{a}$, $\\vec{b}$, and $\\vec{c}$. Swapping any two vectors changes the sign of this volume, reflecting a reversal in orientation.\n",
    "\n",
    "<img src=\"https://www.researchgate.net/profile/A-Arsenovic/publication/318475103/figure/fig3/AS:547679924166656@1507588477660/nterpretation-of-some-elements-within-Geometric-Algebra-highlighting-their-creation.png\" alt=\"drawing\" width=\"600\"/>\n",
    "\n",
    "To combine these geometric primitives, we make use of products. One well-known example is the dot product,\n",
    "$$\n",
    "\\vec{a} \\cdot \\vec{b} = a_1 b_1 + a_2 b_2 + a_3 b_3,\n",
    "$$\n",
    "which gives a scalar measuring the alignment of two 3D vectors. Another is the wedge (exterior) product, denoted $\\wedge$, which indicates the oriented area spanned by two vectors. For instance, in 2D,\n",
    "$$\n",
    "\\vec{a} \\wedge \\vec{b} = (a_1 b_2 - a_2 b_1)\\,\\hat{x}\\hat{y},\n",
    "$$\n",
    "with the result being zero if $\\vec{a}$ and $\\vec{b}$ are parallel. Geometric Algebra merges these two ideas into a single geometric product $\\vec{a}\\,\\vec{b}$:\n",
    "$$\n",
    "\\vec{a}\\,\\vec{b} = \\vec{a} \\cdot \\vec{b} + \\vec{a} \\wedge \\vec{b}.\n",
    "$$\n",
    "This product is associative and distributive over addition, making it a powerful unifying tool for encoding both metric (dot) and orientation (wedge) information at once. When the vectors in question are orthogonal unit vectors like $\\hat{x}$ and $\\hat{y}$, the geometric product simplifies:\n",
    "$$\n",
    "\\hat{x}\\,\\hat{y} = \\hat{x}\\hat{y}, \\quad \\text{and} \\quad \\hat{y}\\,\\hat{x} = -\\,\\hat{x}\\,\\hat{y}.\n",
    "$$\n",
    "In 3D, taking vectors $\\vec{x} = (a_1, b_1, c_1)$ and $\\vec{y} = (a_2, b_2, c_2)$ yields\n",
    "$$\n",
    "\\vec{x}\\,\\vec{y} = (a_1a_2 + b_1b_2 + c_1c_2)\n",
    "\\;+\\; (a_1b_2 - a_2b_1)\\,\\hat{x}\\hat{y}\n",
    "\\;+\\; (b_1c_2 - b_2c_1)\\,\\hat{y}\\hat{z}\n",
    "\\;+\\; (a_1c_2 - a_2c_1)\\,\\hat{x}\\hat{z}.\n",
    "$$\n",
    "This comprehensive form includes both the dot product (the scalar part) and the bivector part (terms with $\\hat{x}\\hat{y}$, $\\hat{y}\\hat{z}$, or $\\hat{x}\\hat{z}$). One can see how the wedge product in 3D is related to the cross product $\\vec{a} \\times \\vec{b}$, which is a perpendicular vector to the plane of $\\vec{a}$ and $\\vec{b}$; in GA, $\\vec{a} \\wedge \\vec{b}$ is a bivector representing that same plane.\n",
    "\n",
    "When we move from pure vectors to general multivectors, we can store scalars, vectors, bivectors, and even trivectors together. In 3D, a typical multivector might be written as\n",
    "$$\n",
    "M = a + \\vec{v} + B + T,\n",
    "$$\n",
    "where $a$ is a scalar, $\\vec{v}$ is a vector, $B$ is a bivector, and $T$ is a trivector. Multiplying two such objects means applying the geometric product across all their parts. But it also remains enlightening to see how things play out in the simpler 2D setting. A 2D multivector\n",
    "$$\n",
    "M = a + b\\,\\hat{x} + c\\,\\hat{y} + d\\,\\hat{x}\\hat{y}\n",
    "$$\n",
    "can encode scalars, vectors, and the plane itself via $\\hat{x}\\hat{y}$. A key observation is\n",
    "$$\n",
    "(\\hat{x}\\hat{y})^2 = -1,\n",
    "$$\n",
    "which matches the defining property of the imaginary unit $i$ in complex numbers. This similarity lets us express plane rotations using exponentials of the bivector $\\hat{x}\\hat{y}$. For a rotation by angle $\\theta$ in 2D,\n",
    "$$\n",
    "e^{\\hat{x}\\hat{y}\\theta} = \\cos\\theta + (\\hat{x}\\hat{y})\\,\\sin\\theta,\n",
    "$$\n",
    "and applying it to a vector $\\vec{v}$ rotates $\\vec{v}$ by $\\theta$. If we represent two 2D multivectors $M_1$ and $M_2$ in polar form, their product behaves much like multiplying two complex numbers:\n",
    "$$\n",
    "M_1 M_2 = |M_1||M_2|\\,e^{\\hat{x}\\hat{y}(\\theta_1+\\theta_2)}.\n",
    "$$\n",
    "------------------\n",
    "Extending this approach to 3D leads us to the idea of a 3D multivector of the form\n",
    "$$\n",
    "M = a + (b\\,\\hat{x} + c\\,\\hat{y} + d\\,\\hat{z})\n",
    "  + (e\\,\\hat{x}\\hat{y} + f\\,\\hat{y}\\hat{z} + g\\,\\hat{x}\\hat{z})\n",
    "  + h\\,\\hat{x}\\hat{y}\\hat{z}.\n",
    "$$\n",
    "\n",
    "Here, we often set \n",
    "$$\n",
    "I = \\hat{x}\\hat{y}\\hat{z},\n",
    "$$\n",
    "which acts like a “3D imaginary unit” because\n",
    "$$\n",
    "I^2 = -1.\n",
    "$$\n",
    "This pseudoscalar $I$ represents an oriented volume element. Sometimes you will see a 3D multivector labeled simply as\n",
    "$$\n",
    "a + \\vec{u} + \\vec{v}\\,I + h\\,I,\n",
    "$$\n",
    "where $\\vec{u}$ and $\\vec{v}$ are ordinary 3D vectors, and $a$ and $h$ remain scalars. This identification bridges the concept of **imaginary units** in 2D (as bivectors) with volumes in 3D.\n",
    "\n",
    "\n",
    "<img src=\"https://www.euclideanspace.com/maths/algebra/clifford/d3/cayleyDigraph1.png\" alt=\"drawing\" width=\"600\"/>\n",
    "\n",
    "---\n",
    "\n",
    "### Quaternions vs. Rotors in GA\n",
    "\n",
    "A key application of 3D Geometric Algebra is describing **rotations** through **rotors**, which are special **unit multivectors** in the even subalgebra (i.e., containing scalars and bivectors but no pure grade-1 or grade-3 parts). If $R$ is such a rotor and $\\vec{v}$ is a 3D vector, the rotated version of $\\vec{v}$ is obtained via the **sandwich product**:\n",
    "\n",
    "$$\n",
    "\\vec{v}' = R \\,\\vec{v}\\, R^{-1}.\n",
    "$$\n",
    "\n",
    "This construction **generalizes** the familiar use of **quaternions** to handle 3D rotations, since quaternions can be viewed as a particular subalgebra of the full 3D Clifford Algebra:\n",
    "\n",
    "- A typical quaternion is written as\n",
    "  $$\n",
    "  q = w + x\\,\\mathbf{i} + y\\,\\mathbf{j} + z\\,\\mathbf{k},\n",
    "  $$\n",
    "  with the well-known rules\n",
    "  $$\n",
    "  \\mathbf{i}^2 = \\mathbf{j}^2 = \\mathbf{k}^2 = \\mathbf{i}\\mathbf{j}\\mathbf{k} = -1.\n",
    "  $$\n",
    "- In Geometric Algebra, these $\\mathbf{i}, \\mathbf{j}, \\mathbf{k}$ can be naturally mapped to specific bivectors like $\\hat{x}\\hat{y}$, $\\hat{y}\\hat{z}$, etc.  \n",
    "  For instance, one common identification is:\n",
    "  $$\n",
    "  \\mathbf{i}\\,\\mapsto\\, \\hat{y}\\hat{z}, \\quad\n",
    "  \\mathbf{j}\\,\\mapsto\\, \\hat{z}\\hat{x}, \\quad\n",
    "  \\mathbf{k}\\,\\mapsto\\, \\hat{x}\\hat{y}.\n",
    "  $$\n",
    "  That means the quaternion part $x\\,\\mathbf{i} + y\\,\\mathbf{j} + z\\,\\mathbf{k}$ corresponds to some linear combination of those three bivectors in GA.\n",
    "\n",
    "Hence, **quaternions** are essentially a **subset** of 3D Clifford Algebra restricted to the scalar plus “imaginary” (bivector) directions. The same logic for rotating a vector via\n",
    "$$\n",
    "\\vec{v}' = q\\,\\vec{v}\\,q^{-1}\n",
    "$$\n",
    "in the quaternion sense translates to a sandwich product in GA.\n",
    "\n",
    "---\n",
    "\n",
    "### Building a Rotor\n",
    "\n",
    "Rotors in GA are exponentials of bivectors. Suppose you want a rotation by angle $\\theta$ around a unit axis $\\hat{u}$. Then the relevant bivector is in the plane perpendicular to $\\hat{u}$, and you can write something like\n",
    "\n",
    "$$\n",
    "R = \\exp\\Bigl(\\tfrac{\\theta}{2}\\,\\hat{u}\\Bigr) \\quad \\text{or, concretely,} \\quad\n",
    "R = \\cos\\Bigl(\\tfrac{\\theta}{2}\\Bigr) - \\sin\\Bigl(\\tfrac{\\theta}{2}\\Bigr)\\,B,\n",
    "$$\n",
    "\n",
    "where $B$ is a bivector encoding that plane. Applying $R$ to a vector $\\vec{v}$ is done by the sandwich product $R\\,\\vec{v}\\,R^{-1}$, preserving the length of $\\vec{v}$ and rotating it as intended. If you only need a rotation around, say, the $\\hat{z}$ axis, then that bivector $B$ is just $\\hat{x}\\hat{y}$. In that case,\n",
    "\n",
    "$$\n",
    "R = \\cos\\Bigl(\\tfrac{\\theta}{2}\\Bigr) - \\sin\\Bigl(\\tfrac{\\theta}{2}\\Bigr)\\,\\hat{x}\\hat{y}.\n",
    "$$\n",
    "\n",
    "\n",
    "\n",
    "------------------\n",
    "\n",
    "Ultimately, one of the most dramatic illustrations of GA’s unifying power appears when we rewrite Maxwell’s equations. Traditionally, these four separate laws govern the electric field $\\vec{E}$, magnetic field $\\vec{B}$, and their interactions with charges and currents:\n",
    "1. $$\\nabla \\cdot \\vec{E} = \\frac{\\rho}{\\epsilon_0},$$\n",
    "2. $$\\nabla \\cdot \\vec{B} = 0,$$\n",
    "3. $$\\nabla \\times \\vec{E} = -\\frac{\\partial\\vec{B}}{\\partial t},$$\n",
    "4. $$\\nabla \\times \\vec{B} = \\mu_0 \\Bigl(\\vec{J} + \\epsilon_0 \\frac{\\partial\\vec{E}}{\\partial t}\\Bigr).$$\n",
    "But in GA, we can group $\\vec{E}$ and $\\vec{B}$ into a single bivector\n",
    "$$\n",
    "F = \\vec{E} + I\\,c\\,\\vec{B},\n",
    "$$\n",
    "where $I = \\hat{x}\\hat{y}\\hat{z}$, and then define a vector derivative\n",
    "$$\n",
    "\\nabla = \\frac{1}{c} \\frac{\\partial}{\\partial t} +  \\hat{x}\\frac{\\partial}{\\partial x} + \\hat{y}\\frac{\\partial}{\\partial y} + \\hat{z}\\frac{\\partial}{\\partial z}.\n",
    "$$\n",
    "All four Maxwell’s equations collapse neatly into one statement:\n",
    "$$\n",
    "\\nabla F = J,\n",
    "$$\n",
    "where $J = \\frac{1}{\\epsilon_0}(c\\rho - \\vec{J})$ represents the electromagnetic current (encompassing charge density and current density). This compact formula elegantly captures the geometry and symmetry of electromagnetism, showing how Geometric Algebra can spotlight the underlying structure of physical laws.\n",
    "\n",
    "By following these threads—from vectors to bivectors and trivectors, from 2D plane rotations to 3D rotors and volumes, and from separate field equations to a single unified electromagnetic statement—we see how GA gracefully weaves together concepts that might otherwise feel disjoint. It stands as a versatile tool not only for clarifying advanced geometry and physics but also for offering a powerful perspective that can be harnessed in modern applications like computer graphics, robotics, and deep learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install numpy matplotlib torch torchvision cliffordlayers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 1: Implementing the Geometric Product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_multivector(s=0, vx=0, vy=0, vz=0, bxy=0, byz=0, bxz=0, txyz=0):\n",
    "    return [s, vx, vy, vz, bxy, byz, bxz, txyz]\n",
    "\n",
    "def print_multivector(M):\n",
    "    # M is [s, vx, vy, vz, bxy, byz, bxz, txyz]\n",
    "    s, vx, vy, vz, bxy, byz, bxz, txyz = M\n",
    "    print(f\"s={s} + v=({vx},{vy},{vz}) + b=({bxy},{byz},{bxz}) + txyz={txyz}\")\n",
    "\n",
    "\n",
    "def geometric_product(M1, M2):\n",
    "    # Unpack\n",
    "    s1, vx1, vy1, vz1, bxy1, byz1, bxz1, t1 = M1\n",
    "    s2, vx2, vy2, vz2, bxy2, byz2, bxz2, t2 = M2\n",
    "    \n",
    "    # Prepare result container\n",
    "    # We'll accumulate results for [scalar, vx, vy, vz, bxy, byz, bxz, trivector]\n",
    "    result = [0.0]*8\n",
    "    \n",
    "    # Compute each pair of terms from M1, M2\n",
    "    # e.g. s1 * s2 -> goes to scalar part\n",
    "    # e.g. s1 * vx2 -> goes to vx of result\n",
    "    # e.g. vx1 * vy2 -> goes to bxy of result (with sign handling)\n",
    "    # and so forth. You will need quite a few lines to handle all combinations.\n",
    "    \n",
    "    # TODO: Fill in systematically. For instance:\n",
    "    # 1) Scalar * Scalar\n",
    "    result[0] += s1*s2\n",
    "    # 2) Scalar * Vector plus Vector * Scalar => goes to vector\n",
    "    result[1] += s1*vx2 + vx1*s2   # x-component\n",
    "    # etc...\n",
    "    \n",
    "    return result\n",
    "\n",
    "# Testing\n",
    "mx = create_multivector(vx=1)  # x\n",
    "my = create_multivector(vy = 1)  # y\n",
    "prod_xy = geometric_product(mx, my)\n",
    "print(\"x*y =\", prod_xy)  # Should reflect b_xy in the 4th slot (index=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exercise 2: Building 2D Rotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "def build_rotZ_rotor(theta):\n",
    "    # cos(theta) -> scalar part\n",
    "    # sin(theta) -> b_xy part\n",
    "    return create_multivector(math.cos(theta), 0, 0, 0, math.sin(theta), 0, 0, 0)\n",
    "\n",
    "def reverse_multivector(M):\n",
    "    s, vx, vy, vz, bxy, byz, bxz, txyz = M\n",
    "    # TODO: reverse the multivector\n",
    "    pass\n",
    "\n",
    "def apply_rotZ(M_vec, R):\n",
    "    R_rev = reverse_multivector(R)\n",
    "    left = geometric_product(R, M_vec)\n",
    "    rotated = geometric_product(left, R_rev)\n",
    "    return rotated\n",
    "\n",
    "# Testing\n",
    "R_90 = build_rotZ_rotor(math.pi/2)\n",
    "vec_x = create_multivector(0, 1, 0, 0, 0, 0, 0, 0)  # (1,0)\n",
    "rotated_vec = apply_rotZ(vec_x, R_90)\n",
    "print(\"Rotated (1,0) by 90 deg around the z axis =>\", rotated_vec)\n",
    "print_multivector(rotated_vec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3: 3D Rotations Around a Given Axis\n",
    "\n",
    "Task\n",
    "1.\tCreate build_3D_rotor(theta) for rotation about the z-axis.\n",
    "2.\tWrite apply_rotor_3D(M_vec, R) to do the sandwich product in 3D.\n",
    "3.\tConfirm that rotating a vector (x,y,z) around $\\hat{z}$ by 90 degrees maps (x,y,z) to (-y,x,z)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key Papers Leveraging Clifford Algebra in Deep Learning\n",
    "\n",
    "In recent years, several groundbreaking works have demonstrated how Clifford (geometric) algebras can be leveraged to build more expressive and symmetry-aware deep learning models. The following key papers introduce novel architectures that exploit multivector representations, geometric products, and group equivariance to address challenging tasks such as dynamical system modeling, PDE simulation, and 3D physical transformations.\n",
    "\n",
    "---\n",
    "\n",
    "### Geometric Clifford Algebra Networks (GCANs)\n",
    "\n",
    "**Citation:**  \n",
    "*David Ruhe, Jayesh K. Gupta, Steven De Keninck, Max Welling, Johannes Brandstetter (2023). Geometric Clifford Algebra Networks. 40th International Conference on Machine Learning (ICML), 2023.*\n",
    "\n",
    "**Summary:**  \n",
    "GCANs are designed for modeling dynamical systems by leveraging symmetry group transformations encoded in geometric (Clifford) algebras. This work begins with a concise review of modern plane-based geometric algebra—built upon isometries represented by elements of the Pin(p, q, r) group—and introduces the concept of group action layers. These layers linearly combine object transformations using pre-specified group actions, while a novel activation and normalization scheme allows these geometric templates to be refined via gradient descent. The approach shows significant improvements in modeling three-dimensional rigid body transformations and large-scale fluid dynamics simulations over traditional methods.\n",
    "\n",
    "**Illustration:**  \n",
    "![GCANs Illustration](https://brandstetter-johannes.github.io/publication/ruhe-2023-cgans/featured_hu09f53f46a7368d6e71b797c2b4bae77a_608561_720x0_resize_lanczos_3.png)\n",
    "\n",
    "**Code:**  \n",
    "[https://github.com/microsoft/cliffordlayers](https://github.com/microsoft/cliffordlayers)\n",
    "\n",
    "---\n",
    "\n",
    "### Clifford Neural Layers for PDE Modeling\n",
    "\n",
    "**Citation:**  \n",
    "*Johannes Brandstetter, Rianne Van Den Berg, Max Welling, Jayesh Gupta (2022). Clifford Neural Layers for PDE Modeling. 11th International Conference on Learning Representations (ICLR), 2023.*\n",
    "\n",
    "**Summary:**  \n",
    "This paper addresses the challenge of simulating complex physical processes described by partial differential equations (PDEs). The authors propose replacing traditional convolution and Fourier operations with Clifford neural layers, which incorporate scalar, vector, and higher-order (bivector, trivector) components into the network. By viewing the evolution of correlated fields as multivector fields, the method captures the intrinsic relationships among different field components, thereby improving the generalization capabilities of neural PDE surrogates. The proposed layers have been successfully applied to tasks such as 2D Navier–Stokes simulations, weather forecasting, and 3D Maxwell equations.\n",
    "\n",
    "**Illustration:**  \n",
    "![Clifford Neural Layers Illustration](https://figures.semanticscholar.org/84959e211a767f902cbf1695ec54a5b50148020f/7-Figure5-1.png)\n",
    "\n",
    "**Code:**  \n",
    "[https://microsoft.github.io/cliffordlayers/](https://microsoft.github.io/cliffordlayers/)\n",
    "\n",
    "---\n",
    "\n",
    "### Clifford Group Equivariant Neural Networks\n",
    "\n",
    "**Citation:**  \n",
    "*David Ruhe, Johannes Brandstetter, Patrick Forré (2023). Clifford Group Equivariant Neural Networks. Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS), 2023.*\n",
    "\n",
    "**Summary:**  \n",
    "This work introduces a novel approach to constructing O(n)- and E(n)-equivariant neural network models by harnessing the properties of the Clifford group—a subgroup within the Clifford algebra that preserves the full multivector grading and the geometric product. The authors prove that every polynomial in multivectors (including its grade projections) defines an equivariant map with respect to the Clifford group, which allows for the design of expressive and dimension-agnostic network layers. The paper reports state-of-the-art performance on diverse tasks such as 3D n-body simulations, 4D Lorentz-equivariant high-energy physics experiments, and 5D convex hull experiments.\n",
    "\n",
    "**Illustration:**  \n",
    "![Clifford Group Equivariant Networks Illustration](https://raw.githubusercontent.com/DavidRuhe/clifford-group-equivariant-neural-networks/master/assets/figure.png)\n",
    "\n",
    "**Code:**  \n",
    "[https://github.com/DavidRuhe/clifford-group-equivariant-neural-networks](https://github.com/DavidRuhe/clifford-group-equivariant-neural-networks)\n",
    "\n",
    "---\n",
    "\n",
    "### Clifford-Steerable Convolutional Neural Networks (CS-CNNs)\n",
    "\n",
    "**Citation:**  \n",
    "*Maksim Zhdanov, David Ruhe, Maurice Weiler, Ana Lucic, Johannes Brandstetter, Patrick Forr'e (2024). Clifford-Steerable Convolutional Neural Networks. International Conference on Machine Learning (ICML), 2024.*\n",
    "\n",
    "**Summary:**  \n",
    "CS-CNNs extend the use of Clifford algebra to convolutional neural networks by processing multivector fields on pseudo-Euclidean spaces $ R^{p,q} $. The proposed architecture achieves equivariance (e.g., E(3)-equivariance in $ \\mathbb{R}^3 $ and Poincaré-equivariance in Minkowski spacetime $ \\mathbb{R}^{1,3} $) through an implicit parametrization of $ O(p, q) $-steerable kernels via Clifford group equivariant networks. This novel approach significantly outperforms traditional CNN methods in tasks involving fluid dynamics and relativistic electrodynamics forecasting.\n",
    "\n",
    "**Illustration:**  \n",
    "![CS-CNNs Illustration](https://figures.semanticscholar.org/60ef44cb0455ebda60233958e8db78df867f03b9/1-Figure1-1.png)\n",
    "\n",
    "**Code:**  \n",
    "[https://github.com/maxxxzdn/clifford-group-equivariant-cnns](https://github.com/maxxxzdn/clifford-group-equivariant-cnns)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Camera Pose predicition using a Neural network based on Geometric Algebra\n",
    "\n",
    "\n",
    "In this part, we aim to predict the **camera pose** (rotation and translation) of a virtual camera viewing a colored cube, leveraging **Clifford Algebra** and **CliffordNet2d** for a deep learning approach.\n",
    "\n",
    "1. **Generate 3D Points from a Colored Cube**\n",
    "   - Sample points on each face of a cube (e.g., faces centered at $x=\\pm1$, $y=\\pm1$, $z=\\pm1$).\n",
    "   - Assign each face a distinct color.\n",
    "\n",
    "2. **Simulate Multiple Camera Views**\n",
    "   - Place a virtual camera on a sphere around the cube.\n",
    "   - For each camera position, generate synthetic 2D images (or 2D point projections) of the cube. These images serve as the training data.\n",
    "   - Associate each image with the “ground truth” rotation (quaternion) and translation that define the camera pose.\n",
    "\n",
    "\n",
    "3. **Multivector MLP in 3D**\n",
    "   - Define an MLP that processes batches of 3D multivectors. \n",
    "   - Internally, use the 8 basis components of a 3D Clifford algebra and specialized linear transformations via Clifford kernels.\n",
    "   - Use group normalization for multivectors (e.g., `CliffordGroupNorm3d`).\n",
    "\n",
    "4. **CliffordNet2d Pipeline**\n",
    "   - Convert a 2D image (shape $(B, C, H, W)$) into a GA-compatible tensor (shape $(B, C, H, W, 3)$) for the **2D** Clifford layers (e.g., `CliffordBasicBlock2d`).\n",
    "   - Pass the image through a **CliffordNet2d** network to learn feature representations in a geometric way. The output maintains shape $(B, C_{\\text{out}}, H, W, 3)$.\n",
    "   - Apply spatial averaging to reduce the image dimensions, then reshape to form **full 3D multivectors** of shape $(B, C_{\\text{out}}, 8)$.\n",
    "   - Feed these 3D multivectors into the **Multivector MLP** to regress the final camera pose.\n",
    "\n",
    "5. **Output and Pose Prediction**\n",
    "   - The model outputs a single multivector (shape $(B, 1, 8)$) per image.\n",
    "   - From these 8 components, extract:\n",
    "     - **Rotation quaternion** from components $[0, 4, 5, 6]$, then normalize.\n",
    "     - **Translation** from components $[1, 2, 3]$.\n",
    "   - Compare predictions to ground truth to train or validate the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a Dataset of Rendered Cube Views\n",
    "\n",
    "In this section we generate a dataset of synthetic images by:\n",
    "\n",
    "1. **Randomly Sampling Points on Each Cube Face:**  \n",
    "   We sample random 3D points on each of the six faces of a cube. Each face is assigned a unique color (e.g. red, green, blue, yellow, cyan, magenta).\n",
    "\n",
    "2. **Projecting the 3D Points into 2D Camera Views:**  \n",
    "   For each view, a camera is placed on a sphere such that it always looks toward the cube center (origin). The camera orientation is defined by a quaternion computed by first aligning the camera’s forward axis to the target direction, and then applying a random roll (spin). Using a pinhole camera model (with provided intrinsics), the 3D points are projected into 2D.\n",
    "\n",
    "3. **Rendering the Projections as an Image:**  \n",
    "   The 2D projections from all faces are rasterized into an image of shape $(3, H, W)$ (with three color channels). Each projected point’s pixel is assigned the face’s color.  \n",
    "   \n",
    "4. **Dataset Formation:**  \n",
    "   Multiple views are generated, and for each view we store the rendered image together with the ground-truth camera pose (position and quaternion).\n",
    "\n",
    "Below is the code that implements these steps, ending with a visualization of a few sample images from the dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# --- Helper functions for quaternions and camera sampling (from previous parts) ---\n",
    "\n",
    "def normalize_quaternion(q):\n",
    "    return q / np.linalg.norm(q)\n",
    "\n",
    "def quaternion_multiply(q1, q2):\n",
    "    w1, x1, y1, z1 = q1\n",
    "    w2, x2, y2, z2 = q2\n",
    "    return np.array([\n",
    "        w1*w2 - x1*x2 - y1*y2 - z1*z2,\n",
    "        w1*x2 + x1*w2 + y1*z2 - z1*y2,\n",
    "        w1*y2 - x1*z2 + y1*w2 + z1*x2,\n",
    "        w1*z2 + x1*y2 - y1*x2 + z1*w2\n",
    "    ])\n",
    "\n",
    "def axis_angle_to_quaternion(axis, angle):\n",
    "    axis = axis / np.linalg.norm(axis)\n",
    "    half_angle = angle / 2.0\n",
    "    s = np.sin(half_angle)\n",
    "    w = np.cos(half_angle)\n",
    "    return np.array([w, *(axis * s)])\n",
    "\n",
    "def quaternion_from_two_vectors(v_from, v_to):\n",
    "    v1 = v_from / np.linalg.norm(v_from)\n",
    "    v2 = v_to / np.linalg.norm(v_to)\n",
    "    dot = np.dot(v1, v2)\n",
    "    if dot < -0.999999:\n",
    "        orth = np.cross(np.array([1,0,0]), v1)\n",
    "        if np.linalg.norm(orth) < 1e-5:\n",
    "            orth = np.cross(np.array([0,1,0]), v1)\n",
    "        orth = orth / np.linalg.norm(orth)\n",
    "        return axis_angle_to_quaternion(orth, np.pi)\n",
    "    if dot > 0.999999:\n",
    "        return np.array([1.0, 0.0, 0.0, 0.0])\n",
    "    axis = np.cross(v1, v2)\n",
    "    w = np.sqrt((1.0 + dot) * 2.0)\n",
    "    q = np.array([w*0.5, axis[0]/w, axis[1]/w, axis[2]/w])\n",
    "    return normalize_quaternion(q)\n",
    "\n",
    "def sample_camera_position(radius=5.0):\n",
    "    theta = np.random.uniform(0, np.pi)\n",
    "    phi = np.random.uniform(0, 2 * np.pi)\n",
    "    x = radius * np.sin(theta) * np.cos(phi)\n",
    "    y = radius * np.sin(theta) * np.sin(phi)\n",
    "    z = radius * np.cos(theta)\n",
    "    return np.array([x, y, z])\n",
    "\n",
    "def sample_camera_pose_on_sphere(radius=5.0):\n",
    "    cam_pos = sample_camera_position(radius)\n",
    "    dir_to_origin = -cam_pos  # since the cube is centered at origin\n",
    "    # look_at_q = quaternion_from_two_vectors(np.array([0, 0, 1]), dir_to_origin)\n",
    "    look_at_q = quaternion_from_two_vectors( dir_to_origin, np.array([0, 0, 1]))\n",
    "    angle = 0#np.random.uniform(0, 2*np.pi)\n",
    "    spin_q = axis_angle_to_quaternion(dir_to_origin, angle)\n",
    "    final_q = quaternion_multiply(spin_q, look_at_q)\n",
    "    return cam_pos, normalize_quaternion(final_q)\n",
    "\n",
    "# Provided function from the first part:\n",
    "def quaternion_to_rotation_matrix(q):\n",
    "    q = normalize_quaternion(q)\n",
    "    q0, q1, q2, q3 = q\n",
    "    R = np.array([\n",
    "        [1 - 2*(q2**2 + q3**2),     2*(q1*q2 - q0*q3),     2*(q1*q3 + q0*q2)],\n",
    "        [    2*(q1*q2 + q0*q3), 1 - 2*(q1**2 + q3**2),     2*(q2*q3 - q0*q1)],\n",
    "        [    2*(q1*q3 - q0*q2),     2*(q2*q3 + q0*q1), 1 - 2*(q1**2 + q2**2)]\n",
    "    ])\n",
    "    return R\n",
    "\n",
    "def project_points(P, q, cam_pos, K):\n",
    "    # t = -quaternion_to_rotation_matrix(q) @ cam_pos\n",
    "    t = np.array([0, 0, np.linalg.norm(cam_pos)])\n",
    "    R = quaternion_to_rotation_matrix(q)\n",
    "    P_homog = np.hstack((P, np.ones((P.shape[0], 1))))\n",
    "    RT = np.hstack((R, t.reshape(3,1)))\n",
    "    projected_homog = (K @ RT) @ P_homog.T\n",
    "    projected = projected_homog[:2, :] / projected_homog[2, :]\n",
    "    return projected.T\n",
    "\n",
    "# --- Functions for cube point sampling and image rendering ---\n",
    "\n",
    "def sample_cube_face_random(constant_value, axis, num_points=100):\n",
    "    \"\"\"\n",
    "    Randomly samples points on a cube face.\n",
    "    For the specified axis (0 for x, 1 for y, 2 for z), the coordinate is fixed at constant_value.\n",
    "    The other two coordinates are sampled uniformly from [-1, 1].\n",
    "    \"\"\"\n",
    "    points = np.random.uniform(-1, 1, (num_points, 3))\n",
    "    points[:, axis] = constant_value\n",
    "    return points\n",
    "\n",
    "# Define a mapping from color names to RGB values (in range 0-1)\n",
    "COLOR_MAP = {\n",
    "    'red': np.array([1, 0, 0]),\n",
    "    'green': np.array([0, 1, 0]),\n",
    "    'blue': np.array([0, 0, 1]),\n",
    "    'yellow': np.array([1, 1, 0]),\n",
    "    'cyan': np.array([0, 1, 1]),\n",
    "    'magenta': np.array([1, 0, 1])\n",
    "}\n",
    "num_points = 10000\n",
    "# Create cube faces with random sampling and distinct colors\n",
    "cube_faces = [\n",
    "    ('red', sample_cube_face_random(1.0, axis=0, num_points=num_points)),    # x = 1\n",
    "    ('green', sample_cube_face_random(-1.0, axis=0, num_points=num_points)),   # x = -1\n",
    "    ('blue', sample_cube_face_random(1.0, axis=1, num_points=num_points)),     # y = 1\n",
    "    ('yellow', sample_cube_face_random(-1.0, axis=1, num_points=num_points)),  # y = -1\n",
    "    ('cyan', sample_cube_face_random(1.0, axis=2, num_points=num_points)),     # z = 1\n",
    "    ('magenta', sample_cube_face_random(-1.0, axis=2, num_points=num_points))  # z = -1\n",
    "]\n",
    "\n",
    "def render_view_with_depth(cube_faces, cam_pos, q, K, H=480, W=640):\n",
    "    \"\"\"\n",
    "    Renders a view of the cube with simple depth sorting.\n",
    "    \"\"\"\n",
    "    image = np.ones((H, W, 3), dtype=np.float32)\n",
    "    \n",
    "    # Compute rotation matrix and translation\n",
    "    R = quaternion_to_rotation_matrix(q)\n",
    "    t = np.array([0, 0, np.linalg.norm(cam_pos)])\n",
    "    \n",
    "    # For each face, compute an average depth (in camera coordinates)\n",
    "    faces_depth = []\n",
    "    for color, points in cube_faces:\n",
    "        # Transform points to camera coordinates: X_cam = R*(X_world) + t\n",
    "        P_cam = (R @ points.T).T + t\n",
    "        avg_depth = np.mean(P_cam[:, 2])  # z coordinate in camera space\n",
    "        faces_depth.append((avg_depth, color, points))\n",
    "    \n",
    "    # Sort faces from farthest to nearest (largest z to smallest z)\n",
    "    faces_depth.sort(key=lambda x: x[0], reverse=True)\n",
    "    \n",
    "    # Render in sorted order\n",
    "    for _, color, points in faces_depth:\n",
    "        pts_2d = project_points(points, q, cam_pos, K)\n",
    "        # print(\"  2D x-range:\", pts_2d[:,0].min(), pts_2d[:,0].max())\n",
    "        # print(\"  2D y-range:\", pts_2d[:,1].min(), pts_2d[:,1].max())\n",
    "        pts_int = np.round(pts_2d).astype(int)\n",
    "        for u, v in pts_int:\n",
    "            if 0 <= u < W and 0 <= v < H:\n",
    "                image[v, u, :] = COLOR_MAP[color]\n",
    "    \n",
    "    return image.transpose(2, 0, 1)\n",
    "\n",
    "def generate_image_dataset_with_depth(cube_faces = None, num_views=10, radius=5.0, K=None, H=480, W=640, num_points = 1000):\n",
    "    if cube_faces is None:\n",
    "        cube_faces = [\n",
    "                    ('red', sample_cube_face_random(1.0, axis=0, num_points=num_points)),    # x = 1\n",
    "                    ('green', sample_cube_face_random(-1.0, axis=0, num_points=num_points)),   # x = -1\n",
    "                    ('blue', sample_cube_face_random(1.0, axis=1, num_points=num_points)),     # y = 1\n",
    "                    ('yellow', sample_cube_face_random(-1.0, axis=1, num_points=num_points)),  # y = -1\n",
    "                    ('cyan', sample_cube_face_random(1.0, axis=2, num_points=num_points)),     # z = 1\n",
    "                    ('magenta', sample_cube_face_random(-1.0, axis=2, num_points=num_points))  # z = -1\n",
    "                ]  \n",
    "    if K is None:\n",
    "        K = np.array([[300, 0, W/2],\n",
    "                      [0, 300, H/2],\n",
    "                      [0, 0, 1]])\n",
    "    \n",
    "    dataset = []\n",
    "    for _ in range(num_views):\n",
    "        cam_pos, q = sample_camera_pose_on_sphere(radius)\n",
    "        img = render_view_with_depth(cube_faces, cam_pos, q, K, H, W)\n",
    "        dataset.append({\n",
    "            'image': img,\n",
    "            'camera_pos': cam_pos,\n",
    "            'quaternion': q\n",
    "        })\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Updated intrinsic matrix with a lower focal length (e.g., 300 instead of 800)\n",
    "W, H = 640, 480\n",
    "fl = 700\n",
    "K = np.array([\n",
    "    [fl,   0, W/2],\n",
    "    [  0, fl, H/2],\n",
    "    [  0,   0,   1]\n",
    "])\n",
    "\n",
    "# Generate a dataset with the updated intrinsic matrix\n",
    "dataset_images = generate_image_dataset_with_depth(cube_faces, num_views=10, radius=5.0, K=K, H=H, W=W)\n",
    "\n",
    "# Visualize a few sample images with the updated intrinsics\n",
    "num_samples = 4\n",
    "fig, axes = plt.subplots(1, num_samples, figsize=(4*num_samples, 4))\n",
    "for i in range(num_samples):\n",
    "    sample = dataset_images[i]\n",
    "    img = sample['image'].transpose(1, 2, 0)\n",
    "    axes[i].imshow(img)\n",
    "    axes[i].set_title(f\"View {i+1}\\nCam Pos: {np.round(sample['camera_pos'],2)}\")\n",
    "    axes[i].axis('off')\n",
    "    # print(sample['camera_pos'])\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing a 3D Multivector Linear Layer and MLP\n",
    "\n",
    "In a geometric algebra (Clifford) setting, each channel is represented as a full multivector with 8 components (in 3D). To process such data, we define a specialized linear layer—**MultivectorLinear**—that:\n",
    "- Accepts an input tensor of shape $(B, C_{\\text{in}}, 8)$.\n",
    "- Uses a Clifford kernel (obtained via `cliffordkernels.get_3d_clifford_kernel`) to map the flattened input (of dimension $8 \\times C_{\\text{in}}$) to an output of dimension $8 \\times C_{\\text{out}}$.\n",
    "- Reshapes the result back to $(B, C_{\\text{out}}, 8)$.\n",
    "\n",
    "Using this module, we then define a **MultivectorMLP** that consists of:\n",
    "1. A MultivectorLinear layer.\n",
    "2. A Clifford-based group normalization (e.g., `gn.CliffordGroupNorm3d`).\n",
    "3. An activation function.\n",
    "4. A second MultivectorLinear layer.\n",
    "5. A final transposition to restore the output shape $(B, C_{\\text{out}}, 8)$.\n",
    "\n",
    "This design modularizes the “Clifford‐linear” transformation so that it can be reused inside deeper architectures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import cliffordlayers.cliffordkernels as ck\n",
    "import cliffordlayers.nn.modules.groupnorm as gn\n",
    "\n",
    "\n",
    "\n",
    "def custom_clifford_norm(x, eps=1e-5):\n",
    "    \"\"\"\n",
    "    Custom normalization for 3D multivector tensors.\n",
    "    \n",
    "    Args:\n",
    "        x: A tensor of shape (B, C, 8) where each channel is a 3D multivector.\n",
    "        eps: A small constant for numerical stability.\n",
    "        \n",
    "    Returns:\n",
    "        x_normalized: The normalized tensor of the same shape as x.\n",
    "    \"\"\"\n",
    "    # Compute mean of the scalar part (first blade) across the batch (dim 0)\n",
    "    # Using x[..., 0:1] preserves the last dimension (shape becomes (1, C, 1))\n",
    "    m = torch.mean(x[..., 0:1], dim=0, keepdims=True)\n",
    "    \n",
    "    # Shift the scalar component: subtract the computed mean\n",
    "    x_shifted_0 = x[..., 0:1] - m  # shape (B, C, 1)\n",
    "    \n",
    "    # Concatenate the shifted scalar with the remaining blades (unchanged)\n",
    "    x_shifted = torch.cat([x_shifted_0, x[..., 1:]], dim=-1)  # shape (B, C, 8)\n",
    "    \n",
    "    # Compute variance across the batch (dim 0) and blade dimension (dim -1)\n",
    "    var = torch.mean(x_shifted**2, dim=[0, 2], keepdims=True)  # shape (1, C, 1)\n",
    "    \n",
    "    # Normalize: divide shifted tensor by the square root of (var + eps)\n",
    "    x_normalized = x_shifted / torch.sqrt(var + eps)\n",
    "    \n",
    "    return x_normalized\n",
    "\n",
    "\n",
    "class MultivectorLinear(nn.Module):\n",
    "    \"\"\"\n",
    "    A linear layer for 3D multivectors.\n",
    "    \n",
    "    Input shape:  (B, C_in, 8)\n",
    "    Output shape: (B, C_out, 8)\n",
    "    \n",
    "    Internally:\n",
    "      1. Transposes the input to (B, 8, C_in) and flattens it to (B, 8 * C_in).\n",
    "      2. Obtains a Clifford kernel from a trainable weight parameter self.W,\n",
    "         via ck.get_3d_clifford_kernel(self.W, g), where self.W has shape (8, C_out, C_in).\n",
    "         This produces a kernel of shape (8, 8, C_out, C_in) that is reshaped to (8 * C_out, 8 * C_in).\n",
    "      3. Multiplies the flattened input by the transposed kernel to yield (B, 8 * C_out).\n",
    "      4. Reshapes the result to (B, 8, C_out) and transposes to (B, C_out, 8).\n",
    "    \"\"\"\n",
    "    def __init__(self, C_in, C_out, g=(1,1,1)):\n",
    "        super().__init__()\n",
    "        self.C_in = C_in\n",
    "        self.C_out = C_out\n",
    "        self.g = g\n",
    "        # Trainable weight parameter for the linear layer.\n",
    "        # Its shape is (8, C_out, C_in); ck.get_3d_clifford_kernel() will expand this.\n",
    "        self.W = nn.Parameter(torch.randn(8, C_out, C_in))\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        x: (B, C_in, 8)\n",
    "        Returns: (B, C_out, 8)\n",
    "        \"\"\"\n",
    "        B = x.shape[0]\n",
    "        # 1) Transpose to (B, 8, C_in)\n",
    "        x = x.transpose(1, 2)  # now (B, 8, C_in)\n",
    "        # Flatten x to (B, 8 * C_in)\n",
    "        x_flat = x.reshape(B, -1)\n",
    "        \n",
    "        # 2) Obtain Clifford kernel\n",
    "        # get_3d_clifford_kernel returns a kernel of shape (8, 8, C_out, C_in)\n",
    "        _, kernel = ck.get_3d_clifford_kernel(self.W, g=self.g)\n",
    "        # Permute to (8, C_out, 8, C_in) then reshape to (8 * C_out, 8 * C_in)\n",
    "        # kernel = kernel.permute(0,2,1,3).reshape(8 * self.C_out, 8 * self.C_in)\n",
    "        \n",
    "        # 3) Multiply: map (B, 8*C_in) -> (B, 8*C_out)\n",
    "        out_flat = x_flat @ kernel.T\n",
    "        \n",
    "        # 4) Reshape to (B, 8, C_out) and then transpose to (B, C_out, 8)\n",
    "        out = out_flat.reshape(B, 8, self.C_out).transpose(1, 2)\n",
    "        return out\n",
    "\n",
    "\n",
    "class MultivectorMLP(nn.Module):\n",
    "    \"\"\"\n",
    "    A Multivector MLP that regresses from an input of shape (B, C_in, 8)\n",
    "    to an output of shape (B, C_out, 8) using Clifford-based linear layers,\n",
    "    group normalization, and activation.\n",
    "    \"\"\"\n",
    "    def __init__(self, C_in, C_out, g=torch.tensor([1,1,1])):\n",
    "        super().__init__()\n",
    "        self.linear1 = MultivectorLinear(C_in, C_out, g=g)\n",
    "        # Apply Clifford group normalization on output channels (B, C_out, 8)\n",
    "        # self.norm = gn.CliffordGroupNorm3d(g=g, num_groups=1, channels=C_out)\n",
    "        self.norm = custom_clifford_norm\n",
    "        self.activation = F.gelu\n",
    "        # A second linear layer for additional transformation\n",
    "        self.linear2 = MultivectorLinear(C_out, C_out, g=g)\n",
    "        self.linear3 = MultivectorLinear(C_out, C_out, g=g)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x: (B, C_in, 8)\n",
    "        x = self.linear1(x)          # => (B, C_out, 8)\n",
    "        x = self.norm(x)             # Normalize in the Clifford domain\n",
    "        x = self.activation(x)       # Apply activation\n",
    "        x = self.linear2(x)          # => (B, C_out, 8)\n",
    "        x = self.norm(x)             # Normalize in the Clifford domain\n",
    "        x = self.activation(x)       # Apply activation\n",
    "        x = self.linear3(x)          # => (B, C_out, 8)\n",
    "        return x\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = 2       # Batch size\n",
    "C_in = 4    # Input channels (each channel is a 3D multivector)\n",
    "C_out = 6   # Output channels\n",
    "x = torch.randn(B, C_in, 8)  # Random input tensor of shape (B, C_in, 8)\n",
    "\n",
    "model = MultivectorMLP(C_in, C_out, g=torch.tensor([1,1,1]))\n",
    "y = model(x)\n",
    "print(\"Input shape:\", x.shape)\n",
    "print(\"Output shape:\", y.shape)  # Expected: (B, C_out, 8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our goal is to build a neural network that processes a 2D image of a colored cube (from a synthetic dataset) and directly regresses the camera’s pose—that is, its rotation (expressed as a quaternion) and translation.\n",
    "\n",
    "### Architecture Overview\n",
    "\n",
    "1. **Input Image:**  \n",
    "   - The network takes a 2D image of shape $(B, C, H, W)$ (for example, a grayscale or RGB image).\n",
    "\n",
    "2. **Convert to Multivector Field:**  \n",
    "   - To prepare the image for 2D Clifford processing, we expand its last dimension by appending two zero channels, yielding a tensor of shape $(B, C, H, W, 3)$.  \n",
    "   - This “multivector field” has 3 components per spatial location, which is required by the 2D Clifford blocks.\n",
    "\n",
    "3. **CliffordNet2d Processing:**  \n",
    "   - The multivector field is fed into a CliffordNet2d (or a series of CliffordBasicBlock2d modules) to extract geometric features.  \n",
    "   - The output from this stage has shape $(B, C_{\\text{out}}, H, W, 3)$.\n",
    "\n",
    "4. **Spatial Averaging:**  \n",
    "   - A spatial average over the height and width dimensions reduces the tensor to $(B, C_{\\text{out}}, 3)$.\n",
    "\n",
    "5. **Expansion to 3D Multivectors:**  \n",
    "   - Each channel is “lifted” to a full 3D multivector with 8 components.  \n",
    "   - This conversion produces a tensor of shape $(B, C_{\\text{out}}, 8)$.  \n",
    "   - In our implementation, this is achieved by concatenating zeros to the 3-component feature vector.\n",
    "\n",
    "6. **Multivector MLP:**  \n",
    "   - The expanded tensor is then passed through a Multivector MLP, which regresses the final pose.  \n",
    "   - The MLP maps the multivector features to an output tensor of shape $(B, 1, 8)$.  \n",
    "   - (Typically, the number of output channels $C''$ is set to 1 for final regression.)\n",
    "\n",
    "7. **Splitting and Normalization:**  \n",
    "   - The final 8 components are split into two parts:\n",
    "     - **Rotation Quaternion:** Indices $[0, 4, 5, 6]$ form the quaternion, which is then normalized to ensure it is a unit quaternion.\n",
    "     - **Translation:** Indices $[1, 2, 3]$ form the translation vector.\n",
    "   - The model returns these as separate tensors, with shapes $(B, 4)$ for the quaternion and $(B, 3)$ for the translation.\n",
    "\n",
    "#### Summary\n",
    "\n",
    "This end-to-end model leverages CliffordNet2d for feature extraction and a dedicated Multivector MLP for regression, providing a fully differentiable pipeline that directly estimates the camera pose from a single 2D image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the 2D Clifford network module (assumed to be available)\n",
    "import cliffordlayers.models.models_2d as models_2d\n",
    "\n",
    "\n",
    "class CameraPoseCliffordModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Camera Pose Regression using CliffordNet2d and a Multivector MLP.\n",
    "    \n",
    "    Input: A 2D image of shape (B, C, H, W)\n",
    "    Output: A rotation quaternion (B, 4) and a translation vector (B, 3)\n",
    "    \n",
    "    Architecture Steps:\n",
    "      1. Convert the image to a multivector field: (B, C, H, W) -> (B, C, H, W, 3)\n",
    "      2. Process via CliffordNet2d to obtain (B, C_out, H, W, 3)\n",
    "      3. Spatial average over (H, W) to get (B, C_out, 3)\n",
    "      4. Convert to a full 3D multivector representation: (B, C_out, 3) -> (B, C_out, 8)\n",
    "      5. Pass through a Multivector MLP to regress to (B, 1, 8)\n",
    "      6. Split the 8 components into quaternion and translation.\n",
    "    \"\"\"\n",
    "    def __init__(self, in_channels=1, hidden_channels=32, mlp_out_channels=1, num_blocks = [2, 2, 2, 2]):\n",
    "        super().__init__()\n",
    "        # Instantiate CliffordNet2d from the library\n",
    "        # Here, we use a typical configuration; adjust parameters as needed.\n",
    "        self.clifford_net = models_2d.CliffordNet2d(\n",
    "            g = [1, 1],\n",
    "            block = models_2d.CliffordBasicBlock2d,\n",
    "            num_blocks = num_blocks,\n",
    "            in_channels = in_channels,\n",
    "            out_channels = hidden_channels,\n",
    "            hidden_channels = hidden_channels,\n",
    "            activation = F.gelu,\n",
    "            norm = True,\n",
    "            rotation = False,\n",
    "        )\n",
    "        self.hidden_channels = hidden_channels\n",
    "        \n",
    "        \n",
    "        \n",
    "        # Instantiate the Multivector MLP to map from (B, hidden_channels, 8) to (B, mlp_out_channels, 8)\n",
    "        self.mlp = MultivectorMLP(C_in=hidden_channels, C_out=mlp_out_channels)\n",
    "\n",
    "    def expand_to_3d(self, x):\n",
    "        out = torch.stack(\n",
    "            [\n",
    "                x[..., 0],\n",
    "                x[..., 1],\n",
    "                x[..., 2],\n",
    "                torch.zeros_like(x[..., 0]),\n",
    "                torch.zeros_like(x[..., 0]),\n",
    "                torch.zeros_like(x[..., 0]),\n",
    "                torch.zeros_like(x[..., 0]),\n",
    "                torch.zeros_like(x[..., 0]),\n",
    "            ], \n",
    "            dim = -1\n",
    "        )\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: Input image tensor of shape (B, C, H, W)\n",
    "        \n",
    "        Returns:\n",
    "            q: Quaternion tensor of shape (B, 4)\n",
    "            t: Translation tensor of shape (B, 3)\n",
    "        \"\"\"\n",
    "        B, C, H, W = x.shape\n",
    "        \n",
    "        # 1. Convert image to multivector field: (B, C, H, W) -> (B, C, H, W, 3)\n",
    "        # Assume original image channels are scalar; we append two zeros.\n",
    "        x = x.unsqueeze(-1)  # (B, C, H, W, 1)\n",
    "        zeros = torch.zeros_like(x)  # (B, C, H, W, 1)\n",
    "        x = torch.cat([x, zeros, zeros], dim=-1)  # (B, C, H, W, 3)\n",
    "        # 2. Pass through CliffordNet2d.\n",
    "        x = self.clifford_net(x)      # Expected output: (B, hidden_channels, 3, H, W)\n",
    "        # 3. Spatial average over (H, W) => (B, hidden_channels, 3)\n",
    "        x = x.mean(dim=[2, 3])\n",
    "\n",
    "        \n",
    "        \n",
    "        # 4. Expand features from 3 to 8 components per channel: (B, hidden_channels, 3) -> (B, hidden_channels, 8)\n",
    "        x = self.expand_to_3d(x)  # Linear transformation applied to last dimension\n",
    "\n",
    "\n",
    "        \n",
    "        \n",
    "        # 5. Pass through the Multivector MLP => (B, mlp_out_channels, 8)\n",
    "        x = self.mlp(x)\n",
    "\n",
    "        \n",
    "        # 6. At this point, we expect x to have shape (B, 1, 8).\n",
    "        # Split the 8 components: indices [0, 4, 5, 6] for quaternion, [1, 2, 3] for translation.\n",
    "        q = x[..., [0, 4, 5, 6]]  # shape: (B, 1, 4)\n",
    "        t = x[..., [1, 2, 3]]      # shape: (B, 1, 3)\n",
    "        \n",
    "        # Normalize the quaternion to ensure it is a unit quaternion.\n",
    "        q = q / torch.norm(q, dim=-1, keepdim=True)\n",
    "        \n",
    "        # Squeeze the channel dimension: final shapes (B, 4) and (B, 3)\n",
    "        q = q.squeeze(1)\n",
    "        t = t.squeeze(1)\n",
    "        \n",
    "        return q, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a random grayscale image (B, C, H, W)\n",
    "B, C, H, W = 2, 1, 64, 64\n",
    "x = torch.randn(B, C, H, W)\n",
    "\n",
    "model = CameraPoseCliffordModel(in_channels=C, hidden_channels=32, mlp_out_channels=1)\n",
    "q, t = model(x)\n",
    "\n",
    "print(\"Input image shape:\", x.shape)\n",
    "print(\"Predicted Quaternion shape:\", q.shape)  # Expected: (B, 4)\n",
    "print(\"Predicted Translation shape:\", t.shape)  # Expected: (B, 3)\n",
    "print(\"Sample Quaternion:\\n\", q[0])\n",
    "print(\"Sample Translation:\\n\", t[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## End-to-End Pose Regression Pipeline and Comparison\n",
    "\n",
    "In this final step, we combine all the snippets and methodology presented above into a complete **pose regression pipeline**. Specifically:\n",
    "\n",
    "1. **Dataset Generation & Splitting:**\n",
    "   - We generate or load approximately 1000 synthetic cube‐view images, each labeled with a ground‐truth camera pose (rotation quaternion + translation).\n",
    "   - The dataset is then split into training (70%), validation (15%), and test (15%) sets.\n",
    "\n",
    "2. **Two Models:**\n",
    "   1. **CameraPoseCliffordModel** (GA‐Based):  \n",
    "      - Takes a 2D image, converts it into a Clifford multivector field, processes it via CliffordNet2d, performs spatial averaging, lifts features to full 3D multivectors, and regresses the final pose with a Multivector MLP.\n",
    "   2. **ResNet‐Based Model** (No GA):  \n",
    "      - A CNN (e.g., ResNet style) that outputs 7 features. The first 4 features form the rotation quaternion (after normalization), and the remaining 3 represent translation.\n",
    "\n",
    "3. **Training Loop & Early Stopping:**\n",
    "   - We train both models on the training set using MSE loss, monitor validation loss for early stopping, and finally evaluate on the test set.\n",
    "\n",
    "4. **Comparison:**\n",
    "   - The test set MSE (or any chosen metric) for both models is reported, providing a direct comparison between the Clifford GA approach and a standard CNN baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "import torchvision.models as models\n",
    "\n",
    "\n",
    "class CubePoseDataset(Dataset):\n",
    "    def __init__(self, data_list):\n",
    "        self.data_list = data_list\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data_list)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data_list[idx]\n",
    "    \n",
    "\n",
    "class ResNetPoseModel(nn.Module):\n",
    "    \"\"\"\n",
    "    A pose regression model using a torchvision ResNet backbone.\n",
    "    Outputs 7 features:\n",
    "      - First 4 => rotation quaternion (normalized)\n",
    "      - Last 3 => translation\n",
    "    \"\"\"\n",
    "    def __init__(self, out_features=7, pretrained=False):\n",
    "        super().__init__()\n",
    "        # Load a ResNet-18 (you can choose resnet50, etc.) from torchvision\n",
    "        # weights=None if pretrained=False, or you can set pretrained weights\n",
    "        self.resnet = models.resnet18(weights=None if not pretrained else 'IMAGENET1K_V1')\n",
    "        \n",
    "        # Modify the final fully-connected layer to produce 'out_features' (7)\n",
    "        in_features = self.resnet.fc.in_features\n",
    "        self.resnet.fc = nn.Linear(in_features, out_features)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: Input image of shape (B, C, H, W).\n",
    "        Returns:\n",
    "            q: Quaternion of shape (B, 4), normalized\n",
    "            t: Translation of shape (B, 3)\n",
    "        \"\"\"\n",
    "        # Forward pass through ResNet => (B, 7)\n",
    "        out = self.resnet(x)\n",
    "        \n",
    "        # Split the first 4 features for quaternion, last 3 for translation\n",
    "        q = out[:, :4]   # (B,4)\n",
    "        t = out[:, 4:]   # (B,3)\n",
    "        \n",
    "        # Normalize the quaternion\n",
    "        q_norm = q / torch.norm(q, dim=1, keepdim=True)\n",
    "        \n",
    "        return q_norm, t\n",
    "\n",
    "\n",
    "\n",
    "def train_epoch(model, loader, optimizer, criterion, device):\n",
    "    model.train()\n",
    "    total_loss = 0.0\n",
    "    for batch in loader:\n",
    "        images = batch['image']\n",
    "        cam_pos = batch['camera_pos'].float()\n",
    "        q = batch['quaternion'].float()\n",
    "        images, cam_pos, q = images.to(device), cam_pos.to(device), q.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        pred_q, pred_t = model(images)\n",
    "        loss = criterion(pred_q, q) + criterion(pred_t, cam_pos)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item() * images.size(0)\n",
    "        print(\"Loss:\", loss.item())\n",
    "    return total_loss / len(loader.dataset)\n",
    "\n",
    "@torch.no_grad()\n",
    "def eval_epoch(model, loader, criterion, device):\n",
    "    model.eval()\n",
    "    total_loss = 0.0\n",
    "    for batch in loader:\n",
    "        images = batch['image']\n",
    "        cam_pos = batch['camera_pos']\n",
    "        q = batch['quaternion']\n",
    "        images, cam_pos, q = images.to(device), cam_pos.to(device), q.to(device)\n",
    "        pred_q, pred_t = model(images)\n",
    "        loss = criterion(pred_q, q) + criterion(pred_t, cam_pos)\n",
    "        total_loss += loss.item() * images.size(0)\n",
    "    return total_loss / len(loader.dataset)\n",
    "\n",
    "class EarlyStopper:\n",
    "    def __init__(self, patience=5):\n",
    "        self.patience = patience\n",
    "        self.best_loss = float('inf')\n",
    "        self.counter = 0\n",
    "        self.stop = False\n",
    "\n",
    "    def step(self, val_loss):\n",
    "        if val_loss < self.best_loss:\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "        else:\n",
    "            self.counter += 1\n",
    "            if self.counter >= self.patience:\n",
    "                self.stop = True\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1) Generate the dataset with your custom function:\n",
    "W, H = 64, 48\n",
    "fl = 70\n",
    "K = np.array([\n",
    "    [fl,   0, W/2],\n",
    "    [  0, fl, H/2],\n",
    "    [  0,   0,   1]\n",
    "])\n",
    "\n",
    "# Generate a dataset with the updated intrinsic matrix\n",
    "data_list = generate_image_dataset_with_depth(cube_faces, num_views=1000, radius=5.0, K=K, H=H, W=W)\n",
    "full_dataset = CubePoseDataset(data_list)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2) Split dataset: train (70%), val (15%), test (15%)\n",
    "total = len(full_dataset)\n",
    "train_size = int(0.7 * total)\n",
    "val_size = int(0.15 * total)\n",
    "test_size = total - train_size - val_size\n",
    "train_ds, val_ds, test_ds = random_split(full_dataset, [train_size, val_size, test_size])\n",
    "\n",
    "train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)\n",
    "val_loader = DataLoader(val_ds, batch_size=32, shuffle=False)\n",
    "test_loader = DataLoader(test_ds, batch_size=32, shuffle=False)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "criterion = nn.MSELoss()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3) Train CameraPoseCliffordModel\n",
    "model_clifford = CameraPoseCliffordModel(in_channels=3, hidden_channels=8, mlp_out_channels=1).to(device)\n",
    "opt_clifford = optim.Adam(model_clifford.parameters(), lr=1e-3)\n",
    "stopper_clifford = EarlyStopper(patience=5)\n",
    "num_epochs = 50\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss = train_epoch(model_clifford, train_loader, opt_clifford, criterion, device)\n",
    "    val_loss = eval_epoch(model_clifford, val_loader, criterion, device)\n",
    "    print(f\"[Clifford] Epoch {epoch+1}/{num_epochs} | Train: {train_loss:.4f}, Val: {val_loss:.4f}\")\n",
    "    stopper_clifford.step(val_loss)\n",
    "    if stopper_clifford.stop:\n",
    "        print(\"Early stopping triggered for Clifford model.\")\n",
    "        break\n",
    "\n",
    "test_loss_clifford = eval_epoch(model_clifford, test_loader, criterion, device)\n",
    "print(\"CameraPoseCliffordModel Test Loss:\", test_loss_clifford)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4) Train ResNet-based model\n",
    "model_resnet = ResNetPoseModel(in_channels=1, num_features=7).to(device)\n",
    "opt_resnet = optim.Adam(model_resnet.parameters(), lr=1e-3)\n",
    "stopper_resnet = EarlyStopper(patience=5)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss = train_epoch(model_resnet, train_loader, opt_resnet, criterion, device)\n",
    "    val_loss = eval_epoch(model_resnet, val_loader, criterion, device)\n",
    "    print(f\"[ResNet] Epoch {epoch+1}/{num_epochs} | Train: {train_loss:.4f}, Val: {val_loss:.4f}\")\n",
    "    stopper_resnet.step(val_loss)\n",
    "    if stopper_resnet.stop:\n",
    "        print(\"Early stopping triggered for ResNet baseline.\")\n",
    "        break\n",
    "\n",
    "test_loss_resnet = eval_epoch(model_resnet, test_loader, criterion, device)\n",
    "print(\"ResNetPoseModel Test Loss:\", test_loss_resnet)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5) Compare final performance\n",
    "print(\"\\nPerformance Comparison:\")\n",
    "print(f\"CameraPoseCliffordModel: {test_loss_clifford:.4f}\")\n",
    "print(f\"ResNetPoseModel       : {test_loss_resnet:.4f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3.11.10",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
