nova_rerun_bridge.dh_robot

  1import numpy as np
  2
  3from nova.api import models
  4
  5
  6class DHRobot:
  7    """A class for handling DH parameters and computing joint positions."""
  8
  9    def __init__(self, dh_parameters: list[models.DHParameter], mounting: models.PlannerPose):
 10        """
 11        Initialize the DHRobot with DH parameters and a mounting pose.
 12        :param dh_parameters: List of DHParameter objects containing all joint configurations.
 13        :param mounting: PlannerPose object representing the mounting orientation and position.
 14        """
 15        self.dh_parameters = dh_parameters
 16        self.mounting = mounting
 17
 18    def pose_to_matrix(self, pose: models.PlannerPose):
 19        """
 20        Convert a PlannerPose (with quaternion orientation) into a 4x4 homogeneous transformation matrix.
 21        :param pose: A PlannerPose object with position: Vector3d and orientation: Quaternion.
 22        :return: A 4x4 numpy array representing the transformation.
 23        """
 24        # Extract translation
 25        if pose.position is None:
 26            x, y, z = 0.0, 0.0, 0.0
 27        else:
 28            x, y, z = pose.position.x, pose.position.y, pose.position.z
 29
 30        # Extract quaternion
 31        if pose.orientation is None:
 32            # If no orientation is provided, assume identity orientation
 33            w, qx, qy, qz = 1.0, 0.0, 0.0, 0.0
 34        else:
 35            w = pose.orientation.w
 36            qx = pose.orientation.x
 37            qy = pose.orientation.y
 38            qz = pose.orientation.z
 39
 40        # Compute rotation matrix from quaternion
 41        # R = [[1 - 2(y²+z²), 2(xy - zw),     2(xz + yw)    ],
 42        #      [2(xy + zw),     1 - 2(x²+z²), 2(yz - xw)    ],
 43        #      [2(xz - yw),     2(yz + xw),   1 - 2(x²+y²)]]
 44
 45        xx = qx * qx
 46        yy = qy * qy
 47        zz = qz * qz
 48        xy = qx * qy
 49        xz = qx * qz
 50        yz = qy * qz
 51        wx = w * qx
 52        wy = w * qy
 53        wz = w * qz
 54
 55        R = np.array(
 56            [
 57                [1.0 - 2.0 * (yy + zz), 2.0 * (xy - wz), 2.0 * (xz + wy)],
 58                [2.0 * (xy + wz), 1.0 - 2.0 * (xx + zz), 2.0 * (yz - wx)],
 59                [2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (xx + yy)],
 60            ]
 61        )
 62
 63        # Construct the full homogeneous transformation matrix
 64        T = np.eye(4)
 65        T[0:3, 0:3] = R
 66        T[0:3, 3] = [x, y, z]
 67
 68        return T
 69
 70    def dh_transform(self, dh_param: models.DHParameter, joint_rotation):
 71        """
 72        Compute the homogeneous transformation matrix for a given DH parameter and joint rotation.
 73        :param dh_param: A single DH parameter.
 74        :param joint_rotation: The joint rotation value (in radians).
 75        :return: A 4x4 homogeneous transformation matrix.
 76        """
 77        # Adjust the angle based on rotation direction
 78        theta = dh_param.theta + joint_rotation * (-1 if dh_param.reverse_rotation_direction else 1)
 79        d = dh_param.d
 80        a = dh_param.a
 81        alpha = dh_param.alpha if dh_param.alpha is not None else 0.0
 82
 83        # Create the homogeneous transformation matrix for this DH parameter
 84        transformation = np.array(
 85            [
 86                [
 87                    np.cos(theta),
 88                    -np.sin(theta) * np.cos(alpha),
 89                    np.sin(theta) * np.sin(alpha),
 90                    a * np.cos(theta),
 91                ],
 92                [
 93                    np.sin(theta),
 94                    np.cos(theta) * np.cos(alpha),
 95                    -np.cos(theta) * np.sin(alpha),
 96                    a * np.sin(theta),
 97                ],
 98                [0, np.sin(alpha), np.cos(alpha), d],
 99                [0, 0, 0, 1],
100            ]
101        )
102        return transformation
103
104    def calculate_joint_positions(self, joint_values):
105        """
106        Compute joint positions based on joint values.
107        :param joint_values: Object containing joint rotation values as a list in joint_values.joints.
108        :return: A list of joint positions as [x, y, z].
109        """
110        # Incorporate the mounting pose at the start
111        accumulated_matrix = self.pose_to_matrix(self.mounting)
112
113        joint_positions = [
114            accumulated_matrix[:3, 3].tolist()
115        ]  # Base position after mounting is applied
116
117        for dh_param, joint_rotation in zip(self.dh_parameters, joint_values.joints, strict=False):
118            transform = self.dh_transform(dh_param, joint_rotation)
119            accumulated_matrix = accumulated_matrix @ transform
120            position = accumulated_matrix[:3, 3]  # Extract translation (x, y, z)
121            joint_positions.append(position.tolist())
122
123        return joint_positions