Source code for simulation.utils.geometry.transform

"""Transformation."""

import math
import numbers
from contextlib import suppress

# Compatible formats
import geometry_msgs.msg as geometry_msgs
import numpy as np
from pyquaternion import Quaternion

from simulation.utils.geometry.vector import Vector

from .frame import validate_and_maintain_frames


[docs]class Transform: """Transformation class consisting of a translation and a rotation. A Transform object can be used to easily interchange between multiple coordinate systems (which can be transformed in one another via a rotation and a translation. Initialization can be done in one of the following ways: Args: 1 (geometry_msgs/Transformation): initialize from geometry_msgs. 2 (Vector, float): Second argument is the transformation's rotation angle in radian. 3 (Vector, pyquaternion.Quaternion): Vector and quaternion. Attributes: tranlation (Vector) rotation (pyquaternion.Quaternion) """ def __init__(self, *args, frame=None): """Transform initialization.""" # Due to recursive calling of the init function, the frame should be set # in the first call within the recursion only. if not hasattr(self, "_frame"): self._frame = frame # Attempt initialization from Vector like and Quaternion like objects with suppress(Exception): args = (args[0], Quaternion(*args[1])) with suppress(Exception): if isinstance(args[1], numbers.Number): args = (args[0], Quaternion(axis=[0, 0, 1], radians=args[1])) pass # Attempt default init with suppress(IndexError, NotImplementedError, TypeError): if type(args[1]) == Quaternion: self.rotation = args[1] self.translation = Vector(args[0]) return # Try to initialize geometry pose with suppress(Exception): # Call this function with values extracted t = Vector(args[0].position) g_quaternion = args[0].orientation q = Quaternion(g_quaternion.w, g_quaternion.x, g_quaternion.y, g_quaternion.z) self.__init__(t, q) return # Try to initialize geometry transform with suppress(Exception): # Call this function with values extracted t = Vector(args[0].translation) g_quaternion = args[0].rotation q = Quaternion(g_quaternion.w, g_quaternion.x, g_quaternion.y, g_quaternion.z) self.__init__(t, q) return with suppress(Exception): # Try to initialize with two vectors translation+rotation t = args[0] rotation_vec = args[1].to_numpy() angle = (-1 if rotation_vec[1] < 0 else 1) * math.acos( np.dot([1, 0, 0], rotation_vec) / np.linalg.norm(rotation_vec) ) self.__init__(t, angle) return # None of the initializations worked raise NotImplementedError( f"Transform initialization not implemented for {type(args[0])}" ) @property @validate_and_maintain_frames def inverse(self) -> "Transform": """Inverse transformation.""" return Transform( -1 * Vector(self.translation).rotated(self.rotation.inverse), self.rotation.inverse, )
[docs] def get_angle(self, axis=Vector(0, 0, 1)) -> float: """Angle of rotation. Args: axis: Axis the rotation is projected onto. Returns: The angle that a vector is rotated, when this transformation is applied. """ # Project the rotation axis onto the rotation axis to get the amount of the rotation # that is in the axis' direction! # Also the quaternions rotation axis is sometimes flipped at which point # the angles flip their sign, # taking the scalar product with the axis fixes that as well return Vector(self.rotation.axis) * axis * self.rotation.radians
[docs] def to_geometry_msg(self) -> geometry_msgs.Transform: """Convert transform to ROS geometry_msg. Returns: This transformation as a geometry_msgs/Transform. """ vector = self.translation.to_geometry_msg() rotation = geometry_msgs.Quaternion( self.rotation.x, self.rotation.y, self.rotation.z, self.rotation.w ) tf = geometry_msgs.Transform() tf.translation = vector tf.rotation = rotation return tf
[docs] def to_affine_matrix(self) -> np.ndarray: """Get transformation as an affine matrix.""" return np.column_stack( ( self.rotation.rotation_matrix, self.translation.to_numpy(), ) )
@validate_and_maintain_frames def __mul__(self, tf: "Transform") -> "Transform": """Multiplication of transformations. The product has to be understood as a single transformation consisting of the right hand transformation applied first and then the left hand transformation. Example: Easily modify a vector multiple times: :math:`(\\text{Tf}_1*\\text{Tf}_2)*\\vec{v} = \\text{Tf}_1*( \\text{Tf}_2*\\vec{v})` Returns: The product transformation. """ if tf.__class__ == self.__class__: return self.__class__( Vector(self.translation) + Vector(tf.translation).rotated(self.rotation), self.rotation * tf.rotation, frame=self._frame, ) return NotImplemented @validate_and_maintain_frames def __eq__(self, tf) -> bool: if self.__class__ != tf.__class__: return NotImplemented return tf.rotation.normalised == self.rotation.normalised and Vector( self.translation ) == Vector(tf.translation) def __repr__(self) -> str: return ( f"Transform(translation={repr(self.translation)}," + f"rotation={repr(self.rotation)}" + (f",frame={self._frame.name}" if self._frame is not None else "") + ")" ) def __hash__(self): return NotImplemented