Source code for concepts.math.frame_utils_xyzw

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : frame_utils_xyzw.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/18/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
from typing import Tuple, TYPE_CHECKING
from concepts.math.rotationlib_xyzw import quat_mul, quat_conjugate, rotate_vector, mat2quat
from concepts.utils.typing_utils import Vec3f, Vec4f

if TYPE_CHECKING:
    from concepts.simulator.pybullet.client import BulletClient




def compose_transformation(pos1: Vec3f, quat1: Vec4f, pos2: Vec3f, quat2: Vec4f) -> Tuple[np.ndarray, np.ndarray]:
    """Compose two transformations.

    The transformations are represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.

    Args:
        pos1: the position of the first transformation.
        quat1: the quaternion of the first transformation.
        pos2: the position of the second transformation.
        quat2: the quaternion of the second transformation.

    Returns:
        the composed transformation, represented as a tuple of position and quaternion.
    """
    pos1 = np.asarray(pos1, dtype=np.float64)
    quat1 = np.asarray(quat1, dtype=np.float64)
    pos2 = np.asarray(pos2, dtype=np.float64)
    quat2 = np.asarray(quat2, dtype=np.float64)

    return pos1 + rotate_vector(pos2, quat1), quat_mul(quat1, quat2)





def inverse_transformation(pos: Vec3f, quat: Vec4f) -> Tuple[np.ndarray, np.ndarray]:
    """Inverse a transformation.

    The transformation is represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos = np.asarray(pos, dtype=np.float64)
    quat = np.asarray(quat, dtype=np.float64)

    inv_pos = -pos
    inv_quat = quat_conjugate(quat)
    return inv_pos, inv_quat





def get_transform_a_to_b(pos1: Vec3f, quat1: Vec4f, pos2: Vec3f, quat2: Vec4f) -> Tuple[np.ndarray, np.ndarray]:
    """Get the transformation from frame A to frame B.

    The transformations are represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos1 = np.asarray(pos1, dtype=np.float64)
    quat1 = np.asarray(quat1, dtype=np.float64)
    pos2 = np.asarray(pos2, dtype=np.float64)
    quat2 = np.asarray(quat2, dtype=np.float64)

    inv_quat1 = quat_conjugate(quat1)
    a_to_b_pos = rotate_vector(pos2 - pos1, inv_quat1)
    a_to_b_quat = quat_mul(inv_quat1, quat2)
    return a_to_b_pos, a_to_b_quat





def frame_mul(pos_a: Vec3f, quat_a: Vec4f, a_to_b: Tuple[Vec3f, Vec4f]) -> Tuple[np.ndarray, np.ndarray]:
    """Multiply a frame with a transformation.

    The frame is represented as a tuple of position and quaternion.
    The transformation is represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos_a = np.asarray(pos_a, dtype=np.float64)
    quat_a = np.asarray(quat_a, dtype=np.float64)
    transform_pos = np.asarray(a_to_b[0], dtype=np.float64)
    transform_quat = np.asarray(a_to_b[1], dtype=np.float64)

    pos_b = pos_a + rotate_vector(transform_pos, quat_a)
    quat_b = quat_mul(quat_a, transform_quat)
    return pos_b, quat_b





def frame_inv(pos_b: Vec3f, quat_b: Vec4f, a_to_b: Tuple[Vec3f, Vec4f]) -> Tuple[np.ndarray, np.ndarray]:
    """Inverse a frame with a transformation.

    The frame is represented as a tuple of position and quaternion.
    The transformation is represented as a tuple of position and quaternion.
    The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
    """
    pos_b = np.asarray(pos_b, dtype=np.float64)
    quat_b = np.asarray(quat_b, dtype=np.float64)
    transform_pos = np.asarray(a_to_b[0], dtype=np.float64)
    transform_quat = np.asarray(a_to_b[1], dtype=np.float64)

    inv_transform_quat = quat_conjugate(transform_quat)
    quat_a = quat_mul(quat_b, inv_transform_quat)
    pos_a = pos_b - rotate_vector(transform_pos, quat_a)
    return pos_a, quat_a





def solve_tool_from_ee(ee_pos: Vec3f, ee_quat: Vec4f, ee_to_tool: Tuple[Vec3f, Vec4f]) -> Tuple[np.ndarray, np.ndarray]:
    return frame_mul(ee_pos, ee_quat, ee_to_tool)





def solve_ee_from_tool(target_tool_pos: Vec3f, target_tool_quat: Vec4f, ee_to_tool: Tuple[Vec3f, Vec4f]) -> Tuple[np.ndarray, np.ndarray]:
    """Solve for the end-effector position and orientation given the tool position and orientation."""
    return frame_inv(target_tool_pos, target_tool_quat, ee_to_tool)





def compute_ee_rotation_mat_from_directions(u: Vec3f, v: Vec3f) -> np.ndarray:
    """Compute the rotation matrix from two directions (the "down" direction for the end effector and the "forward" direction for the end effector).

    Args:
        u: The "down" direction for the end effector.
        v: The "forward" direction for the end effector.
    """
    u = np.asarray(u)
    u = u / np.linalg.norm(u)
    v = np.asarray(v)
    v = v - np.dot(u, v) * u
    v = v / np.linalg.norm(v)
    w = np.cross(u, v)
    return np.array([u, v, w]).T





def compute_ee_quat_from_directions(u: Vec3f, v: Vec3f, default_quat: Vec4f = (0, 1, 0, 0)) -> np.ndarray:
    """Compute the quaternion from two directions (the "down" direction for the end effector and the "forward" direction for the end effector).

    Args:
        u: the "down" direction for the end effector.
        v: the "forward" direction for the end effector.
        default_quat: the default quaternion to be multiplied. This is the quaternion that represents the rotation for the default end effector orientation,
            facing downwards and the forward direction is along the x-axis.
    """
    mat = compute_ee_rotation_mat_from_directions(u, v)
    mat_reference = compute_ee_rotation_mat_from_directions((0, 0, -1), (1, 0, 0))
    quat = mat2quat(np.dot(mat, np.linalg.inv(mat_reference)))
    return quat_mul(quat, default_quat)