# Source code for kornia.geometry.liegroup.se2

```# kornia.geometry.se2 module inspired by Sophus-sympy.
# https://github.com/strasdat/Sophus/blob/master/sympy/sophus/se2.py
from __future__ import annotations

from kornia.core import (
Device,
Dtype,
Module,
Parameter,
Tensor,
concatenate,
rand,
stack,
tensor,
where,
zeros_like,
)
from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_SAME_DEVICES, KORNIA_CHECK_TYPE
from kornia.geometry.liegroup._utils import check_se2_omega_shape, check_se2_t_shape, check_v_shape
from kornia.geometry.liegroup.so2 import So2
from kornia.geometry.vector import Vector2

def _check_se2_r_t_shape(r: So2, t: Tensor) -> None:
z_shape = r.z.shape
if ((len(z_shape) == 1) and (len(t.shape) == 2)) or ((len(z_shape) == 0) and len(t.shape) == 1):
check_se2_t_shape(t)
else:
raise ValueError(
f"Invalid input, both the inputs should be either batched or unbatched. Got: {r.z.shape} and {t.shape}"
)

[docs]class Se2(Module):
r"""Base class to represent the Se2 group.

The SE(2) is the group of rigid body transformations about the origin of two-dimensional Euclidean
space :math:`R^2` under the operation of composition.
See more:

Example:
>>> so2 = So2.identity(1)
>>> t = torch.ones((1, 2))
>>> se2 = Se2(so2, t)
>>> se2
rotation: Parameter containing:
translation: Parameter containing:
"""

[docs]    def __init__(self, rotation: So2, translation: Vector2 | Tensor) -> None:
"""Constructor for the base class.

Internally represented by a complex number `z` and a translation 2-vector.

Args:
rotation: So2 group encompassing a rotation.
translation: translation vector with the shape of :math:`(B, 2)`.

Example:
>>> so2 = So2.identity(1)
>>> t = torch.ones((1, 2))
>>> se2 = Se2(so2, t)
>>> se2
rotation: Parameter containing:
translation: Parameter containing:
"""
super().__init__()
KORNIA_CHECK_TYPE(rotation, So2)
# TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support
# KORNIA_CHECK_TYPE(translation, (Vector3, Tensor))
if not isinstance(translation, (Vector2, Tensor)):
raise TypeError(f"translation type is {type(translation)}")
self._translation: Vector2 | Parameter
self._rotation: So2 = rotation
if isinstance(translation, Tensor):
_check_se2_r_t_shape(rotation, translation)  # TODO remove
self._translation = Parameter(translation)
else:
self._translation = translation

[docs]    def __repr__(self) -> str:
return f"rotation: {self.r}\ntranslation: {self.t}"

def __getitem__(self, idx: int | slice) -> Se2:
return Se2(self._rotation[idx], self._translation[idx])

def __mul__(self, right: Se2) -> Se2:
...

def __mul__(self, right: Tensor) -> Tensor:
...

[docs]    def __mul__(self, right: Se2 | Tensor) -> Se2 | Tensor:
"""Compose two Se2 transformations.

Args:
right: the other Se2 transformation.

Return:
The resulting Se2 transformation.
"""
so2 = self.so2
t = self.t
if isinstance(right, Se2):
KORNIA_CHECK_TYPE(right, Se2)
_r = so2 * right.so2
_t = t + so2 * right.t
return Se2(_r, _t)
elif isinstance(right, (Vector2, Tensor)):
# TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support
# _check_se2_r_t_shape(so2, risght)
return so2 * right + t
else:
raise TypeError(f"Unsupported type: {type(right)}")

@property
def so2(self) -> So2:
"""Return the underlying rotation(So2)."""
return self._rotation

@property
def r(self) -> So2:
"""Return the underlying rotation(So2)."""
return self._rotation

@property
def t(self) -> Vector2 | Parameter:
"""Return the underlying translation vector of shape :math:`(B,2)`."""
return self._translation

@property
def rotation(self) -> So2:
"""Return the underlying rotation(So2)."""
return self._rotation

@property
def translation(self) -> Vector2 | Parameter:
"""Return the underlying translation vector of shape :math:`(B,2)`."""
return self._translation

[docs]    @staticmethod
def exp(v: Tensor) -> Se2:
"""Converts elements of lie algebra to elements of lie group.

Args:
v: vector of shape :math:`(B, 3)`.

Example:
>>> v = torch.ones((1, 3))
>>> s = Se2.exp(v)
>>> s.r
Parameter containing:
>>> s.t
Parameter containing:
"""
check_v_shape(v)
theta = v[..., 2]
so2 = So2.exp(theta)
z = tensor(0.0, device=v.device, dtype=v.dtype)
theta_nonzeros = theta != 0.0
a = where(theta_nonzeros, so2.z.imag / theta, z)
b = where(theta_nonzeros, (1.0 - so2.z.real) / theta, z)
x = v[..., 0]
y = v[..., 1]
t = stack((a * x - b * y, b * x + a * y), -1)
return Se2(so2, t)

[docs]    def log(self) -> Tensor:
"""Converts elements of lie group  to elements of lie algebra.

Example:
>>> v = torch.ones((1, 3))
>>> s = Se2.exp(v).log()
>>> s
"""
theta = self.so2.log()
half_theta = 0.5 * theta
denom = self.so2.z.real - 1
a = where(
denom != 0, -(half_theta * self.so2.z.imag) / denom, tensor(0.0, device=theta.device, dtype=theta.dtype)
)
row0 = stack((a, half_theta), -1)
row1 = stack((-half_theta, a), -1)
V_inv = stack((row0, row1), -2)
upsilon = V_inv @ self.t.data[..., None]
return stack((upsilon[..., 0, 0], upsilon[..., 1, 0], theta), -1)

[docs]    @staticmethod
def hat(v: Tensor) -> Tensor:
"""Converts elements from vector space to lie algebra. Returns matrix of shape :math:`(B, 3, 3)`.

Args:
v: vector of shape:math:`(B, 3)`.

Example:
>>> theta = torch.tensor(3.1415/2)
>>> So2.hat(theta)
tensor([[0.0000, 1.5707],
[1.5707, 0.0000]])
"""
# TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support
check_v_shape(v)
upsilon = stack((v[..., 0], v[..., 1]), -1)
theta = v[..., 2]
col0 = concatenate((So2.hat(theta), upsilon.unsqueeze(-2)), -2)

[docs]    @staticmethod
def vee(omega: Tensor) -> Tensor:
"""Converts elements from lie algebra to vector space.

Args:
omega: 3x3-matrix representing lie algebra of shape :math:`(B, 3, 3)`.

Returns:
vector of shape :math:`(B, 3)`.

Example:
>>> v = torch.ones(3)
>>> omega_hat = Se2.hat(v)
>>> Se2.vee(omega_hat)
tensor([1., 1., 1.])
"""
# TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support
check_se2_omega_shape(omega)
upsilon = omega[..., 2, :2]
theta = So2.vee(omega[..., :2, :2])
return concatenate((upsilon, theta[..., None]), -1)

[docs]    @classmethod
def identity(cls, batch_size: int | None = None, device: Device | None = None, dtype: Dtype = None) -> Se2:
"""Create a Se2 group representing an identity rotation and zero translation.

Args:
batch_size: the batch size of the underlying data.

Example:
>>> s = Se2.identity(1)
>>> s.r
Parameter containing:
>>> s.t
x: tensor([0.])
y: tensor([0.])
"""
t: Tensor = tensor([0.0, 0.0], device=device, dtype=dtype)
if batch_size is not None:
KORNIA_CHECK(batch_size >= 1, msg="batch_size must be positive")
t = t.repeat(batch_size, 1)
return cls(So2.identity(batch_size, device, dtype), Vector2(t))

[docs]    def matrix(self) -> Tensor:
"""Returns the matrix representation of shape :math:`(B, 3, 3)`.

Example:
>>> s = Se2(So2.identity(1), torch.ones(1, 2))
>>> s.matrix()
tensor([[[1., -0., 1.],
[0., 1., 1.],
"""
rt = concatenate((self.r.matrix(), self.t.data[..., None]), -1)
rt_3x3 = pad(rt, (0, 0, 0, 1))  # add last row zeros
rt_3x3[..., -1, -1] = 1.0
return rt_3x3

[docs]    @classmethod
def from_matrix(cls, matrix: Tensor) -> Se2:
"""Create an Se2 group from a matrix.

Args:
matrix: tensor of shape :math:`(B, 3, 3)`.

Example:
>>> s = Se2.from_matrix(torch.eye(3).repeat(2, 1, 1))
>>> s.r
Parameter containing:
>>> s.t
Parameter containing:
tensor([[0., 0.],
"""
# KORNIA_CHECK_SHAPE(matrix, ["B", "3", "3"])  # FIXME: resolve shape bugs. @edgarriba
r = So2.from_matrix(matrix[..., :2, :2])
t = matrix[..., :2, -1]
return cls(r, t)

[docs]    def inverse(self) -> Se2:
"""Returns the inverse transformation.

Example:
>>> s = Se2(So2.identity(1), torch.ones(1,2))
>>> s_inv = s.inverse()
>>> s_inv.r
Parameter containing:
>>> s_inv.t
Parameter containing:
"""
r_inv: So2 = self.r.inverse()
_t = -1 * self.t
if isinstance(_t, int):
raise TypeError('Unexpected integer from `-1 * translation`')

return Se2(r_inv, r_inv * _t)

[docs]    @classmethod
def random(cls, batch_size: int | None = None, device: Device | None = None, dtype: Dtype = None) -> Se2:
"""Create a Se2 group representing a random transformation.

Args:
batch_size: the batch size of the underlying data.

Example:
>>> s = Se2.random()
>>> s = Se2.random(batch_size=3)
"""
r = So2.random(batch_size, device, dtype)
shape: tuple[int, ...]
if batch_size is None:
shape = (2,)
else:
KORNIA_CHECK(batch_size >= 1, msg="batch_size must be positive")
shape = (batch_size, 2)
return cls(r, Vector2(rand(shape, device=device, dtype=dtype)))

[docs]    @classmethod
def trans(cls, x: Tensor, y: Tensor) -> Se2:
"""Construct a translation only Se2 instance.

Args:
x: the x-axis translation.
y: the y-axis translation.
"""
KORNIA_CHECK(x.shape == y.shape)
KORNIA_CHECK_SAME_DEVICES([x, y])
batch_size = x.shape[0] if len(x.shape) > 0 else None
rotation = So2.identity(batch_size, x.device, x.dtype)
return cls(rotation, stack((x, y), -1))

[docs]    @classmethod
def trans_x(cls, x: Tensor) -> Se2:
"""Construct a x-axis translation.

Args:
x: the x-axis translation.
"""
zs = zeros_like(x)
return cls.trans(x, zs)

[docs]    @classmethod
def trans_y(cls, y: Tensor) -> Se2:
"""Construct a y-axis translation.

Args:
y: the y-axis translation.
"""
zs = zeros_like(y)
return cls.trans(zs, y)