def get_joint_matrix()

in robogym/mujoco/forward_kinematics.py [0:0]

• 35 lines of code
• 10 McCabe index (conditional complexity)

def get_joint_matrix(pos, angle, axis):
    def transform_rot_x_matrix(pos, angle):
        """
        Optimization - create a homogeneous matrix where rotation submatrix
        rotates around the X axis by given angle in radians
        """
        m = np.eye(4)
        m[1, 1] = m[2, 2] = np.cos(angle)
        s = np.sin(angle)
        m[1, 2] = -s
        m[2, 1] = s
        m[:3, 3] = pos
        return m

    def transform_rot_y_matrix(pos, angle):
        """
        Optimization - create a homogeneous matrix where rotation submatrix
        rotates around the Y axis by given angle in radians
        """
        m = np.eye(4)
        m[0, 0] = m[2, 2] = np.cos(angle)
        s = np.sin(angle)
        m[0, 2] = s
        m[2, 0] = -s
        m[:3, 3] = pos
        return m

    def transform_rot_z_matrix(pos, angle):
        """
        Optimization - create a homogeneous matrix where rotation submatrix
        rotates around the Z axis by given angle in radians
        """
        m = np.eye(4)
        m[0, 0] = m[1, 1] = np.cos(angle)
        s = np.sin(angle)
        m[0, 1] = -s
        m[1, 0] = s
        m[:3, 3] = pos
        return m

    if abs(axis[0]) == 1.0 and axis[1] == 0.0 and axis[2] == 0.0:
        return transform_rot_x_matrix(pos, angle * axis[0])
    elif axis[0] == 0.0 and abs(axis[1]) == 1.0 and axis[2] == 0.0:
        return transform_rot_y_matrix(pos, angle * axis[1])
    elif axis[0] == 0.0 and axis[1] == 0.0 and abs(axis[2]) == 1.0:
        return transform_rot_z_matrix(pos, angle * axis[2])
    else:
        return homogeneous_matrix_from_pos_mat(
            pos, rot.quat2mat(rot.quat_from_angle_and_axis(angle, axis))
        )