Source code for ManipulaPy.utils

#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-or-later
"""
Utilities Module - ManipulaPy

This module contains essential utility functions for working with rigid body motions,
transformations, and related operations in robotic manipulation. Functions cover
screw theory, matrix operations, Lie algebra computations, time scaling, and
coordinate transformations for kinematics and dynamics calculations.

The functions in this module support:
- Extracting and manipulating screw vectors and twists
- Computing transformation matrices from twists and joint angles
- Matrix logarithms and exponentials for SE(3) and SO(3)
- Converting between rotation matrices and Euler angles
- Skew-symmetric matrix operations
- Time scaling functions for trajectory generation

Copyright (c) 2025 Mohamed Aboelnasr

This file is part of ManipulaPy.

ManipulaPy is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ManipulaPy is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with ManipulaPy. If not, see <https://www.gnu.org/licenses/>.
"""

from typing import Optional, Tuple

import numpy as np
from numpy.typing import NDArray
from scipy.linalg import expm




def extract_r_list(Slist) -> NDArray[np.float64]:
    """
    Extracts the r_list from the given Slist.

    Parameters:
        Slist (list): A list of S vectors representing the joint screws.

    Returns:
        np.ndarray: An array of r vectors.
    """
    # Handle None or improperly shaped input
    if Slist is None or not hasattr(np.array(Slist), "T"):
        return np.array([])

    # Continue with the original function
    r_list = []
    for S in np.array(Slist).T:
        omega = S[:3]
        v = S[3:]
        if np.linalg.norm(omega) != 0:
            r = -np.cross(omega, v) / np.linalg.norm(omega) ** 2
            r_list.append(r)
        else:
            r_list.append([0, 0, 0])  # For prismatic joints
    return np.array(r_list)





def extract_omega_list(Slist) -> NDArray[np.float64]:
    """
    Extracts the first three elements from each sublist in the given list and returns them as a numpy array.

    Parameters:
        Slist (list): A list of sublists.

    Returns:
        np.array: A numpy array containing the first three elements from each sublist.
    """
    return np.array(Slist)[:, :3]





def extract_screw_list(omega_list, r_list) -> Optional[NDArray[np.float64]]:
    """
    Build a 6xn screw-axis matrix from (3xn) angular velocities 'omega_list'
    and (3xn) positions 'r_list'.

    For each column ``i``, ``S[:3, i]`` is set to ``omega_list[:, i]`` and
    ``S[3:, i]`` is set to ``-(omega x r)``. Returns a 6xn array of
    ``[wx, wy, wz, vx, vy, vz]`` in each column.
    """
    if omega_list is None or r_list is None:
        return None

    # Convert to numpy arrays if not already
    omega_list = np.asarray(omega_list)
    r_list = np.asarray(r_list)

    # Handle case where r_list is empty or 1D
    if r_list.size == 0:
        # Create a default r_list with zeros
        if omega_list.ndim == 2:
            n_joints = omega_list.shape[1]
        else:
            n_joints = omega_list.shape[0] // 3 if omega_list.ndim == 1 else 1
        r_list = np.zeros((3, n_joints))
    elif r_list.ndim == 1:
        # If r_list is 1D, reshape or handle appropriately
        if r_list.size == 0:
            if omega_list.ndim == 2:
                n_joints = omega_list.shape[1]
            else:
                n_joints = 1
            r_list = np.zeros((3, n_joints))
        elif r_list.size == 3:
            # Single position vector, reshape to (3, 1)
            r_list = r_list.reshape(3, 1)
        else:
            # Multiple positions in 1D array, try to reshape
            if r_list.size % 3 == 0:
                n_positions = r_list.size // 3
                r_list = r_list.reshape(3, n_positions)
            else:
                raise ValueError(
                    f"Cannot reshape r_list of size {r_list.size} into (3, n) format"
                )

    # Ensure omega_list is also 2D
    if omega_list.ndim == 1:
        if omega_list.size % 3 == 0:
            n_joints = omega_list.size // 3
            omega_list = omega_list.reshape(3, n_joints)
        else:
            raise ValueError(
                f"Cannot reshape omega_list of size {omega_list.size} into (3, n) format"
            )

    w_rows, w_cols = omega_list.shape
    r_rows, r_cols = r_list.shape

    if w_rows != 3 or r_rows != 3:
        raise ValueError("omega_list and r_list must each have 3 rows.")
    if w_cols != r_cols:
        # Try to broadcast if one has only one column
        if r_cols == 1 and w_cols > 1:
            r_list = np.tile(r_list, (1, w_cols))
            r_cols = w_cols
        elif w_cols == 1 and r_cols > 1:
            omega_list = np.tile(omega_list, (1, r_cols))
            w_cols = r_cols
        else:
            raise ValueError(
                f"omega_list and r_list must have the same number of columns. Got {w_cols} and {r_cols}."
            )

    S = np.zeros((6, w_cols), dtype=float)
    for i in range(w_cols):
        w = omega_list[:, i]
        r = r_list[:, i]
        v = np.cross(-w, r)
        S[:3, i] = w
        S[3:, i] = v
    return S





def NearZero(z: float) -> bool:
    """
    Determines if a given number is near zero.

    Parameters:
        z (float): The number to check.

    Returns:
        bool: True if the number is near zero, False otherwise.
    """
    return abs(z) < 1e-6





def skew_symmetric(v) -> NDArray[np.float64]:
    """
    Returns the skew symmetric matrix of a 3D vector.

    Parameters:
        v (array-like): A 3D vector.

    Returns:
        np.ndarray: The corresponding skew symmetric matrix.
    """
    return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])





def transform_from_twist(S, theta: float) -> NDArray[np.float64]:
    """
    Computes the transformation matrix from a twist and a joint angle.

    Parameters:
        S (array-like): A 6D twist vector.
        theta (float): The joint angle.

    Returns:
        np.ndarray: The corresponding transformation matrix.
    """
    omega = S[:3]
    v = S[3:]
    if np.linalg.norm(omega) == 0:  # Prismatic joint
        T = np.eye(4)
        T[:3, 3] = v * theta
        return T
    else:  # Revolute joint
        skew_omega = skew_symmetric(omega)
        R = (
            np.eye(3)
            + np.sin(theta) * skew_omega
            + (1 - np.cos(theta)) * np.dot(skew_omega, skew_omega)
        )
        p = np.dot(
            np.eye(3) * theta
            + (1 - np.cos(theta)) * skew_omega
            + (theta - np.sin(theta)) * np.dot(skew_omega, skew_omega),
            v,
        )
        return np.vstack((np.hstack((R, p.reshape(-1, 1))), [0, 0, 0, 1]))





def adjoint_transform(T) -> NDArray[np.float64]:
    """
    Computes the adjoint transformation matrix for a given transformation matrix.

    Parameters:
        T (np.ndarray): A 4x4 transformation matrix.

    Returns:
        np.ndarray: The corresponding adjoint transformation matrix.
    """
    R = T[0:3, 0:3]
    p = T[0:3, 3]
    skew_p = skew_symmetric(p)
    return np.vstack((np.hstack((R, np.zeros((3, 3)))), np.hstack((skew_p @ R, R))))





def logm(T) -> NDArray[np.float64]:
    """
    Computes the logarithm of a transformation matrix.

    Parameters:
        T (np.ndarray): A 4x4 transformation matrix.

    Returns:
        np.ndarray: The logarithm of the transformation matrix.
    """
    R = T[0:3, 0:3]
    p = T[0:3, 3]
    omega, theta = rotation_logm(R)
    if theta < 1e-6:
        return np.hstack((np.zeros(3), p))
    G_inv = (
        1 / theta * np.eye(3)
        - 0.5 * skew_symmetric(omega)
        + (1 / theta - 0.5 / np.tan(theta / 2))
        * np.dot(skew_symmetric(omega), skew_symmetric(omega))
    )
    v = theta * np.dot(G_inv, p)
    return np.hstack((omega * theta, v))





def rotation_logm(R) -> Tuple[NDArray[np.float64], float]:
    """
    Computes the logarithm of a rotation matrix.

    Parameters:
        R (np.ndarray): A 3x3 rotation matrix.

    Returns:
        tuple: A tuple containing the rotation vector and the angle.
    """
    # Validate input shape
    R = np.asarray(R)
    if R.shape != (3, 3):
        raise ValueError(
            f"rotation_logm requires a 3x3 rotation matrix, got shape {R.shape}. "
            f"Matrix:\n{R}"
        )

    # Ensure we're computing with scalar values
    trace_val = np.trace(R)
    # Clamp to [-1, 1] to avoid numerical issues with arccos
    cos_theta = np.clip((trace_val - 1) / 2, -1.0, 1.0)
    theta = np.arccos(cos_theta)

    # Convert to Python scalar (handles numpy scalars, 0-d arrays, and 1-d arrays)
    try:
        theta_scalar = float(theta)
    except (TypeError, ValueError):
        # Handle cases where theta is an array
        theta_arr = np.asarray(theta).flatten()
        if theta_arr.size == 0:
            theta_scalar = 0.0
        else:
            theta_scalar = float(theta_arr[0])

    # Check if rotation is very small (identity or near-identity)
    if theta_scalar < 1e-6:
        return np.zeros(3), theta_scalar
    elif theta_scalar > np.pi - 1e-6:
        # theta ~ pi: R is (near-)symmetric so R - R^T -> 0 while 1/(2 sin theta)
        # blows up, collapsing the generic formula to the zero vector. Extract
        # the axis from the most positive diagonal term instead (mirrors
        # MatrixLog3 / Modern Robotics). The threshold is kept tight (1e-6):
        # the generic formula stays accurate until extremely close to pi, and
        # the diagonal extraction assumes exactly theta = pi, so widening the
        # band would trade a smaller floating-point error for a larger
        # axis-approximation error.
        if not NearZero(1 + R[2, 2]):
            omega = (1.0 / np.sqrt(2 * (1 + R[2, 2]))) * np.array(
                [R[0, 2], R[1, 2], 1 + R[2, 2]]
            )
        elif not NearZero(1 + R[1, 1]):
            omega = (1.0 / np.sqrt(2 * (1 + R[1, 1]))) * np.array(
                [R[0, 1], 1 + R[1, 1], R[2, 1]]
            )
        else:
            omega = (1.0 / np.sqrt(2 * (1 + R[0, 0]))) * np.array(
                [1 + R[0, 0], R[1, 0], R[2, 0]]
            )
        return omega, theta_scalar
    else:
        omega = (
            1
            / (2 * np.sin(theta_scalar))
            * np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])
        )
        return omega, theta_scalar





def logm_to_twist(logm) -> NDArray[np.float64]:
    """
    Convert the logarithm of a transformation matrix to a twist vector.

    Parameters:
        logm (np.ndarray): The logarithm of a transformation matrix.

    Returns:
        np.ndarray: The corresponding twist vector.
    """
    if logm.shape != (4, 4):
        raise ValueError("logm must be a 4x4 matrix.")

    omega_matrix = logm[0:3, 0:3]
    omega = skew_symmetric_to_vector(omega_matrix)
    v = logm[0:3, 3]
    return np.hstack((omega, v))





def skew_symmetric_to_vector(skew_symmetric) -> NDArray[np.float64]:
    """
    Convert a skew-symmetric matrix to a vector.

    Parameters:
        skew_symmetric (np.ndarray): A 3x3 skew-symmetric matrix.

    Returns:
        np.ndarray: The corresponding vector.
    """
    return np.array([skew_symmetric[2, 1], skew_symmetric[0, 2], skew_symmetric[1, 0]])





def se3ToVec(se3_matrix) -> NDArray[np.float64]:
    """
    Convert an se(3) matrix to a twist vector.

    Parameters:
        se3_matrix (np.ndarray): A 4x4 matrix from the se(3) Lie algebra.

    Returns:
        np.ndarray: A 6-dimensional twist vector.
    """
    if se3_matrix.shape != (4, 4):
        raise ValueError("Input matrix must be a 4x4 matrix.")

    omega = np.array([se3_matrix[2, 1], se3_matrix[0, 2], se3_matrix[1, 0]])
    v = se3_matrix[0:3, 3]
    return np.hstack((omega, v))





def TransToRp(T) -> Tuple[NDArray[np.float64], NDArray[np.float64]]:
    """
    Converts a homogeneous transformation matrix into a rotation matrix and position vector.

    Parameters:
        T (np.ndarray): A 4x4 transformation matrix.

    Returns:
        tuple: A tuple containing the rotation matrix and position vector.
    """
    R = T[0:3, 0:3]
    p = T[0:3, 3]
    return R, p





def TransInv(T) -> NDArray[np.float64]:
    """
    Inverts a homogeneous transformation matrix.

    Parameters:
        T (np.ndarray): A 4x4 transformation matrix.

    Returns:
        np.ndarray: The inverse of the transformation matrix.
    """
    R, p = TransToRp(T)
    Rt = R.T
    return np.vstack((np.hstack((Rt, -Rt @ p.reshape(-1, 1))), [0, 0, 0, 1]))





def MatrixLog6(T) -> NDArray[np.float64]:
    """
    Compute the matrix logarithm of a given transformation matrix T.

    Parameters:
        T (np.ndarray): The transformation matrix of shape (4, 4).

    Returns:
        np.ndarray: The matrix logarithm of T, with shape (4, 4).
    """
    R, p = TransToRp(T)
    omega, theta = rotation_logm(R)

    # Pure translation or extremely small rotation
    if abs(theta) < 1e-6:
        return np.vstack(
            (np.hstack((np.zeros((3, 3)), p.reshape(-1, 1))), [0, 0, 0, 0])
        )

    # Use the standard G^{-1}(theta) from Modern Robotics (Eq. 3.88):
    # G_inv = I/theta - 0.5*w_hat + (1/theta - 0.5*cot(theta/2)) * w_hat^2
    w_hat = skew_symmetric(omega)
    G_inv = (
        (np.eye(3) / theta)
        - 0.5 * w_hat
        + (1 / theta - 0.5 / np.tan(theta / 2)) * (w_hat @ w_hat)
    )

    # se(3) log has w_hat*theta in the rotation block
    omega_mat_scaled = theta * w_hat
    v = theta * (G_inv @ p)
    return np.vstack((np.hstack((omega_mat_scaled, v.reshape(-1, 1))), [0, 0, 0, 0]))





def MatrixExp6(se3mat) -> NDArray[np.float64]:
    """
    Computes the matrix exponential of a matrix in se(3).

    Parameters:
        se3mat (np.ndarray): A 4x4 matrix representing a twist in se(3).

    Returns:
        np.ndarray: The corresponding 4x4 transformation matrix in SE(3).
    """
    if se3mat.shape != (4, 4):
        raise ValueError("Input matrix must be of shape (4, 4)")

    # Extract rotation (so(3)) and translation components
    omega_theta_vec = np.array([se3mat[2, 1], se3mat[0, 2], se3mat[1, 0]])
    v = se3mat[0:3, 3]
    theta = np.linalg.norm(omega_theta_vec)

    # Pure translation case
    if theta < 1e-6:
        T = np.eye(4)
        T[:3, 3] = v
        return T

    # Unit rotation axis hat matrix
    omega_hat = se3mat[0:3, 0:3] / theta

    # Rotation component
    R = MatrixExp3(se3mat[0:3, 0:3])

    # Compute the matrix G(theta) that maps body velocity to translation
    G = (
        np.eye(3) * theta
        + (1 - np.cos(theta)) * omega_hat
        + (theta - np.sin(theta)) * (omega_hat @ omega_hat)
    )

    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = (G @ v) / theta
    return T





def MatrixLog3(R) -> NDArray[np.float64]:
    """
    Computes the matrix logarithm of a rotation matrix.

    Parameters:
        R (np.ndarray): A 3x3 rotation matrix.

    Returns:
        np.ndarray: The matrix logarithm of the rotation matrix.
    """
    acosinput = (np.trace(R) - 1) / 2.0
    if acosinput >= 1:
        return np.zeros((3, 3))
    elif acosinput <= -1:
        if not NearZero(1 + R[2][2]):
            omega = (1.0 / np.sqrt(2 * (1 + R[2][2]))) * np.array(
                [R[0][2], R[1][2], 1 + R[2][2]]
            )
        elif not NearZero(1 + R[1][1]):
            omega = (1.0 / np.sqrt(2 * (1 + R[1][1]))) * np.array(
                [R[0][1], 1 + R[1][1], R[2][1]]
            )
        else:
            omega = (1.0 / np.sqrt(2 * (1 + R[0][0]))) * np.array(
                [1 + R[0][0], R[1][0], R[2][0]]
            )
        return VecToso3(np.pi * omega)
    else:
        theta = np.arccos(acosinput)
        return theta / 2.0 / np.sin(theta) * (R - np.array(R).T)





def VecToso3(omega) -> NDArray[np.float64]:
    """
    Converts a 3D angular velocity vector to a skew-symmetric matrix.

    Parameters:
        omega (array-like): A 3D angular velocity vector.

    Returns:
        np.ndarray: The corresponding skew-symmetric matrix.
    """
    return np.array(
        [[0, -omega[2], omega[1]], [omega[2], 0, -omega[0]], [-omega[1], omega[0], 0]]
    )





def VecTose3(V) -> NDArray[np.float64]:
    """
    Converts a spatial velocity vector to an se(3) matrix.

    Parameters:
        V (array-like): A 6D spatial velocity vector.

    Returns:
        np.ndarray: The corresponding 4x4 matrix in se(3).
    """
    return np.r_[np.c_[VecToso3(V[:3]), V[3:].reshape(3, 1)], np.zeros((1, 4))]





def MatrixExp3(so3mat) -> NDArray[np.float64]:
    """
    Computes the matrix exponential of a matrix in so(3).

    Parameters:
        so3mat (np.ndarray): A 3x3 skew-symmetric matrix.

    Returns:
        np.ndarray: The corresponding 3x3 rotation matrix in SO(3).
    """
    return expm(so3mat)





def CubicTimeScaling(Tf: float, t: float) -> float:
    """
    Compute the cubic time scaling factor.

    Parameters:
        Tf (float): The total time of the motion.
        t (float): The current time.

    Returns:
        float: The cubic time scaling factor.
    """
    return 3 * (t / Tf) ** 2 - 2 * (t / Tf) ** 3





def QuinticTimeScaling(Tf: float, t: float) -> float:
    """
    Compute the quintic time scaling factor.

    Parameters:
        Tf (float): The total time of the motion.
        t (float): The current time.

    Returns:
        float: The quintic time scaling factor.
    """
    return 10 * (t / Tf) ** 3 - 15 * (t / Tf) ** 4 + 6 * (t / Tf) ** 5





def rotation_matrix_to_euler_angles(R) -> NDArray[np.float64]:
    """
    Convert a rotation matrix to Euler angles (roll, pitch, yaw).

    Parameters:
        R (numpy.ndarray): A 3x3 rotation matrix.

    Returns:
        numpy.ndarray: A 3-element array representing the Euler angles (roll, pitch, yaw).
    """
    assert R.shape == (3, 3), f"Expected 3x3 rotation matrix, got shape {R.shape}"
    sy = np.sqrt(R[0, 0] ** 2 + R[1, 0] ** 2)

    singular = sy < 1e-6

    if not singular:
        x = np.arctan2(R[2, 1], R[2, 2])
        y = np.arctan2(-R[2, 0], sy)
        z = np.arctan2(R[1, 0], R[0, 0])
    else:
        x = np.arctan2(-R[1, 2], R[1, 1])
        y = np.arctan2(-R[2, 0], sy)
        z = 0

    return np.array([x, y, z])





def euler_to_rotation_matrix(euler_deg) -> NDArray[np.float64]:
    """
    Convert Euler angles (roll_deg, pitch_deg, yaw_deg) in degrees
    to a 3x3 rotation matrix.
    ZYX convention is typical in robotics: yaw -> pitch -> roll,
    but adapt as needed.

    Parameters:
        euler_deg (array-like): [roll_deg, pitch_deg, yaw_deg]

    Returns:
        np.ndarray: A 3x3 rotation matrix (float64 by default).
    """
    roll_deg, pitch_deg, yaw_deg = euler_deg
    # Convert degrees to radians
    roll = np.radians(roll_deg)
    pitch = np.radians(pitch_deg)
    yaw = np.radians(yaw_deg)

    # Example Z-Y-X convention (yaw→pitch→roll).
    # If your code uses X→Y→Z or another sequence, adapt these multiplications.
    Rz = np.array(
        [[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]],
        dtype=np.float64,
    )

    Ry = np.array(
        [
            [np.cos(pitch), 0, np.sin(pitch)],
            [0, 1, 0],
            [-np.sin(pitch), 0, np.cos(pitch)],
        ],
        dtype=np.float64,
    )

    Rx = np.array(
        [[1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)]],
        dtype=np.float64,
    )

    # Multiply in the correct order for your convention.
    R = Rz @ Ry @ Rx
    return R