# kornia.geometry.se2 module inspired by Sophus-sympy.
# https://github.com/strasdat/Sophus/blob/master/sympy/sophus/se2.py
from typing import Optional, Tuple, overload
from kornia.core import Module, Parameter, Tensor, concatenate, pad, rand, stack, tensor, where, zeros_like
from kornia.geometry.liegroup._utils import check_se2_omega_shape, check_se2_r_t_shape, check_v_shape
from kornia.geometry.liegroup.so2 import So2
from kornia.testing import KORNIA_CHECK, KORNIA_CHECK_SAME_DEVICES, KORNIA_CHECK_TYPE
[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:
tensor([1.+0.j], requires_grad=True)
translation: Parameter containing:
tensor([[1., 1.]], requires_grad=True)
"""
[docs] def __init__(self, rotation: So2, translation: 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:
tensor([1.+0.j], requires_grad=True)
translation: Parameter containing:
tensor([[1., 1.]], requires_grad=True)
"""
super().__init__()
KORNIA_CHECK_TYPE(rotation, So2)
# TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support
check_se2_r_t_shape(rotation, translation)
self._rotation = rotation
self._translation = Parameter(translation)
[docs] def __repr__(self) -> str:
return f"rotation: {self.r}\ntranslation: {self.t}"
def __getitem__(self, idx) -> 'Se2':
return Se2(self._rotation[idx], self._translation[idx])
@overload
def __mul__(self, right: 'Se2') -> 'Se2':
...
@overload
def __mul__(self, right: Tensor) -> Tensor:
...
[docs] def __mul__(self, right):
"""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, Tensor):
KORNIA_CHECK_TYPE(right, Tensor)
# TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support
check_se2_r_t_shape(so2, right)
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) -> Tensor:
"""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) -> Tensor:
"""Return the underlying translation vector of shape :math:`(B,2)`."""
return self._translation
[docs] @staticmethod
def exp(v) -> '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:
tensor([0.5403+0.8415j], requires_grad=True)
>>> s.t
Parameter containing:
tensor([[0.3818, 1.3012]], requires_grad=True)
"""
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
tensor([[1.0000, 1.0000, 1.0000]], grad_fn=<StackBackward0>)
"""
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[..., None]
return stack((upsilon[..., 0, 0], upsilon[..., 1, 0], theta), -1)
[docs] @staticmethod
def hat(v):
"""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)
return pad(col0, (0, 1))
[docs] @staticmethod
def vee(omega) -> 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: Optional[int] = None, device=None, 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:
tensor([1.+0.j], requires_grad=True)
>>> s.t
Parameter containing:
tensor([[0., 0.]], requires_grad=True)
"""
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), 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.],
[0., 0., 1.]]], grad_fn=<CopySlices>)
"""
rt = concatenate((self.r.matrix(), self.t[..., None]), -1)
rt_3x3 = pad(rt, (0, 0, 0, 1)) # add last row zeros
rt_3x3[..., -1, -1] = 1.0
return rt_3x3
[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:
tensor([1.+0.j], requires_grad=True)
>>> s_inv.t
Parameter containing:
tensor([[-1., -1.]], requires_grad=True)
"""
r_inv: So2 = self.r.inverse()
t_inv: Tensor = r_inv * (-1 * self.t)
return Se2(r_inv, t_inv)
[docs] @classmethod
def random(cls, batch_size: Optional[int] = None, device=None, 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, 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)
[docs] def adjoint(self) -> Tensor:
"""Returns the adjoint matrix of shape :math:`(B, 3, 3)`.
Example:
>>> s = Se2.identity()
>>> s.adjoint()
tensor([[1., -0., 0.],
[0., 1., -0.],
[0., 0., 1.]], grad_fn=<CopySlices>)
"""
rt = self.matrix()
rt[..., 0:2, 2] = stack((self.t[..., 1], -self.t[..., 0]), -1)
return rt