Repository: kornia/kornia

Issue:
add `Quaternion` api
## 🚀 Feature

Add `Quaternion` class that is compatible with `Eigen::Quaternion`
https://eigen.tuxfamily.org/dox/classEigen_1_1Quaternion.html

Initial implementation here: https://github.com/kornia/kornia/blob/liegroup-algebra/kornia/geometry/liegroup/quaternion.py

## Motivation

We want a simple api that can be easily combined with `kornia.geometry.conversions` and the upcoming `kornia.geometry.liegroup`

TBD: where this api should live

______________________________________________________________________

#### Consider also to contribute to Kornia universe projects :)

<sub>

- [**Tutorials**](https://github.com/kornia/tutorials): our repository containing the tutorials.

</sub>



Discussion:

Resolving patch:
diff --git a/kornia/core/__init__.py b/kornia/core/__init__.py
--- a/kornia/core/__init__.py
+++ b/kornia/core/__init__.py
@@ -1,3 +1,3 @@
-from ._backend import Module, Parameter, Tensor, concatenate, normalize, stack
+from ._backend import Module, Parameter, Tensor, as_tensor, concatenate, normalize, rand, stack
 
-__all__ = ["concatenate", "Module", "Tensor", "Parameter", "normalize", "stack"]
+__all__ = ["concatenate", "Module", "Tensor", "Parameter", "normalize", "stack", "as_tensor", "rand"]
diff --git a/kornia/core/_backend.py b/kornia/core/_backend.py
--- a/kornia/core/_backend.py
+++ b/kornia/core/_backend.py
@@ -10,3 +10,9 @@
 concatenate = torch.cat
 stack = torch.stack
 normalize = F.normalize
+
+# constructors
+as_tensor = torch.as_tensor
+
+# random
+rand = torch.rand
diff --git a/kornia/geometry/quaternion.py b/kornia/geometry/quaternion.py
new file mode 100644
--- /dev/null
+++ b/kornia/geometry/quaternion.py
@@ -0,0 +1,327 @@
+# kornia.geometry.quaternion module inspired by Eigen, Sophus-sympy, and PyQuaternion.
+# https://github.com/strasdat/Sophus/blob/master/sympy/sophus/quaternion.py
+# https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py
+# https://gitlab.com/libeigen/eigen/-/blob/master/Eigen/src/Geometry/Quaternion.h
+from math import pi
+from typing import Tuple, Union
+
+from kornia.core import Module, Parameter, Tensor, as_tensor, concatenate, rand, stack
+from kornia.geometry.conversions import (
+    QuaternionCoeffOrder,
+    normalize_quaternion,
+    quaternion_to_rotation_matrix,
+    rotation_matrix_to_quaternion,
+)
+from kornia.testing import KORNIA_CHECK, KORNIA_CHECK_SHAPE, KORNIA_CHECK_TYPE
+
+
+class Quaternion(Module):
+    r"""Base class to represent a Quaternion.
+
+    A quaternion is a four dimensional vector representation of a rotation transformation in 3d.
+    See more: https://en.wikipedia.org/wiki/Quaternion
+
+    The general definition of a quaternion is given by:
+
+    .. math ::
+
+        Q = a + b \cdot \mathbf{i} + c \cdot \mathbf{j} + d \cdot \mathbf{k}
+
+    Thus, we represent a rotation quaternion as a contiguous tensor structure to
+    perform rigid bodies transformations:
+
+    .. math ::
+
+        Q = \begin{bmatrix} q_w & q_x & q_y & q_z \end{bmatrix}
+
+    Example:
+        >>> q = Quaternion.identity(batch_size=4)
+        >>> q.data
+        Parameter containing:
+        tensor([[1., 0., 0., 0.],
+                [1., 0., 0., 0.],
+                [1., 0., 0., 0.],
+                [1., 0., 0., 0.]], requires_grad=True)
+        >>> q.real
+        tensor([[1.],
+                [1.],
+                [1.],
+                [1.]], grad_fn=<SliceBackward0>)
+        >>> q.vec
+        tensor([[0., 0., 0.],
+                [0., 0., 0.],
+                [0., 0., 0.],
+                [0., 0., 0.]], grad_fn=<SliceBackward0>)
+    """
+
+    def __init__(self, data: Tensor) -> None:
+        """Constructor for the base class.
+
+        Args:
+            data: tensor containing the quaternion data with the sape of :math:`(B, 4)`.
+
+        Example:
+            >>> data = torch.rand(2, 4)
+            >>> q = Quaternion(data)
+            >>> q.shape
+            (2, 4)
+        """
+        super().__init__()
+        KORNIA_CHECK_SHAPE(data, ["B", "4"])
+        self._data = Parameter(data)
+
+    def __repr__(self) -> str:
+        return f"real: {self.real} \nvec: {self.vec}"
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+    def __neg__(self) -> 'Quaternion':
+        """Inverts the sign of the quaternion data.
+
+        Example:
+            >>> q = Quaternion.identity(batch_size=1)
+            >>> -q.data
+            tensor([[-1., -0., -0., -0.]], grad_fn=<NegBackward0>)
+        """
+        return Quaternion(-self.data)
+
+    def __add__(self, right: 'Quaternion') -> 'Quaternion':
+        """Add a given quaternion.
+
+        Args:
+            right: the quaternion to add.
+
+        Example:
+            >>> q1 = Quaternion.identity(batch_size=1)
+            >>> q2 = Quaternion(Tensor([[2., 0., 1., 1.]]))
+            >>> q3 = q1 + q2
+            >>> q3.data
+            Parameter containing:
+            tensor([[3., 0., 1., 1.]], requires_grad=True)
+        """
+        KORNIA_CHECK_TYPE(right, Quaternion)
+        return Quaternion(self.data + right.data)
+
+    def __sub__(self, right: 'Quaternion') -> 'Quaternion':
+        """Subtract a given quaternion.
+
+        Args:
+            right: the quaternion to subtract.
+
+        Example:
+            >>> q1 = Quaternion(Tensor([[2., 0., 1., 1.]]))
+            >>> q2 = Quaternion.identity(batch_size=1)
+            >>> q3 = q1 - q2
+            >>> q3.data
+            Parameter containing:
+            tensor([[1., 0., 1., 1.]], requires_grad=True)
+        """
+        KORNIA_CHECK_TYPE(right, Quaternion)
+        return Quaternion(self.data - right.data)
+
+    def __mul__(self, right: 'Quaternion') -> 'Quaternion':
+        KORNIA_CHECK_TYPE(right, Quaternion)
+        # NOTE: borrowed from sophus sympy. Produce less multiplications compared to others.
+        # https://github.com/strasdat/Sophus/blob/785fef35b7d9e0fc67b4964a69124277b7434a44/sympy/sophus/quaternion.py#L19
+        new_real = self.real * right.real - self._batched_squared_norm(self.vec, right.vec)
+        new_vec = self.real * right.vec + right.real * self.vec + self.vec.cross(right.vec)
+        return Quaternion(concatenate((new_real, new_vec), -1))
+
+    def __div__(self, right: Union[Tensor, 'Quaternion']) -> 'Quaternion':
+        if isinstance(right, Tensor):
+            return Quaternion(self.data / right)
+        KORNIA_CHECK_TYPE(right, Quaternion)
+        return self * right.inv()
+
+    def __truediv__(self, right: 'Quaternion') -> 'Quaternion':
+        return self.__div__(right)
+
+    @property
+    def data(self) -> Tensor:
+        """Return the underlying data with shape :math:`(B,4).`"""
+        return self._data
+
+    @property
+    def coeffs(self) -> Tensor:
+        """Return the underlying data with shape :math:`(B,4)`.
+
+        Alias for :func:`~kornia.geometry.quaternion.Quaternion.data`
+        """
+        return self._data
+
+    @property
+    def real(self) -> Tensor:
+        """Return the real part with shape :math:`(B,1)`.
+
+        Alias for :func:`~kornia.geometry.quaternion.Quaternion.w`
+        """
+        return self.w
+
+    @property
+    def vec(self) -> Tensor:
+        """Return the vector with the imaginary part with shape :math:`(B,3)`."""
+        return self.data[..., 1:]
+
+    @property
+    def q(self) -> Tensor:
+        """Return the underlying data with shape :math:`(B,4)`.
+
+        Alias for :func:`~kornia.geometry.quaternion.Quaternion.data`
+        """
+        return self.data
+
+    @property
+    def scalar(self) -> Tensor:
+        """Return a scalar with the real with shape :math:`(B,1)`.
+
+        Alias for :func:`~kornia.geometry.quaternion.Quaternion.w`
+        """
+        return self.real
+
+    @property
+    def w(self) -> Tensor:
+        """Return the :math:`q_w` with shape :math:`(B,1)`."""
+        return self.data[..., 0:1]
+
+    @property
+    def x(self) -> Tensor:
+        """Return the :math:`q_x` with shape :math:`(B,1)`."""
+        return self.data[..., 1:2]
+
+    @property
+    def y(self) -> Tensor:
+        """Return the :math:`q_y` with shape :math:`(B,1)`."""
+        return self.data[..., 2:3]
+
+    @property
+    def z(self) -> Tensor:
+        """Return the :math:`q_z` with shape :math:`(B,1)`."""
+        return self.data[..., 3:4]
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        """Return the shape of the underlying data with shape :math:`(B,4)`."""
+        return tuple(self.data.shape)
+
+    @property
+    def polar_angle(self) -> Tensor:
+        """Return the polar angle with shape :math:`(B,1)`.
+
+        Example:
+            >>> q = Quaternion.identity(batch_size=1)
+            >>> q.polar_angle
+            tensor([[0.]], grad_fn=<AcosBackward0>)
+        """
+        return (self.scalar / self.norm()).acos()
+
+    def matrix(self) -> Tensor:
+        """Convert the quaternion to a rotation matrix of shape :math:`(B,3,3)`.
+
+        Example:
+            >>> q = Quaternion.identity(batch_size=1)
+            >>> m = q.matrix()
+            >>> m
+            tensor([[[1., 0., 0.],
+                     [0., 1., 0.],
+                     [0., 0., 1.]]], grad_fn=<ViewBackward0>)
+        """
+        return quaternion_to_rotation_matrix(self.data, order=QuaternionCoeffOrder.WXYZ)
+
+    @classmethod
+    def from_matrix(cls, matrix: Tensor) -> 'Quaternion':
+        """Create a quaternion from a rotation matrix.
+
+        Args:
+            matrix: the rotation matrix to convert of shape :math:`(B,3,3)`.
+
+        Example:
+            >>> m = torch.eye(3)[None]
+            >>> q = Quaternion.from_matrix(m)
+            >>> q.data
+            Parameter containing:
+            tensor([[1., 0., 0., 0.]], requires_grad=True)
+        """
+        return cls(rotation_matrix_to_quaternion(matrix, order=QuaternionCoeffOrder.WXYZ))
+
+    @classmethod
+    def identity(cls, batch_size: int) -> 'Quaternion':
+        """Create a quaternion representing an identity rotation.
+
+        Args:
+            batch_size: the batch size of the underlying data.
+
+        Example:
+            >>> q = Quaternion.identity(batch_size=2)
+            >>> q.data
+            Parameter containing:
+            tensor([[1., 0., 0., 0.],
+                    [1., 0., 0., 0.]], requires_grad=True)
+        """
+        data: Tensor = as_tensor([1.0, 0.0, 0.0, 0.0])
+        data = data.repeat(batch_size, 1)
+        return cls(data)
+
+    @classmethod
+    def from_coeffs(cls, w: float, x: float, y: float, z: float) -> 'Quaternion':
+        """Create a quaternion from the data coefficients.
+
+        Args:
+            w: a float representing the :math:`q_w` component.
+            x: a float representing the :math:`q_x` component.
+            y: a float representing the :math:`q_y` component.
+            z: a float representing the :math:`q_z` component.
+
+        Example:
+            >>> q = Quaternion.from_coeffs(1., 0., 0., 0.)
+            >>> q.data
+            Parameter containing:
+            tensor([[1., 0., 0., 0.]], requires_grad=True)
+        """
+        return cls(as_tensor([[w, x, y, z]]))
+
+    @classmethod
+    def random(cls, batch_size: int) -> 'Quaternion':
+        """Create a random unit quaternion of shape :math:`(B,4)`.
+
+        Uniformly distributed across the rotation space as per: http://planning.cs.uiuc.edu/node198.html
+
+        Args:
+            batch_size: the batch size of the underlying data.
+
+        Example:
+            >>> q = Quaternion.random(batch_size=2)
+            >>> q.norm()
+            tensor([1.0000, 1.0000], grad_fn=<NormBackward1>)
+        """
+        r1, r2, r3 = rand(3, batch_size)
+        q1 = (1.0 - r1).sqrt() * ((2 * pi * r2).sin())
+        q2 = (1.0 - r1).sqrt() * ((2 * pi * r2).cos())
+        q3 = r1.sqrt() * (2 * pi * r3).sin()
+        q4 = r1.sqrt() * (2 * pi * r3).cos()
+        return cls(stack((q1, q2, q3, q4), -1))
+
+    def norm(self) -> Tensor:
+        return self.data.norm(p=2, dim=-1)
+
+    def normalize(self) -> 'Quaternion':
+        return Quaternion(normalize_quaternion(self.data))
+
+    def conj(self) -> 'Quaternion':
+        return Quaternion(concatenate((self.real, -self.vec), -1))
+
+    def inv(self) -> 'Quaternion':
+        return self.conj() / self.squared_norm()
+
+    def squared_norm(self) -> Tensor:
+        return self._batched_squared_norm(self.vec) + self.real**2
+
+    def _batched_squared_norm(self, x, y=None):
+        if y is None:
+            y = x
+        KORNIA_CHECK(x.shape == y.shape)
+        return (x[..., None, :] @ y[..., :, None])[..., 0]
+
+    # TODO: implement me
+    def slerp(self):
+        raise NotImplementedError


Test patch:
diff --git a/kornia/testing/__init__.py b/kornia/testing/__init__.py
--- a/kornia/testing/__init__.py
+++ b/kornia/testing/__init__.py
@@ -214,6 +214,11 @@ def KORNIA_UNWRAP(maybe_obj, typ):
     return cast(typ, maybe_obj)
 
 
+def KORNIA_CHECK_TYPE(x, typ, msg: Optional[str] = None):
+    if not isinstance(x, typ):
+        raise TypeError(f"Invalid type: {type(x)}.\n{msg}")
+
+
 def KORNIA_CHECK_IS_TENSOR(x, msg: Optional[str] = None):
     if not isinstance(x, Tensor):
         raise TypeError(f"Not a Tensor type. Got: {type(x)}.\n{msg}")
diff --git a/test/geometry/test_quaternion.py b/test/geometry/test_quaternion.py
new file mode 100644
--- /dev/null
+++ b/test/geometry/test_quaternion.py
@@ -0,0 +1,180 @@
+import pytest
+import torch
+
+from kornia.geometry.quaternion import Quaternion
+from kornia.testing import assert_close
+
+
+class TestQuaternion:
+    def assert_close(self, actual, expected, rtol=None, atol=None):
+        if isinstance(actual, Quaternion):
+            actual = actual.data.data
+        elif isinstance(actual, torch.nn.Parameter):
+            actual = actual.data
+        if isinstance(expected, Quaternion):
+            expected = expected.data.data
+        elif isinstance(expected, torch.nn.Parameter):
+            expected = expected.data
+
+        assert_close(actual, expected, rtol=rtol, atol=atol)
+
+    def test_smoke(self, device, dtype):
+        q = Quaternion.from_coeffs(1.0, 0.0, 0.0, 0.0)
+        q = q.to(device, dtype)
+        q_data = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=device, dtype=dtype)
+        assert isinstance(q, Quaternion)
+        assert q.shape == (1, 4)
+        self.assert_close(q.data, q_data)
+        self.assert_close(q.q, q_data)
+        self.assert_close(q.real, q_data[..., :1])
+        self.assert_close(q.scalar, q_data[..., :1])
+        self.assert_close(q.vec, q_data[..., 1:])
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_init(self, device, dtype, batch_size):
+        q1 = Quaternion.identity(batch_size)
+        q1 = q1.to(device, dtype)
+        q2 = Quaternion(q1.data)
+        assert isinstance(q2, Quaternion)
+        self.assert_close(q1, q2)
+
+    def test_init_fail(self, device, dtype):
+        with pytest.raises(Exception):
+            _ = Quaternion("q")
+
+        with pytest.raises(Exception):
+            _ = Quaternion([1, 0, 0, 0])
+
+        with pytest.raises(Exception):
+            _ = Quaternion(1, [0, 0, 0])
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_random(self, device, dtype, batch_size):
+        q = Quaternion.random(batch_size)
+        q = q.to(device, dtype)
+        q_n = q.normalize().norm()
+        self.assert_close(q_n, q_n.new_ones(q_n.shape))
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_minus(self, device, dtype, batch_size):
+        data = torch.rand(batch_size, 4, device=device, dtype=dtype)
+        q = Quaternion(data)
+        q = q.to(device, dtype)
+        self.assert_close(-q, -data)
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_add(self, device, dtype, batch_size):
+        d1 = torch.rand(batch_size, 4, device=device, dtype=dtype)
+        d2 = torch.rand(batch_size, 4, device=device, dtype=dtype)
+        q1 = Quaternion(d1).to(device, dtype)
+        q2 = Quaternion(d2).to(device, dtype)
+        q3 = q1 + q2
+        assert isinstance(q3, Quaternion)
+        self.assert_close(q3, d1 + d2)
+        q1 += q2
+        self.assert_close(q1, q3)
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_subtract(self, device, dtype, batch_size):
+        d1 = torch.rand(batch_size, 4, device=device, dtype=dtype)
+        d2 = torch.rand(batch_size, 4, device=device, dtype=dtype)
+        q1 = Quaternion(d1).to(device, dtype)
+        q2 = Quaternion(d2).to(device, dtype)
+        q3 = q1 - q2
+        assert isinstance(q3, Quaternion)
+        self.assert_close(q3, d1 - d2)
+        q1 -= q2
+        self.assert_close(q1, q3)
+
+    def test_multiplication_of_bases(self, device, dtype):
+        one = Quaternion.from_coeffs(1.0, 0.0, 0.0, 0.0).to(device, dtype)
+        i = Quaternion.from_coeffs(0.0, 1.0, 0.0, 0.0).to(device, dtype)
+        j = Quaternion.from_coeffs(0.0, 0.0, 1.0, 0.0).to(device, dtype)
+        k = Quaternion.from_coeffs(0.0, 0.0, 0.0, 1.0).to(device, dtype)
+
+        self.assert_close(i * i, j * j)
+        self.assert_close(j * j, k * k)
+        self.assert_close(k * k, i * j * k)
+        self.assert_close(i * j * k, -one)
+
+        self.assert_close(i * j, k)
+        self.assert_close(i * i, -one)
+        self.assert_close(i * k, -j)
+        self.assert_close(j * i, -k)
+        self.assert_close(j * j, -one)
+        self.assert_close(j * k, i)
+        self.assert_close(k * i, j)
+        self.assert_close(k * j, -i)
+        self.assert_close(k * k, -one)
+        self.assert_close(i * j * k, -one)
+
+    def test_division_of_bases(self, device, dtype):
+        one = Quaternion.from_coeffs(1.0, 0.0, 0.0, 0.0).to(device, dtype)
+        i = Quaternion.from_coeffs(0.0, 1.0, 0.0, 0.0).to(device, dtype)
+        j = Quaternion.from_coeffs(0.0, 0.0, 1.0, 0.0).to(device, dtype)
+        k = Quaternion.from_coeffs(0.0, 0.0, 0.0, 1.0).to(device, dtype)
+
+        self.assert_close(i / i, j / j)
+        self.assert_close(j / j, k / k)
+        self.assert_close(k / k, one)
+        self.assert_close(k / -k, -one)
+
+        self.assert_close(i / j, -k)
+        self.assert_close(i / i, one)
+        self.assert_close(i / k, j)
+        self.assert_close(j / i, k)
+        self.assert_close(j / j, one)
+        self.assert_close(j / k, -i)
+        self.assert_close(k / i, -j)
+        self.assert_close(k / j, i)
+        self.assert_close(k / k, one)
+        self.assert_close(i / -j, k)
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_inverse(self, device, dtype, batch_size):
+        q1 = Quaternion.random(batch_size)
+        q1 = q1.to(device, dtype)
+        q2 = Quaternion.identity(batch_size)
+        q2 = q2.to(device, dtype)
+        self.assert_close(q1 * q1.inv(), q2, rtol=1e-4, atol=1e-4)
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_conjugate(self, device, dtype, batch_size):
+        q1 = Quaternion.random(batch_size)
+        q1 = q1.to(device, dtype)
+        q2 = Quaternion.random(batch_size)
+        q2 = q2.to(device, dtype)
+        self.assert_close((q1 * q2).conj(), q2.conj() * q1.conj())
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_double_conjugate(self, device, dtype, batch_size):
+        q1 = Quaternion.random(batch_size)
+        q1 = q1.to(device, dtype)
+        self.assert_close(q1, q1.conj().conj())
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_norm(self, device, dtype, batch_size):
+        q1 = Quaternion.random(batch_size)
+        q1 = q1.to(device, dtype)
+        q2 = Quaternion.random(batch_size)
+        q2 = q2.to(device, dtype)
+        self.assert_close((q1 * q2).norm(), q1.norm() * q2.norm())
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_normalize(self, device, dtype, batch_size):
+        q1 = Quaternion.random(batch_size)
+        q1 = q1.to(device, dtype)
+        q1_n = q1.normalize().norm()
+        self.assert_close(q1_n, q1_n.new_ones(q1_n.shape))
+
+    @pytest.mark.parametrize("batch_size", (1, 2, 5))
+    def test_matrix(self, device, dtype, batch_size):
+        q1 = Quaternion.random(batch_size)
+        q1 = q1.to(device, dtype)
+        m1 = q1.matrix()
+        q2 = Quaternion.from_matrix(m1)
+        for (qq1, qq2) in zip(q1.data, q2.data):
+            try:
+                self.assert_close(qq1, qq2)
+            except Exception:
+                self.assert_close(qq1, -qq2)