Code Generation from Unified Robot Description Format (URDF) for Accelerated Robotics
Code Generation from Unified Robot Description Format (URDF) for Accelerated Robotics Code Generation from Unified Robot Description Format (URDF) for Accelerated Robotics
Code Generation from Unified Robot Description Format (URDF) for Accelerated Robotics
Code Generation from Unified Robot Description Format (URDF) for Accelerated Robotics

Introduction

  1. I am a first time presenter and attender at CppCon
  2. About me: I work as a robotics scientist at PickNik Robotics
  3. I earned my PhD in Computer Science on robot learning from the University of New Hampshire

Outline

  • Background
  • Motivation
  • Build Process
  • Code Implementation
  • Results
  • Summary: lessons learned

The Background

Motivation

Concept video

Motivation

  • Some algorithms need to run at incredible speeds for optimal performance
  • This is especially true for common subroutines used by expensive algorithms
  • For example, robot path planning with RRT depends on nearest neighbor lookup and forward kinematics, and collision detection

Motivation

  • Performance improvements of more than 500x over the state-of-the-art
  • Compiler takes in standard Unified Robot Description Format (URDF) files and generates optimized code
  • Setup data structure to optimize SIMD execution
  • Skip unneeded computations like combining fixed joints, unrolling loops

Motivation

  • Software written by hardware manufactures has an inherit advantage because it can be written with very specific platforms in mind and take advantage of hardware specific knowledge
  • For example, robot design with spherical wrist can decompose inverse kinematics into two stages 1) find the wrist angles to achieve the desired orientation and 2) find the position value needed to find reach the desired Cartesian position
  • Generic IK solvers cannot take advantage of this kind of knowledge in general
  • One downside of hardware specific implementations is scalability
  • Hardware generally does not change very often, other than attachments, like cameras and end effectors

Motivation

  • Compiled code can be tested for memory allocations and real-time compatibility. Safety critical applications, such as surgical robots need to meet these requirements
  • Allows for hardware specific CPU optimizations to be enabled since the code is source built on every machine
  • Some advantages of code generation as opposed to a templated robot builder are 1) less likely to make a mistake building defining the robot 2) do not have to maintain code for each of your robots

The Build Process


cmake_minimum_required(VERSION 3.22)
project(fast_forward_kinematics)

# create executable
add_executable(forward_kinematics_test forwardKinematicsTest.cpp)

# setup targets to generate C++ code
set(URDF_FILE "data/robot.urdf")
set(ROOT "base_link")
set(TIP "tool0")
generate_fast_forward_kinematics_library(${URDF_FILE} ${ROOT} ${TIP})

# link against generated code
target_link_libraries(forward_kinematics_test PUBLIC fast_forward_kinematics_library)

          

function(generate_fast_forward_kinematics_library URDF_FILE ROOT_LINK TIP_LINK)

    find_package(Python REQUIRED COMPONENTS Interpreter)
    if (Python_Interpreter_FOUND)
        message(STATUS "Python executable: ${Python_EXECUTABLE}")
    else ()
        message(FATAL_ERROR "Python interpreter not found.")
    endif ()

    execute_process(
            COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/get_num_joints.py
            ${URDF_FILE}
            ${ROOT_LINK}
            ${TIP_LINK}
            OUTPUT_VARIABLE FAST_FK_NUMBER_OF_JOINTS
            OUTPUT_STRIP_TRAILING_WHITESPACE
            COMMAND_ECHO STDOUT
    )

    include(ExternalProject)
    ExternalProject_Add(
            LBFGSpp
            PREFIX ${CMAKE_BINARY_DIR}/LBFGSpp
            GIT_REPOSITORY https://github.com/yixuan/LBFGSpp.git
            GIT_TAG v0.3.0
            CMAKE_ARGS -DCMAKE_INSTALL_PREFIX=${CMAKE_BINARY_DIR}
    )
    ExternalProject_Get_Property(LBFGSpp source_dir)
    set(LBFGSppIncludeDir ${source_dir}/include)

    add_custom_command(
            OUTPUT forward_kinematics_lib.cpp
            COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/robot_gen.py
            ${URDF_FILE}
            ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
            ${CMAKE_CURRENT_BINARY_DIR}/forward_kinematics_lib.cpp
            ${ROOT_LINK}
            ${TIP_LINK}
            DEPENDS ${URDF_FILE} ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
            COMMENT
            "Running `${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/robot_gen.py
                      ${URDF_FILE}
                      ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
                      ${CMAKE_CURRENT_BINARY_DIR}/forward_kinematics_test.cpp
                      ${ROOT_LINK}
                      ${TIP_LINK}`"
            VERBATIM
    )
    add_custom_target(code_generation DEPENDS forward_kinematics_lib.cpp)

    add_library(fast_forward_kinematics_library SHARED forward_kinematics_lib.cpp)
    add_dependencies(fast_forward_kinematics_library code_generation)
    add_dependencies(fast_forward_kinematics_library LBFGSpp)

    find_package(Eigen3 3.3 NO_MODULE)
    target_compile_definitions(fast_forward_kinematics_library PUBLIC "${FAST_FK_NUMBER_OF_JOINTS}")
    target_include_directories(fast_forward_kinematics_library PUBLIC ${CMAKE_SOURCE_DIR}/include)
    target_include_directories(fast_forward_kinematics_library PUBLIC ${LBFGSppIncludeDir})
    target_link_libraries(fast_forward_kinematics_library PUBLIC Eigen3::Eigen)

endfunction()
















#
          

from urdf_parser_py import urdf
import argparse


def run():
    parser = argparse.ArgumentParser()
    parser.add_argument('urdf_file')
    parser.add_argument('root_link_name')
    parser.add_argument('tip_link_name')
    args = parser.parse_args()

    root_link_name = args.root_link_name
    tip_link_name = args.tip_link_name
    with open(args.urdf_file) as f:
        robot = urdf.Robot.from_xml_string(f.read())

    joint_names = []
    while tip_link_name != root_link_name:
        tip_joint_name, tip_link_name = robot.parent_map[tip_link_name]
        joint_names.append(tip_joint_name)

    print(f"FAST_FK_NUMBER_OF_JOINTS={len(joint_names)}", end="")


if __name__ == "__main__":
    run()

          

<?xml version="1.0" ?>
<robot name="ur5e">
  <link name="world"/>
  <link name="base_link"/>
  <link name="base_link_inertia">
    <visual>
      <origin rpy="0 0 3.141592653589793" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/visual/base.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 3.141592653589793" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/collision/base.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="4.0"/>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <inertia ixx="0.00443333156" ixy="0.0" ixz="0.0" iyy="0.00443333156" iyz="0.0" izz="0.0072"/>
    </inertial>
  </link>
  <link name="shoulder_link">
    <visual>
      <origin rpy="0 0 3.141592653589793" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/visual/shoulder.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 3.141592653589793" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/collision/shoulder.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="3.7"/>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <inertia ixx="0.010267495893" ixy="0.0" ixz="0.0" iyy="0.010267495893" iyz="0.0" izz="0.00666"/>
    </inertial>
  </link>
  <link name="upper_arm_link">
    <visual>
      <origin rpy="1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0.138"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/visual/upperarm.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0.138"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/collision/upperarm.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="8.393"/>
      <origin rpy="0 1.5707963267948966 0" xyz="-0.2125 0.0 0.138"/>
      <inertia ixx="0.1338857818623325" ixy="0.0" ixz="0.0" iyy="0.1338857818623325" iyz="0.0" izz="0.0151074"/>
    </inertial>
  </link>
  <link name="forearm_link">
    <visual>
      <origin rpy="1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0.007"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/visual/forearm.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0.007"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/collision/forearm.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="2.275"/>
      <origin rpy="0 1.5707963267948966 0" xyz="-0.1961 0.0 0.007"/>
      <inertia ixx="0.031209355099586295" ixy="0.0" ixz="0.0" iyy="0.031209355099586295" iyz="0.0" izz="0.004095"/>
    </inertial>
  </link>
  <link name="wrist_1_link">
    <visual>
      <origin rpy="1.5707963267948966 0 0" xyz="0 0 -0.127"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/visual/wrist1.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="1.5707963267948966 0 0" xyz="0 0 -0.127"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/collision/wrist1.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="1.219"/>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <inertia ixx="0.0025598989760400002" ixy="0.0" ixz="0.0" iyy="0.0025598989760400002" iyz="0.0" izz="0.0021942"/>
    </inertial>
  </link>
  <link name="wrist_2_link">
    <visual>
      <origin rpy="0 0 0" xyz="0 0 -0.0997"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/visual/wrist2.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 0" xyz="0 0 -0.0997"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/collision/wrist2.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="1.219"/>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <inertia ixx="0.0025598989760400002" ixy="0.0" ixz="0.0" iyy="0.0025598989760400002" iyz="0.0" izz="0.0021942"/>
    </inertial>
  </link>
  <link name="wrist_3_link">
    <visual>
      <origin rpy="1.5707963267948966 0 0" xyz="0 0 -0.0989"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/visual/wrist3.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="1.5707963267948966 0 0" xyz="0 0 -0.0989"/>
      <geometry>
        <mesh filename="package://ur_description/meshes/ur5e/collision/wrist3.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.1879"/>
      <origin rpy="0 0 0" xyz="0.0 0.0 -0.0229"/>
      <inertia ixx="9.890410052167731e-05" ixy="0.0" ixz="0.0" iyy="9.890410052167731e-05" iyz="0.0" izz="0.0001321171875"/>
    </inertial>
  </link>
  <!-- base_joint fixes base_link to the environment -->
  <joint name="base_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0 0"/>
    <parent link="world"/>
    <child link="base_link"/>
  </joint>
  <!-- joints - main serial chain -->
  <joint name="base_link-base_link_inertia" type="fixed">
    <parent link="base_link"/>
    <child link="base_link_inertia"/>
    <!-- 'base_link' is REP-103 aligned (so X+ forward), while the internal
           frames of the robot/controller have X+ pointing backwards.
           Use the joint between 'base_link' and 'base_link_inertia' (a dummy
           link/frame) to introduce the necessary rotation over Z (of pi rad).
      -->
    <origin rpy="0 0 3.141592653589793" xyz="0 0 0"/>
  </joint>
  <joint name="shoulder_pan_joint" type="revolute">
    <parent link="base_link_inertia"/>
    <child link="shoulder_link"/>
    <origin rpy="0 0 4.898391345320441e-09" xyz="0 0 0.1625702965797758"/>
    <axis xyz="0 0 1"/>
    <limit effort="150.0" lower="-3.141592653589793" upper="3.141592653589793" velocity="0.7853981633974483"/>
    <dynamics damping="0" friction="0"/>
  </joint>
  <joint name="shoulder_lift_joint" type="revolute">
    <parent link="shoulder_link"/>
    <child link="upper_arm_link"/>
    <origin rpy="1.571497675314641 0 3.984082409699119e-06" xyz="0.000182214465989093 0 0"/>
    <axis xyz="0 0 1"/>
    <limit effort="150.0" lower="-4.71238898038469" upper="1.5707963267948966" velocity="0.8726646259971648"/>
    <dynamics damping="0" friction="0"/>
  </joint>
  <joint name="elbow_joint" type="revolute">
    <parent link="upper_arm_link"/>
    <child link="forearm_link"/>
    <origin rpy="0.001047955138343435 -0.001169294971873293 4.425876114326764e-06" xyz="-0.4249817627044961 0 0"/>
    <axis xyz="0 0 1"/>
    <limit effort="150.0" lower="-3.141592653589793" upper="3.141592653589793" velocity="0.8726646259971648"/>
    <dynamics damping="0" friction="0"/>
  </joint>
  <joint name="wrist_1_joint" type="revolute">
    <parent link="forearm_link"/>
    <child link="wrist_1_link"/>
    <origin rpy="0.007478051175221102 0.0004306268150536495 -7.033675124629128e-06" xyz="-0.3921666446509172 -0.0009975307642066673 0.1333919956383524"/>
    <axis xyz="0 0 1"/>
    <limit effort="28.0" lower="-3.141592653589793" upper="3.141592653589793" velocity="3.141592653589793"/>
    <dynamics damping="0" friction="0"/>
  </joint>
  <joint name="wrist_2_joint" type="revolute">
    <parent link="wrist_1_link"/>
    <child link="wrist_2_link"/>
    <origin rpy="1.568983665155081 0 -1.012784017796753e-07" xyz="2.398975523480517e-06 -0.09959821611958637 0.0001805380634879481"/>
    <axis xyz="0 0 1"/>
    <limit effort="28.0" lower="-3.141592653589793" upper="3.141592653589793" velocity="3.141592653589793"/>
    <dynamics damping="0" friction="0"/>
  </joint>
  <joint name="wrist_3_joint" type="revolute">
    <parent link="wrist_2_link"/>
    <child link="wrist_3_link"/>
    <origin rpy="1.571939944776703 3.141592653589793 -3.1415926331844" xyz="7.603964784130673e-05 0.09950302422228456 0.0001137934973534554"/>
    <axis xyz="0 0 1"/>
    <limit effort="28.0" lower="-3.141592653589793" upper="3.141592653589793" velocity="3.141592653589793"/>
    <dynamics damping="0" friction="0"/>
  </joint>
  <link name="ft_frame"/>
  <joint name="wrist_3_link-ft_frame" type="fixed">
    <parent link="wrist_3_link"/>
    <child link="ft_frame"/>
    <origin rpy="3.141592653589793 0 0" xyz="0 0 0"/>
  </joint>
  <!-- ROS-Industrial 'base' frame - base_link to UR 'Base' Coordinates transform -->
  <link name="base"/>
  <joint name="base_link-base_fixed_joint" type="fixed">
    <!-- Note the rotation over Z of pi radians - as base_link is REP-103
           aligned (i.e., has X+ forward, Y+ left and Z+ up), this is needed
           to correctly align 'base' with the 'Base' coordinate system of
           the UR controller.
      -->
    <origin rpy="0 0 3.141592653589793" xyz="0 0 0"/>
    <parent link="base_link"/>
    <child link="base"/>
  </joint>
  <!-- ROS-Industrial 'flange' frame - attachment point for EEF models -->
  <link name="flange"/>
  <joint name="wrist_3-flange" type="fixed">
    <parent link="wrist_3_link"/>
    <child link="flange"/>
    <origin rpy="0 -1.5707963267948966 -1.5707963267948966" xyz="0 0 0"/>
  </joint>
  <!-- ROS-Industrial 'tool0' frame - all-zeros tool frame -->
  <link name="tool0"/>
  <joint name="flange-tool0" type="fixed">
    <!-- default toolframe - X+ left, Y+ up, Z+ front -->
    <origin rpy="1.5707963267948966 0 1.5707963267948966" xyz="0 0 0"/>
    <parent link="flange"/>
    <child link="tool0"/>
  </joint>
  <joint name="tool_changer_joint" type="fixed">
    <!-- The parent link must be read from the robot model it is attached to. -->
    <parent link="tool0"/>
    <child link="tool_changer_link"/>
    <origin rpy="0 0 0" xyz="0 0 0"/>
  </joint>
  <link name="tool_changer_link">
    <visual>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://picknik_accessories/descriptions/brackets/mtc_ur_tool_changer/mtc_ur3510_ur_toolchanger.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://picknik_accessories/descriptions/brackets/mtc_ur_tool_changer/mtc_ur3510_ur_toolchanger_collision.stl"/>
      </geometry>
    </collision>
  </link>
  <joint name="tool_side_joint" type="fixed">
    <parent link="tool_changer_link"/>
    <child link="tool_changer_tool0"/>
    <origin rpy="0 0 0" xyz="0 0 0.0331"/>
  </joint>
  <link name="tool_changer_tool0"/>
  <joint name="realsense_camera_adapter_joint" type="fixed">
    <!-- The parent link must be read from the robot model it is attached to. -->
    <parent link="tool_changer_tool0"/>
    <child link="realsense_camera_adapter_link"/>
    <origin rpy="0 0 0" xyz="0 0 0"/>
  </joint>
  <link name="realsense_camera_adapter_link">
    <visual>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://picknik_accessories/descriptions/brackets/ur_realsense_camera_adapter/picknik_ur5_realsense_camera_adapter_rev2.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://picknik_accessories/descriptions/brackets/ur_realsense_camera_adapter/picknik_ur5_realsense_camera_adapter_rev2_collision.stl"/>
      </geometry>
    </collision>
  </link>
  <link name="realsense_camera_adapter_tool0"/>
  <joint name="realsense_camera_adapter_tool0_joint" type="fixed">
    <parent link="realsense_camera_adapter_link"/>
    <child link="realsense_camera_adapter_tool0"/>
    <origin rpy="0 0 0" xyz="0 0 0.007"/>
  </joint>
  <link name="d415_mount_link"/>
  <joint name="d415_mount_joint" type="fixed">
    <parent link="realsense_camera_adapter_link"/>
    <child link="d415_mount_link"/>
    <origin rpy="0 -1.4835298641951802 1.5707963267948966" xyz="0 -0.067 0.0171"/>
  </joint>
  <link name="l515_mount_link"/>
  <joint name="l515_mount_joint" type="fixed">
    <parent link="realsense_camera_adapter_link"/>
    <child link="l515_mount_link"/>
    <origin rpy="0 -1.4835298641951802 1.5707963267948966" xyz="0 -0.0406 0.0173"/>
  </joint>
  <!-- camera body, with origin at bottom screw mount -->
  <joint name="wrist_mounted_camera_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0 0"/>
    <parent link="d415_mount_link"/>
    <child link="wrist_mounted_camera_bottom_screw_frame"/>
  </joint>
  <link name="wrist_mounted_camera_bottom_screw_frame"/>
  <joint name="wrist_mounted_camera_link_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0.02 0.0115"/>
    <parent link="wrist_mounted_camera_bottom_screw_frame"/>
    <child link="wrist_mounted_camera_link"/>
  </joint>
  <link name="wrist_mounted_camera_link">
    <visual>
      <origin rpy="1.5707963267948966 0 1.5707963267948966" xyz="0.00987 -0.02 0"/>
      <geometry>
        <mesh filename="package://realsense2_description/meshes/d415.stl"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 0" xyz="0 -0.02 0"/>
      <geometry>
        <box size="0.02005 0.099 0.023"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.072"/>
      <!-- The following are not reliable values, and should not be used for modeling -->
      <origin xyz="0 0 0"/>
      <inertia ixx="0.003881243" ixy="0.0" ixz="0.0" iyy="0.000498940" iyz="0.0" izz="0.003879257"/>
    </inertial>
  </link>
  <!-- camera depth joints and links -->
  <joint name="wrist_mounted_camera_depth_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0 0"/>
    <parent link="wrist_mounted_camera_link"/>
    <child link="wrist_mounted_camera_depth_frame"/>
  </joint>
  <link name="wrist_mounted_camera_depth_frame"/>
  <joint name="wrist_mounted_camera_depth_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="wrist_mounted_camera_depth_frame"/>
    <child link="wrist_mounted_camera_depth_optical_frame"/>
  </joint>
  <link name="wrist_mounted_camera_depth_optical_frame"/>
  <!-- camera left IR joints and links -->
  <joint name="wrist_mounted_camera_infra1_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0.0 0"/>
    <parent link="wrist_mounted_camera_link"/>
    <child link="wrist_mounted_camera_infra1_frame"/>
  </joint>
  <link name="wrist_mounted_camera_infra1_frame"/>
  <joint name="wrist_mounted_camera_infra1_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="wrist_mounted_camera_infra1_frame"/>
    <child link="wrist_mounted_camera_infra1_optical_frame"/>
  </joint>
  <link name="wrist_mounted_camera_infra1_optical_frame"/>
  <!-- camera right IR joints and links -->
  <joint name="wrist_mounted_camera_infra2_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 -0.055 0"/>
    <parent link="wrist_mounted_camera_link"/>
    <child link="wrist_mounted_camera_infra2_frame"/>
  </joint>
  <link name="wrist_mounted_camera_infra2_frame"/>
  <joint name="wrist_mounted_camera_infra2_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="wrist_mounted_camera_infra2_frame"/>
    <child link="wrist_mounted_camera_infra2_optical_frame"/>
  </joint>
  <link name="wrist_mounted_camera_infra2_optical_frame"/>
  <!-- camera color joints and links -->
  <joint name="wrist_mounted_camera_color_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0.015 0"/>
    <parent link="wrist_mounted_camera_link"/>
    <child link="wrist_mounted_camera_color_frame"/>
  </joint>
  <link name="wrist_mounted_camera_color_frame"/>
  <joint name="wrist_mounted_camera_color_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="wrist_mounted_camera_color_frame"/>
    <child link="wrist_mounted_camera_color_optical_frame"/>
  </joint>
  <link name="wrist_mounted_camera_color_optical_frame"/>
  <joint name="ur_to_robotiq_joint" type="fixed">
    <!-- The parent link must be read from the robot model it is attached to. -->
    <parent link="realsense_camera_adapter_tool0"/>
    <child link="ur_to_robotiq_link"/>
    <origin rpy="0 0 0" xyz="0 0 0"/>
  </joint>
  <link name="ur_to_robotiq_link">
    <visual>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/ur_to_robotiq_adapter.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 0" xyz="0 0 0"/>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/ur_to_robotiq_adapter.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.01"/>
      <origin xyz="0 0 0"/>
      <inertia ixx="0.000044" ixy="0.0" ixz="0.0" iyy="0.000027" iyz="0.0" izz="0.000027"/>
    </inertial>
  </link>
  <joint name="gripper_side_joint" type="fixed">
    <parent link="ur_to_robotiq_link"/>
    <child link="gripper_mount_link"/>
    <origin rpy="0 -1.5707963267948966 1.5707963267948966" xyz="0 0 0.011"/>
  </joint>
  <link name="gripper_mount_link"/>
  <joint name="robotiq_85_base_joint" type="fixed">
    <parent link="gripper_mount_link"/>
    <child link="robotiq_85_base_link"/>
    <origin rpy="0 0 0" xyz="0 0 0"/>
  </joint>
  <link name="robotiq_85_base_link">
    <!-- Visual aid for grasping -->
    <visual>
      <origin rpy="0 0 0" xyz="0.16 0 0"/>
      <geometry>
        <sphere radius="0.007"/>
      </geometry>
      <material name="">
        <color rgba="0 1.0 0 1.0"/>
      </material>
    </visual>
    <!-- Gripper body -->
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_base_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_base_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.636951"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000380" ixy="0.000000" ixz="0.000000" iyx="0.000000" iyy="0.001110" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.001171"/>
    </inertial>
  </link>
  <joint name="robotiq_85_left_knuckle_joint" type="revolute">
    <parent link="robotiq_85_base_link"/>
    <child link="robotiq_85_left_knuckle_link"/>
    <axis xyz="0 0 1"/>
    <origin rpy="3.141592653589793 0.0 0.0" xyz="0.05490451627 0.03060114443 0.0"/>
    <limit effort="50" lower="0.0" upper="0.804" velocity="0.5"/>
    <dynamics damping="100.0"/>
  </joint>
  <joint name="robotiq_85_right_knuckle_joint" type="continuous">
    <parent link="robotiq_85_base_link"/>
    <child link="robotiq_85_right_knuckle_link"/>
    <axis xyz="0 0 1"/>
    <origin rpy="0.0 0.0 0.0" xyz="0.05490451627 -0.03060114443 0.0"/>
    <limit effort="50" lower="-3.14" upper="3.14" velocity="100.0"/>
    <mimic joint="robotiq_85_left_knuckle_joint"/>
    <dynamics damping="100.0"/>
  </joint>
  <link name="robotiq_85_left_knuckle_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_knuckle_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_knuckle_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.018491"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000009" ixy="-0.000001" ixz="0.000000" iyx="-0.000001" iyy="0.000001" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000010"/>
    </inertial>
  </link>
  <link name="robotiq_85_right_knuckle_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_knuckle_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_knuckle_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.018491"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000009" ixy="-0.000001" ixz="0.000000" iyx="-0.000001" iyy="0.000001" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000010"/>
    </inertial>
  </link>
  <joint name="robotiq_85_left_finger_joint" type="fixed">
    <parent link="robotiq_85_left_knuckle_link"/>
    <child link="robotiq_85_left_finger_link"/>
    <origin rpy="0 0 0" xyz="-0.00408552455 -0.03148604435 0.0"/>
  </joint>
  <joint name="robotiq_85_right_finger_joint" type="fixed">
    <parent link="robotiq_85_right_knuckle_link"/>
    <child link="robotiq_85_right_finger_link"/>
    <origin rpy="0 0 0" xyz="-0.00408552455 -0.03148604435 0.0"/>
  </joint>
  <link name="robotiq_85_left_finger_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_finger_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_finger_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.027309"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000003" ixy="-0.000002" ixz="0.000000" iyx="-0.000002" iyy="0.000021" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000020"/>
    </inertial>
  </link>
  <link name="robotiq_85_right_finger_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_finger_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_finger_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.027309"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000003" ixy="-0.000002" ixz="0.000000" iyx="-0.000002" iyy="0.000021" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000020"/>
    </inertial>
  </link>
  <joint name="robotiq_85_left_inner_knuckle_joint" type="continuous">
    <parent link="robotiq_85_base_link"/>
    <child link="robotiq_85_left_inner_knuckle_link"/>
    <axis xyz="0 0 1"/>
    <origin rpy="3.141592653589793 0.0 0.0" xyz="0.06142 0.0127 0"/>
    <limit effort="0.1" lower="-3.14" upper="3.14" velocity="100.0"/>
    <mimic joint="robotiq_85_left_knuckle_joint" offset="0"/>
  </joint>
  <link name="robotiq_85_left_inner_knuckle_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_inner_knuckle_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_inner_knuckle_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.029951"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000039" ixy="0.000000" ixz="0.000000" iyx="0.000000" iyy="0.000005" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000035"/>
    </inertial>
  </link>
  <link name="robotiq_85_right_inner_knuckle_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_inner_knuckle_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_inner_knuckle_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.029951"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000039" ixy="0.000000" ixz="0.000000" iyx="0.000000" iyy="0.000005" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000035"/>
    </inertial>
  </link>
  <joint name="robotiq_85_left_finger_tip_joint" type="continuous">
    <parent link="robotiq_85_left_inner_knuckle_link"/>
    <child link="robotiq_85_left_finger_tip_link"/>
    <axis xyz="0 0 1"/>
    <origin rpy="0.0 0.0 0.0" xyz="0.04303959807 -0.03759940821 0.0"/>
    <limit effort="0.1" lower="-3.14" upper="3.14" velocity="100.0"/>
    <mimic joint="robotiq_85_left_knuckle_joint" multiplier="-1"/>
  </joint>
  <joint name="robotiq_85_right_finger_tip_joint" type="continuous">
    <parent link="robotiq_85_right_inner_knuckle_link"/>
    <child link="robotiq_85_right_finger_tip_link"/>
    <axis xyz="0 0 1"/>
    <origin rpy="0.0 0.0 0.0" xyz="0.04303959807 -0.03759940821  0.0"/>
    <limit effort="0.1" lower="-3.14" upper="3.14" velocity="100.0"/>
    <mimic joint="robotiq_85_left_knuckle_joint" multiplier="-1"/>
  </joint>
  <link name="robotiq_85_left_finger_tip_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_finger_tip_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_finger_tip_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.019555"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000002" ixy="0.000000" ixz="0.000000" iyx="0.000000" iyy="0.000005" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000006"/>
    </inertial>
  </link>
  <link name="robotiq_85_right_finger_tip_link">
    <visual>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/visual/robotiq_85_finger_tip_link.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="package://robotiq_85_description/meshes/collision/robotiq_85_finger_tip_link.stl"/>
      </geometry>
    </collision>
    <inertial>
      <mass value="0.019555"/>
      <origin xyz="0.0 0.0 0.0"/>
      <inertia ixx="0.000002" ixy="0.000000" ixz="0.000000" iyx="0.000000" iyy="0.000005" iyz="0.000000" izx="0.000000" izy="0.000000" izz="0.000006"/>
    </inertial>
  </link>
  <link name="grasp_link"/>
  <joint name="grasp_link_joint" type="fixed">
    <parent link="gripper_mount_link"/>
    <child link="grasp_link"/>
    <origin rpy="1.5707963267948966 0.0 1.5707963267948966" xyz="0.134 0.0 0.0"/>
  </joint>
  <!-- Moveit Studio requires a link name called manual_grasp_link to perform quick tasks such as inspect surface -->
  <link name="manual_grasp_link"/>
  <joint name="manual_grasp_joint" type="fixed">
    <parent link="grasp_link"/>
    <child link="manual_grasp_link"/>
    <origin rpy="0.08726646259971647 0 0" xyz="0 0 0"/>
  </joint>
  <link name="table">
    <visual>
      <geometry>
        <box size="2.0 2.0 0.25 "/>
      </geometry>
      <material name="blue">
        <color rgba="0 0 1 0.3"/>
      </material>
    </visual>
    <collision>
      <geometry>
        <box size="2.0 2.0  0.25 "/>
      </geometry>
    </collision>
  </link>
  <joint name="table" type="fixed">
    <parent link="base"/>
    <child link="table"/>
    <origin rpy="0 0 0" xyz="1.4 0 -0.125"/>
  </joint>
  <!-- Camera External -->
  <link name="external_camera_link"/>
  <joint name="external_camera_joint" type="fixed">
    <parent link="world"/>
    <child link="external_camera_link"/>
    <origin rpy="0.0 0.4 0" xyz="-0.3 0.3 1.0"/>
  </joint>
  <!-- camera body, with origin at bottom screw mount -->
  <joint name="scene_camera_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0 0"/>
    <parent link="external_camera_link"/>
    <child link="scene_camera_bottom_screw_frame"/>
  </joint>
  <link name="scene_camera_bottom_screw_frame"/>
  <joint name="scene_camera_link_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0.010600000000000002 0.0175 0.0125"/>
    <parent link="scene_camera_bottom_screw_frame"/>
    <child link="scene_camera_link"/>
  </joint>
  <link name="scene_camera_link">
    <visual>
      <!-- the mesh origin is at front plate in between the two infrared camera axes -->
      <origin rpy="1.5707963267948966 0 1.5707963267948966" xyz="0.0043 -0.0175 0"/>
      <geometry>
        <!-- <box size="${d435_cam_width} ${d435_cam_height} ${d435_cam_depth}"/> -->
        <mesh filename="package://realsense2_description/meshes/d435.dae"/>
      </geometry>
    </visual>
    <collision>
      <origin rpy="0 0 0" xyz="0 -0.0175 0"/>
      <geometry>
        <box size="0.02505 0.09 0.025"/>
      </geometry>
    </collision>
    <inertial>
      <!-- The following are not reliable values, and should not be used for modeling -->
      <mass value="0.072"/>
      <origin xyz="0 0 0"/>
      <inertia ixx="0.003881243" ixy="0.0" ixz="0.0" iyy="0.000498940" iyz="0.0" izz="0.003879257"/>
    </inertial>
  </link>
  <!-- camera depth joints and links -->
  <joint name="scene_camera_depth_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0 0"/>
    <parent link="scene_camera_link"/>
    <child link="scene_camera_depth_frame"/>
  </joint>
  <link name="scene_camera_depth_frame"/>
  <joint name="scene_camera_depth_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="scene_camera_depth_frame"/>
    <child link="scene_camera_depth_optical_frame"/>
  </joint>
  <link name="scene_camera_depth_optical_frame"/>
  <!-- camera left IR joints and links -->
  <joint name="scene_camera_infra1_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0.0 0"/>
    <parent link="scene_camera_link"/>
    <child link="scene_camera_infra1_frame"/>
  </joint>
  <link name="scene_camera_infra1_frame"/>
  <joint name="scene_camera_infra1_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="scene_camera_infra1_frame"/>
    <child link="scene_camera_infra1_optical_frame"/>
  </joint>
  <link name="scene_camera_infra1_optical_frame"/>
  <!-- camera right IR joints and links -->
  <joint name="scene_camera_infra2_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 -0.05 0"/>
    <parent link="scene_camera_link"/>
    <child link="scene_camera_infra2_frame"/>
  </joint>
  <link name="scene_camera_infra2_frame"/>
  <joint name="scene_camera_infra2_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="scene_camera_infra2_frame"/>
    <child link="scene_camera_infra2_optical_frame"/>
  </joint>
  <link name="scene_camera_infra2_optical_frame"/>
  <!-- camera color joints and links -->
  <joint name="scene_camera_color_joint" type="fixed">
    <origin rpy="0 0 0" xyz="0 0.015 0"/>
    <parent link="scene_camera_link"/>
    <child link="scene_camera_color_frame"/>
  </joint>
  <link name="scene_camera_color_frame"/>
  <joint name="scene_camera_color_optical_joint" type="fixed">
    <origin rpy="-1.5707963267948966 0 -1.5707963267948966" xyz="0 0 0"/>
    <parent link="scene_camera_color_frame"/>
    <child link="scene_camera_color_optical_frame"/>
  </joint>
  <link name="scene_camera_color_optical_frame"/>
</robot>

          

function(generate_fast_forward_kinematics_library URDF_FILE ROOT_LINK TIP_LINK)

    find_package(Python REQUIRED COMPONENTS Interpreter)
    if (Python_Interpreter_FOUND)
        message(STATUS "Python executable: ${Python_EXECUTABLE}")
    else ()
        message(FATAL_ERROR "Python interpreter not found.")
    endif ()

    execute_process(
            COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/get_num_joints.py
            ${URDF_FILE}
            ${ROOT_LINK}
            ${TIP_LINK}
            OUTPUT_VARIABLE FAST_FK_NUMBER_OF_JOINTS
            OUTPUT_STRIP_TRAILING_WHITESPACE
            COMMAND_ECHO STDOUT
    )

    include(ExternalProject)
    ExternalProject_Add(
            LBFGSpp
            PREFIX ${CMAKE_BINARY_DIR}/LBFGSpp
            GIT_REPOSITORY https://github.com/yixuan/LBFGSpp.git
            GIT_TAG v0.3.0
            CMAKE_ARGS -DCMAKE_INSTALL_PREFIX=${CMAKE_BINARY_DIR}
    )
    ExternalProject_Get_Property(LBFGSpp source_dir)
    set(LBFGSppIncludeDir ${source_dir}/include)

    add_custom_command(
            OUTPUT forward_kinematics_lib.cpp
            COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/robot_gen.py
            ${URDF_FILE}
            ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
            ${CMAKE_CURRENT_BINARY_DIR}/forward_kinematics_lib.cpp
            ${ROOT_LINK}
            ${TIP_LINK}
            DEPENDS ${URDF_FILE} ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
            COMMENT
            "Running `${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/robot_gen.py
                      ${URDF_FILE}
                      ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
                      ${CMAKE_CURRENT_BINARY_DIR}/forward_kinematics_test.cpp
                      ${ROOT_LINK}
                      ${TIP_LINK}`"
            VERBATIM
    )
    add_custom_target(code_generation DEPENDS forward_kinematics_lib.cpp)

    add_library(fast_forward_kinematics_library SHARED forward_kinematics_lib.cpp)
    add_dependencies(fast_forward_kinematics_library code_generation)
    add_dependencies(fast_forward_kinematics_library LBFGSpp)

    find_package(Eigen3 3.3 NO_MODULE)
    target_compile_definitions(fast_forward_kinematics_library PUBLIC "${FAST_FK_NUMBER_OF_JOINTS}")
    target_include_directories(fast_forward_kinematics_library PUBLIC ${CMAKE_SOURCE_DIR}/include)
    target_include_directories(fast_forward_kinematics_library PUBLIC ${LBFGSppIncludeDir})
    target_link_libraries(fast_forward_kinematics_library PUBLIC Eigen3::Eigen)

endfunction()
















#
          

from jinja2 import Template
from urdf_parser_py import urdf

import argparse
import numpy as np


def run():
    parser = argparse.ArgumentParser()
    parser.add_argument('urdf_file')
    parser.add_argument('fk_template')
    parser.add_argument('fk_output_file')
    parser.add_argument('root_link_name')
    parser.add_argument('tip_link_name')
    args = parser.parse_args()

    root_link_name = args.root_link_name
    tip_link_name = args.tip_link_name
    with open(args.urdf_file) as f:
        robot = urdf.Robot.from_xml_string(f.read())

    def get_transform(joint):
        rpy = joint.origin.rpy
        xyz = joint.origin.xyz
        T = np.eye(4)

        yaw = np.array([[np.cos(rpy[2]), np.sin(rpy[2]), 0],
                        [-np.sin(rpy[2]), np.cos(rpy[2]), 0],
                        [0, 0, 1]])

        pitch = np.array([[np.cos(rpy[1]), 0, np.sin(rpy[1])],
                          [0, 1, 0],
                          [-np.sin(rpy[1]), 0, np.cos(rpy[1])]])

        roll = np.array([[1, 0, 0],
                         [0, np.cos(rpy[0]), np.sin(rpy[0])],
                         [0, -np.sin(rpy[0]), np.cos(rpy[0])]])
        T[:3, :3] = yaw @ pitch @ roll
        T[:3, 3] = xyz

        return T

    rotations = []
    offsets = []
    types = []
    joint_names = []
    while tip_link_name != root_link_name:
        tip_joint_name, tip_link_name = robot.parent_map[tip_link_name]
        joint_names.append(tip_joint_name)

    joint_names.reverse()

    T_fixed = np.eye(4)
    for joint_name in joint_names:
        joint = robot.joint_map[joint_name]
        if joint.type == 'fixed':
            T_fixed = T_fixed * get_transform(joint)
        elif joint.type == 'revolute':
            # TODO: need to add joint limits?
            T = get_transform(joint) @ T_fixed
            T_fixed = np.eye(4)
            rotations.append(T[:3, :3])
            offsets.append(T[:3, 3])
            types.append('revolute')
        elif joint.type == 'continuous':
            T = get_transform(joint) @ T_fixed
            T_fixed = np.eye(4)
            rotations.append(T[:3, :3])
            offsets.append(T[:3, 3])
            types.append('revolute')
        elif joint.type == 'prismatic':
            # TODO: need to add joint limits?
            T = get_transform(joint) @ T_fixed
            T_fixed = np.eye(4)
            rotations.append(T[:3, :3])
            offsets.append(T[:3, 3])
            types.append('prismatic')
        else:
            raise Exception(f"joint type {joint.type} in URDF not supported")

    # truncate near zero values
    offsets = [val * (np.abs(val) > 1E-5) for val in offsets]
    rotations = [val * (np.abs(val) > 1E-5) for val in rotations]

    with open(args.fk_template, 'r') as f:
        j2_template = Template(f.read())
        code = j2_template.render({'rotations': rotations, 'offsets': offsets, 'types': types}, trim_blocks=True)

    with open(args.fk_output_file, 'w') as f:
        f.write(code)


if __name__ == "__main__":
    run()

          

#include "fast_kinematics.hpp"

namespace fast_fk::internal {
    // sin(t) cos(t)  px py pz R11, R12, R13...
    constexpr int data_size = 2 + 3 + 9+2; // 16

  void forward_kinematics_internal(float *input_data, std::size_t /*size*/) {

      // row major
      constexpr std::array<float, {{rotations|length}}* 9> R_all = {
              {%- for rotation in rotations %}
              {%- for row in rotation %}
              {%- for val in row %} {{val}},
              {%- endfor %}
              {% endfor %}
              {%- endfor -%}
      };

      constexpr std::array<float, {{rotations|length}}* 3> offset_all = {
              {%- for offset in offsets %}
              {%- for val in offset %} {{val}},
              {%- endfor %}
              {% endfor -%}
      };

      float tmp1;
      float tmp2;
      float tmp3;

      {% for type in types %}
      {
          constexpr std::size_t ind = {{loop.index0}};

          // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
          const float &R11_fixed = R_all[ind * 9 + 0];
          const float &R21_fixed = R_all[ind * 9 + 1];
          const float &R31_fixed = R_all[ind * 9 + 2];
          const float &R12_fixed = R_all[ind * 9 + 3];
          const float &R22_fixed = R_all[ind * 9 + 4];
          const float &R32_fixed = R_all[ind * 9 + 5];
          const float &R13_fixed = R_all[ind * 9 + 6];
          const float &R23_fixed = R_all[ind * 9 + 7];
          const float &R33_fixed = R_all[ind * 9 + 8];

          // offset is the offset of joint_i+1 in frame joint_i
          const float &offset_x = offset_all[ind * 3 + 0];
          const float &offset_y = offset_all[ind * 3 + 1];
          const float &offset_z = offset_all[ind * 3 + 2];

          // R is the rotation from joint_i+1 in the base frame (row major)
          float &R11 = input_data[ind * data_size + 5];
          float &R12 = input_data[ind * data_size + 6];
          float &R13 = input_data[ind * data_size + 7];
          float &R21 = input_data[ind * data_size + 8];
          float &R22 = input_data[ind * data_size + 9];
          float &R23 = input_data[ind * data_size + 10];
          float &R31 = input_data[ind * data_size + 11];
          float &R32 = input_data[ind * data_size + 12];
          float &R33 = input_data[ind * data_size + 13];

          {% if type == 'revolute' %}
          const float &sin_theta = input_data[ind * data_size + 0];
          const float &cos_theta = input_data[ind * data_size + 1];
          {% if loop.index0 == 0 %}
          // apply revolute rotation to R_fixed and store in joint to base rotation R
          R11 = R11_fixed * cos_theta + R12_fixed * sin_theta;
          R12 = -R11_fixed * sin_theta + R12_fixed * cos_theta;
          R13 = R13_fixed;
          R21 = R21_fixed * cos_theta + R22_fixed * sin_theta;
          R22 = -R21_fixed * sin_theta + R22_fixed * cos_theta;
          R23 = R23_fixed;
          R31 = R31_fixed * cos_theta + R32_fixed * sin_theta;
          R32 = -R31_fixed * sin_theta + R32_fixed * cos_theta;
          R33 = R33_fixed;

          // rotate offset to be in base frame
          input_data[ind * data_size + 2] = offset_x;
          input_data[ind * data_size + 3] = offset_y;
          input_data[ind * data_size + 4] = offset_z;

          {% else %}
          // R_old is the rotation from joint_i in base frame (column major)
          float &R11_old = input_data[(ind - 1) * data_size + 5];
          float &R12_old = input_data[(ind - 1) * data_size + 6];
          float &R13_old = input_data[(ind - 1) * data_size + 7];
          float &R21_old = input_data[(ind - 1) * data_size + 8];
          float &R22_old = input_data[(ind - 1) * data_size + 9];
          float &R23_old = input_data[(ind - 1) * data_size + 10];
          float &R31_old = input_data[(ind - 1) * data_size + 11];
          float &R32_old = input_data[(ind - 1) * data_size + 12];
          float &R33_old = input_data[(ind - 1) * data_size + 13];

          // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
          float &R11_tmp = input_data[ind * data_size + 5];
          float &R12_tmp = input_data[ind * data_size + 6];
          float &R13_tmp = input_data[ind * data_size + 7];
          float &R21_tmp = input_data[ind * data_size + 8];
          float &R22_tmp = input_data[ind * data_size + 9];
          float &R23_tmp = input_data[ind * data_size + 10];
          float &R31_tmp = input_data[ind * data_size + 11];
          float &R32_tmp = input_data[ind * data_size + 12];
          float &R33_tmp = input_data[ind * data_size + 13];

          // apply revolute rotation to R11_fixed and store in R11_tmp
          R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
          R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
          R13_tmp = R13_fixed;
          R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
          R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
          R23_tmp = R23_fixed;
          R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
          R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
          R33_tmp = R33_fixed;

          // apply R11_old rotation R11_tmp, to make it in base frame
          tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
          tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
          tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
          R11 = tmp1;
          R21 = tmp2;
          R31 = tmp3;
          tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
          tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
          tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
          R12 = tmp1;
          R22 = tmp2;
          R32 = tmp3;
          tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
          tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
          tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
          R13 = tmp1;
          R23 = tmp2;
          R33 = tmp3;

          // p_old is the position of joint_i in base frame
          const float &px_old = input_data[(ind - 1) * data_size + 2];
          const float &py_old = input_data[(ind - 1) * data_size + 3];
          const float &pz_old = input_data[(ind - 1) * data_size + 4];

          // rotate offset to be in base frame and add p_old
          tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
          tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
          tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
          input_data[ind * data_size + 2] = tmp1;
          input_data[ind * data_size + 3] = tmp2;
          input_data[ind * data_size + 4] = tmp3;
          {% endif %}

          {% elif type == 'prismatic' %}
          {% if loop.index0 == 0 %}
          // copy fixed rotation into joint to base rotation
          R11 = R11_fixed;
          R12 = R12_fixed;
          R13 = R13_fixed;
          R21 = R21_fixed;
          R22 = R22_fixed;
          R23 = R23_fixed;
          R31 = R31_fixed;
          R32 = R32_fixed;
          R33 = R33_fixed;

          // add offset and prismatic actuation in base frame
          const float &linear = input_data[ind * data_size + 0];
          input_data[ind * data_size + 2] = R13_fixed * linear + offset_x;
          input_data[ind * data_size + 3] = R23_fixed * linear + offset_y;
          input_data[ind * data_size + 4] = R33_fixed * linear + offset_z;

          {% else %}
          // R_old is the rotation from joint_i in base frame (row major)
          float &R11_old = input_data[(ind - 1) * data_size + 5];
          float &R12_old = input_data[(ind - 1) * data_size + 6];
          float &R13_old = input_data[(ind - 1) * data_size + 7];
          float &R21_old = input_data[(ind - 1) * data_size + 8];
          float &R22_old = input_data[(ind - 1) * data_size + 9];
          float &R23_old = input_data[(ind - 1) * data_size + 10];
          float &R31_old = input_data[(ind - 1) * data_size + 11];
          float &R32_old = input_data[(ind - 1) * data_size + 12];
          float &R33_old = input_data[(ind - 1) * data_size + 13];

          // apply R11_old rotation R11_fixed, to make it in base frame
          tmp1 = R11_old * R11_fixed + R12_old * R21_fixed + R13_old * R31_fixed;
          tmp2 = R21_old * R11_fixed + R22_old * R21_fixed + R23_old * R31_fixed;
          tmp3 = R31_old * R11_fixed + R32_old * R21_fixed + R33_old * R31_fixed;
          R11 = tmp1;
          R21 = tmp2;
          R31 = tmp3;
          tmp1 = R11_old * R12_fixed + R12_old * R22_fixed + R13_old * R32_fixed;
          tmp2 = R21_old * R12_fixed + R22_old * R22_fixed + R23_old * R32_fixed;
          tmp3 = R31_old * R12_fixed + R32_old * R22_fixed + R33_old * R32_fixed;
          R12 = tmp1;
          R22 = tmp2;
          R32 = tmp3;
          tmp1 = R11_old * R13_fixed + R12_old * R23_fixed + R13_old * R33_fixed;
          tmp2 = R21_old * R13_fixed + R22_old * R23_fixed + R23_old * R33_fixed;
          tmp3 = R31_old * R13_fixed + R32_old * R23_fixed + R33_old * R33_fixed;
          R13 = tmp1;
          R23 = tmp2;
          R33 = tmp3;

          // rotate offset to be in base frame
          tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z;
          tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z;
          tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z;

          // add offset, prismatic actuation, and old position in base frame
          const float &px_old = input_data[(ind - 1) * data_size + 2];
          const float &py_old = input_data[(ind - 1) * data_size + 3];
          const float &pz_old = input_data[(ind - 1) * data_size + 4];
          const float &linear = input_data[ind * data_size + 0];
          input_data[ind * data_size + 2] = tmp1 + R13_old * linear + px_old;
          input_data[ind * data_size + 3] = tmp2 + R23_old * linear + py_old;
          input_data[ind * data_size + 4] = tmp3 + R33_old * linear + pz_old;
          {% endif %}
          {% endif %}
      }
      {% endfor %}
  }

}


namespace fast_fk::internal {
    constexpr float axis_scale = 1;

    float InverseKinematics::operator()(const Eigen::VectorX<float> &q, Eigen::VectorX<float> &grad) {
        // construct input_data from q
        float tmp1;
        float tmp2;

        {% for type in types %}
        {% if type == 'revolute' -%}
        tmp1 = q[{{loop.index0}}];
        tmp2 = q[{{loop.index0}}];
        joint_data[{{loop.index0}}][0] = std::sin(tmp1);
        joint_data[{{loop.index0}}][1] = std::cos(tmp2);
        {%- elif type == 'prismatic' -%}
        joint_data[{{loop.index0}}][0] = q[{{loop.index0}}];
        {%- endif %}
        {%- endfor %}

        float *input_data = joint_data.data()->data();
        internal::forward_kinematics_internal(input_data, data_size * {{rotations | length}});

        Eigen::Vector3<float> target_x_axis_ = target_rot_.block<3, 1>(0, 0);
        Eigen::Vector3<float> target_y_axis_ = target_rot_.block<3, 1>(0, 1);
        Eigen::Vector3<float> target_z_axis_ = target_rot_.block<3, 1>(0, 2);
        Eigen::Vector3<float> target_x_ = target_pose_ + axis_scale * target_x_axis_;
        Eigen::Vector3<float> target_y_ = target_pose_ + axis_scale * target_y_axis_;
        Eigen::Vector3<float> target_z_ = target_pose_ + axis_scale * target_z_axis_;


        Eigen::Vector<float, 3> ee_pose;
        ee_pose << input_data[({{ rotations | length }} - 1) *data_size + 2], input_data[({{ rotations | length }} - 1)
        *data_size + 3], input_data[({{ rotations | length }} - 1) *data_size + 4];

        // x point
        Eigen::Vector<float, 3> x_axis;
        x_axis << input_data[({{ rotations | length }} - 1) *data_size + 5], input_data[({{ rotations | length }} - 1)
        *data_size + 8], input_data[({{ rotations | length }} - 1) *data_size + 11];
        Eigen::Vector<float, 3> current_x = ee_pose + x_axis;

        // y point
        Eigen::Vector<float, 3> y_axis;
        y_axis << input_data[({{ rotations | length }} - 1) *data_size + 6], input_data[({{ rotations | length }} - 1)
        *data_size + 9], input_data[({{ rotations | length }} - 1) *data_size + 12];
        Eigen::Vector<float, 3> current_y = ee_pose + y_axis;

        // z point
        Eigen::Vector<float, 3> z_axis;
        z_axis << input_data[({{ rotations | length }} - 1) *data_size + 7], input_data[({{ rotations | length }} - 1)
        *data_size + 10], input_data[({{ rotations | length }} - 1) *data_size + 13];
        Eigen::Vector<float, 3> current_z = ee_pose + z_axis;

        float fx = 0.0;
        Eigen::Vector3<float> delta_x = target_x_.array() - current_x.array();
        Eigen::Vector3<float> delta_y = target_y_.array() - current_y.array();
        Eigen::Vector3<float> delta_z = target_z_.array() - current_z.array();
        fx += (delta_x.array() * delta_x.array()).sum();
        fx += (delta_y.array() * delta_y.array()).sum();
        fx += (delta_z.array() * delta_z.array()).sum();

        // J =  (delta_x)^2
        //    + (delta_y)^2
        //    + (delta_z)^2

        // dJ/dq = dJ1/dq + dJ2/dq + dJ3/dq

        // dJ1/dq = dJ/ddelta_x * ddelta_x/dq
        // dJ2/dq = dJ/ddelta_y * ddelta_y/dq
        // dJ3/dq = dJ/ddelta_z * ddelta_z/dq

        // dJ/ddelta_x : 1x3 2*(x_target - x_current)
        // dJ/ddelta_y : 1x3 2*(y_target - y_current)
        // dJ/ddelta_z : 1x3 2*(z_target - z_current)

        // ddelta_x/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x x_axis
        // ddelta_y/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x y_axis
        // ddelta_z/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x z_axis


        // gradient
        grad *= 0;
        Eigen::Vector3<float> dJ_ddelta_x = 2 * delta_x;
        Eigen::Vector3<float> dJ_ddelta_y = 2 * delta_y;
        Eigen::Vector3<float> dJ_ddelta_z = 2 * delta_z;

        Eigen::Vector<float, 3> joint_pose;
        Eigen::Vector<float, 3> delta;
        Eigen::Vector<float, 3> joint_axis;
        Eigen::Vector<float, 3> jac_linear;
        Eigen::Vector<float, 3> jac_angular;
        Eigen::Vector3<float> ddelta_x_dq_ind;
        Eigen::Vector3<float> ddelta_y_dq_ind;
        Eigen::Vector3<float> ddelta_z_dq_ind;

        {% for type in types %}
        {
            constexpr std::size_t ind = {{loop.index0}};
            {% if type == 'revolute' %}
            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);

            {% elif type == 'prismatic' %}
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            // jac_linear is the joint axis
            // x point
            grad[ind] -= dJ_ddelta_x.dot(joint_axis);

            // y point
            grad[ind] -= dJ_ddelta_y.dot(joint_axis);

            // z point
            grad[ind] -= dJ_ddelta_z.dot(joint_axis);
            {% endif %}
        }
        {% endfor %}

        return fx;
    }


}
          

#include "fast_kinematics.hpp"

namespace fast_fk::internal {
    // sin(t) cos(t)  px py pz R11, R12, R13...
    constexpr int data_size = 2 + 3 + 9 + 2; // 16

    void forward_kinematics_internal(float *input_data, std::size_t /*size*/) {

        // row major
        constexpr std::array<float, 6 * 9> R_all = { -1.0, -0.0, 0.0,
                                                    0.0, -1.0, 0.0,
                                                    0.0, 0.0, 1.0,
                                                    0.9999999999920636, -0.0, 0.0,
                                                    -0.0, -0.0007013484622411433, 0.9999997540472005,
                                                    0.0, -0.999999754055137, -0.0007013484622467095,
                                                    0.9999993163649181, 0.0, -0.001169289425224364,
                                                    -0.0, 0.9999994508798468, 0.0010479601216713834,
                                                    0.001169294705420077, -0.00104795423012248, 0.9999987672701519,
                                                    0.9999999072555382, -1.0253697707014652e-05, 0.00043056216350321623,
                                                    0.0, 0.9999720394582239, 0.0074779845070412285,
                                                    -0.0004306268017444494, -0.007477980785065483, 0.9999719467884766,
                                                    0.9999999999999949, -0.0, -0.0,
                                                    0.0, 0.001812660647159262, 0.9999983571293345,
                                                    0.0, -0.9999983571293396, 0.0018126606471592713,
                                                    0.9999999999999998, 0.0, -0.0,
                                                    -0.0, 0.0011436177325240319, -0.9999993460690269,
                                                    -0.0, 0.9999993460690271, 0.001143617732524032,
        };

        constexpr std::array<float, 6 * 3> offset_all = { 0.0, 0.0, 0.1625702965797758,
                                                         0.000182214465989093, 0.0, 0.0,
                                                         -0.4249817627044961, 0.0, 0.0,
                                                         -0.3921666446509172, -0.0009975307642066673, 0.1333919956383524,
                                                         0.0, -0.09959821611958637, 0.0001805380634879481,
                                                         7.603964784130673e-05, 0.09950302422228456, 0.0001137934973534554,
        };

        float tmp1;
        float tmp2;
        float tmp3;


        {
            constexpr std::size_t ind = 0;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // apply revolute rotation to R_fixed and store in joint to base rotation R
            R11 = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12 = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13 = R13_fixed;
            R21 = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22 = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23 = R23_fixed;
            R31 = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32 = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33 = R33_fixed;

            // rotate offset to be in base frame
            input_data[ind * data_size + 2] = offset_x;
            input_data[ind * data_size + 3] = offset_y;
            input_data[ind * data_size + 4] = offset_z;




        }

        {
            constexpr std::size_t ind = 1;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 2;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 3;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 4;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 5;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

    }

}


namespace fast_fk::internal {
    constexpr float axis_scale = 1;

    float InverseKinematics::operator()(const Eigen::VectorX<float> &q, Eigen::VectorX<float> &grad) {
        // construct input_data from q
        float tmp1;
        float tmp2;

        tmp1 = q[0];
        tmp2 = q[0];
        joint_data[0][0] = std::sin(tmp1);
        joint_data[0][1] = std::cos(tmp2);
        tmp1 = q[1];
        tmp2 = q[1];
        joint_data[1][0] = std::sin(tmp1);
        joint_data[1][1] = std::cos(tmp2);
        tmp1 = q[2];
        tmp2 = q[2];
        joint_data[2][0] = std::sin(tmp1);
        joint_data[2][1] = std::cos(tmp2);
        tmp1 = q[3];
        tmp2 = q[3];
        joint_data[3][0] = std::sin(tmp1);
        joint_data[3][1] = std::cos(tmp2);
        tmp1 = q[4];
        tmp2 = q[4];
        joint_data[4][0] = std::sin(tmp1);
        joint_data[4][1] = std::cos(tmp2);
        tmp1 = q[5];
        tmp2 = q[5];
        joint_data[5][0] = std::sin(tmp1);
        joint_data[5][1] = std::cos(tmp2);

        float *input_data = joint_data.data()->data();
        internal::forward_kinematics_internal(input_data, data_size * 6);

        Eigen::Vector3<float> target_x_axis_ = target_rot_.block<3, 1>(0, 0);
        Eigen::Vector3<float> target_y_axis_ = target_rot_.block<3, 1>(0, 1);
        Eigen::Vector3<float> target_z_axis_ = target_rot_.block<3, 1>(0, 2);
        Eigen::Vector3<float> target_x_ = target_pose_ + axis_scale * target_x_axis_;
        Eigen::Vector3<float> target_y_ = target_pose_ + axis_scale * target_y_axis_;
        Eigen::Vector3<float> target_z_ = target_pose_ + axis_scale * target_z_axis_;


        Eigen::Vector<float, 3> ee_pose;
        ee_pose << input_data[(6 - 1) *data_size + 2], input_data[(6 - 1)
                                                                  *data_size + 3], input_data[(6 - 1) *data_size + 4];

        // x point
        Eigen::Vector<float, 3> x_axis;
        x_axis << input_data[(6 - 1) *data_size + 5], input_data[(6 - 1)
                                                                 *data_size + 8], input_data[(6 - 1) *data_size + 11];
        Eigen::Vector<float, 3> current_x = ee_pose + x_axis;

        // y point
        Eigen::Vector<float, 3> y_axis;
        y_axis << input_data[(6 - 1) *data_size + 6], input_data[(6 - 1)
                                                                 *data_size + 9], input_data[(6 - 1) *data_size + 12];
        Eigen::Vector<float, 3> current_y = ee_pose + y_axis;

        // z point
        Eigen::Vector<float, 3> z_axis;
        z_axis << input_data[(6 - 1) *data_size + 7], input_data[(6 - 1)
                                                                 *data_size + 10], input_data[(6 - 1) *data_size + 13];
        Eigen::Vector<float, 3> current_z = ee_pose + z_axis;

        float fx = 0.0;
        Eigen::Vector3<float> delta_x = target_x_.array() - current_x.array();
        Eigen::Vector3<float> delta_y = target_y_.array() - current_y.array();
        Eigen::Vector3<float> delta_z = target_z_.array() - current_z.array();
        fx += (delta_x.array() * delta_x.array()).sum();
        fx += (delta_y.array() * delta_y.array()).sum();
        fx += (delta_z.array() * delta_z.array()).sum();

        // J =  (delta_x)^2
        //    + (delta_y)^2
        //    + (delta_z)^2

        // dJ/dq = dJ1/dq + dJ2/dq + dJ3/dq

        // dJ1/dq = dJ/ddelta_x * ddelta_x/dq
        // dJ2/dq = dJ/ddelta_y * ddelta_y/dq
        // dJ3/dq = dJ/ddelta_z * ddelta_z/dq

        // dJ/ddelta_x : 1x3 2*(x_target - x_current)
        // dJ/ddelta_y : 1x3 2*(y_target - y_current)
        // dJ/ddelta_z : 1x3 2*(z_target - z_current)

        // ddelta_x/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x x_axis
        // ddelta_y/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x y_axis
        // ddelta_z/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x z_axis


        // gradient
        grad *= 0;
        Eigen::Vector3<float> dJ_ddelta_x = 2 * delta_x;
        Eigen::Vector3<float> dJ_ddelta_y = 2 * delta_y;
        Eigen::Vector3<float> dJ_ddelta_z = 2 * delta_z;

        Eigen::Vector<float, 3> joint_pose;
        Eigen::Vector<float, 3> delta;
        Eigen::Vector<float, 3> joint_axis;
        Eigen::Vector<float, 3> jac_linear;
        Eigen::Vector<float, 3> jac_angular;
        Eigen::Vector3<float> ddelta_x_dq_ind;
        Eigen::Vector3<float> ddelta_y_dq_ind;
        Eigen::Vector3<float> ddelta_z_dq_ind;


        {
            constexpr std::size_t ind = 0;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 1;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 2;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 3;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 4;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 5;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }


        return fx;
    }


}
          

#include <cstdio>

namespace fast_fk::internal {
    // sin(t) cos(t)  px py pz R11, R12, R13...
    constexpr int data_size = 2 + 3 + 9 + 2; // 16
    constexpr int num_of_joints = {{rotations|length}}; // 16

    __global__ void forward_kinematics_internal_kernel(float *input_data, std::size_t total_size) {
        unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx * data_size * num_of_joints >= total_size) {
            return;
        }
        input_data += (idx * data_size*num_of_joints);

        // row major
        constexpr float R_all[{{rotations|length}}* 9] = {
                {%- for rotation in rotations %}
                {%- for row in rotation %}
                {%- for val in row %} {{val}},
                {%- endfor %}
                {% endfor %}
                {%- endfor -%}
        };

        constexpr float offset_all[{{rotations|length}}* 3] = {
                {%- for offset in offsets %}
                {%- for val in offset %} {{val}},
                {%- endfor %}
                {% endfor -%}
        };

        float tmp1;
        float tmp2;
        float tmp3;

        {% for type in types %}
        {
            constexpr std::size_t ind = {{loop.index0}};

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];

            {% if type == 'revolute' %}
            const float &sin_t = input_data[ind * data_size + 0];
            const float &cos_t = input_data[ind * data_size + 1];
            {% if loop.index0 == 0 %}
            // apply revolute rotation to R_fixed and store in joint to base rotation R
            R11 = R11_fixed * cos_t + R12_fixed * sin_t;
            R12 = -R11_fixed * sin_t + R12_fixed * cos_t;
            R13 = R13_fixed;
            R21 = R21_fixed * cos_t + R22_fixed * sin_t;
            R22 = -R21_fixed * sin_t + R22_fixed * cos_t;
            R23 = R23_fixed;
            R31 = R31_fixed * cos_t + R32_fixed * sin_t;
            R32 = -R31_fixed * sin_t + R32_fixed * cos_t;
            R33 = R33_fixed;

            // rotate offset to be in base frame
            input_data[ind * data_size + 2] = offset_x;
            input_data[ind * data_size + 3] = offset_y;
            input_data[ind * data_size + 4] = offset_z;

            {% else %}
            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_t + R12_fixed * sin_t;
            R12_tmp = -R11_fixed * sin_t + R12_fixed * cos_t;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_t + R22_fixed * sin_t;
            R22_tmp = -R21_fixed * sin_t + R22_fixed * cos_t;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_t + R32_fixed * sin_t;
            R32_tmp = -R31_fixed * sin_t + R32_fixed * cos_t;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;
            {% endif %}

            {% elif type == 'prismatic' %}
            {% if loop.index0 == 0 %}
            // copy fixed rotation into joint to base rotation
            R11 = R11_fixed;
            R12 = R12_fixed;
            R13 = R13_fixed;
            R21 = R21_fixed;
            R22 = R22_fixed;
            R23 = R23_fixed;
            R31 = R31_fixed;
            R32 = R32_fixed;
            R33 = R33_fixed;

            // add offset and prismatic actuation in base frame
            const float &linear = input_data[ind * data_size + 0];
            input_data[ind * data_size + 2] = R13_fixed * linear + offset_x;
            input_data[ind * data_size + 3] = R23_fixed * linear + offset_y;
            input_data[ind * data_size + 4] = R33_fixed * linear + offset_z;

            {% else %}
            // R_old is the rotation from joint_i in base frame (row major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // apply R11_old rotation R11_fixed, to make it in base frame
            tmp1 = R11_old * R11_fixed + R12_old * R21_fixed + R13_old * R31_fixed;
            tmp2 = R21_old * R11_fixed + R22_old * R21_fixed + R23_old * R31_fixed;
            tmp3 = R31_old * R11_fixed + R32_old * R21_fixed + R33_old * R31_fixed;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_fixed + R12_old * R22_fixed + R13_old * R32_fixed;
            tmp2 = R21_old * R12_fixed + R22_old * R22_fixed + R23_old * R32_fixed;
            tmp3 = R31_old * R12_fixed + R32_old * R22_fixed + R33_old * R32_fixed;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_fixed + R12_old * R23_fixed + R13_old * R33_fixed;
            tmp2 = R21_old * R13_fixed + R22_old * R23_fixed + R23_old * R33_fixed;
            tmp3 = R31_old * R13_fixed + R32_old * R23_fixed + R33_old * R33_fixed;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // rotate offset to be in base frame
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z;

            // add offset, prismatic actuation, and old position in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];
            const float &linear = input_data[ind * data_size + 0];
            input_data[ind * data_size + 2] = tmp1 + R13_old * linear + px_old;
            input_data[ind * data_size + 3] = tmp2 + R23_old * linear + py_old;
            input_data[ind * data_size + 4] = tmp3 + R33_old * linear + pz_old;
            {% endif %}
            {% endif %}
        }
        {% endfor %}
    }

    static float *dev_data = nullptr;
    static std::size_t last_size = 0;
    void forward_kinematics_internal(float *input_data, std::size_t size) {
        if (dev_data== nullptr || size > last_size){
            if (dev_data != nullptr){
                cudaFree(dev_data);
            }
            last_size = size;
            cudaMalloc(&dev_data, size * sizeof(float));
        }
        cudaMemcpy(dev_data, input_data, size * sizeof(float), cudaMemcpyKind::cudaMemcpyHostToDevice);
        // Invoke the CUDA kernel
        constexpr size_t block_size = 256;
        size_t num_threads = size/(data_size*num_of_joints);
        size_t grid_size = (num_threads + block_size - 1) / block_size;
        forward_kinematics_internal_kernel<<<grid_size, block_size>>>(dev_data, size);
        cudaDeviceSynchronize();
        cudaMemcpy(input_data, dev_data, size * sizeof(float), cudaMemcpyKind::cudaMemcpyDeviceToHost);

        // Error handling
        cudaError_t cudaerr = cudaDeviceSynchronize();
        if (cudaerr != cudaSuccess)
            std::printf("kernel launch failed with error \"%s\".\n", cudaGetErrorString(cudaerr));
    }
}
          

function(generate_fast_forward_kinematics_library URDF_FILE ROOT_LINK TIP_LINK)

    find_package(Python REQUIRED COMPONENTS Interpreter)
    if (Python_Interpreter_FOUND)
        message(STATUS "Python executable: ${Python_EXECUTABLE}")
    else ()
        message(FATAL_ERROR "Python interpreter not found.")
    endif ()

    execute_process(
            COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/get_num_joints.py
            ${URDF_FILE}
            ${ROOT_LINK}
            ${TIP_LINK}
            OUTPUT_VARIABLE FAST_FK_NUMBER_OF_JOINTS
            OUTPUT_STRIP_TRAILING_WHITESPACE
            COMMAND_ECHO STDOUT
    )

    include(ExternalProject)
    ExternalProject_Add(
            LBFGSpp
            PREFIX ${CMAKE_BINARY_DIR}/LBFGSpp
            GIT_REPOSITORY https://github.com/yixuan/LBFGSpp.git
            GIT_TAG v0.3.0
            CMAKE_ARGS -DCMAKE_INSTALL_PREFIX=${CMAKE_BINARY_DIR}
    )
    ExternalProject_Get_Property(LBFGSpp source_dir)
    set(LBFGSppIncludeDir ${source_dir}/include)

    add_custom_command(
            OUTPUT forward_kinematics_lib.cpp
            COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/robot_gen.py
            ${URDF_FILE}
            ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
            ${CMAKE_CURRENT_BINARY_DIR}/forward_kinematics_lib.cpp
            ${ROOT_LINK}
            ${TIP_LINK}
            DEPENDS ${URDF_FILE} ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
            COMMENT
            "Running `${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/robot_gen.py
                      ${URDF_FILE}
                      ${CMAKE_SOURCE_DIR}/scripts/robot_config.cpp.template
                      ${CMAKE_CURRENT_BINARY_DIR}/forward_kinematics_test.cpp
                      ${ROOT_LINK}
                      ${TIP_LINK}`"
            VERBATIM
    )
    add_custom_target(code_generation DEPENDS forward_kinematics_lib.cpp)

    add_library(fast_forward_kinematics_library SHARED forward_kinematics_lib.cpp)
    add_dependencies(fast_forward_kinematics_library code_generation)
    add_dependencies(fast_forward_kinematics_library LBFGSpp)

    find_package(Eigen3 3.3 NO_MODULE)
    target_compile_definitions(fast_forward_kinematics_library PUBLIC "${FAST_FK_NUMBER_OF_JOINTS}")
    target_include_directories(fast_forward_kinematics_library PUBLIC ${CMAKE_SOURCE_DIR}/include)
    target_include_directories(fast_forward_kinematics_library PUBLIC ${LBFGSppIncludeDir})
    target_link_libraries(fast_forward_kinematics_library PUBLIC Eigen3::Eigen)

endfunction()
















#

The Code


#pragma once

#include <concepts>
#include <iostream>

#include <Eigen/Core>

namespace fk_interface {

    struct IKSolverStats {
        float fx = 0;
        int niter = 0;
        float grad_norm = 0;
        bool success = false;
        std::string what;
    };

    template<typename T>
    concept KinematicInterfaceConstraint = requires(T obj, size_t ind, float value,
                                                    const Eigen::Vector<float, T::get_num_joints()> &values,
                                                    Eigen::Vector<float, T::get_num_joints()> &values_non_const,
                                                    Eigen::VectorX<float> &q_guess) {
        { T::get_num_joints() };
        { obj.set_joint(ind, value) };
        { obj.set_joints(values) };
        { obj.set_joint(ind, value, value) };
        { obj.set_joints(values, values) };
        { obj.get_joint(ind) };
        { obj.get_joints(values_non_const) };
    };


    template<typename T>
    concept ForwardKinematicInterfaceConstraint = requires(T obj, size_t ind, Eigen::Matrix<float, 4, 4> &transform) {
        { obj.get_frame(ind, transform) };
        { obj.forward_kinematics() };
    };

    template<typename T>
    concept InverseInterfaceConstraint = requires(T obj, Eigen::VectorX<float> &q_guess,
                                                  Eigen::Matrix<float, 4, 4> &transform) {
        { obj.inverse_kinematics(transform, q_guess) } -> std::same_as<IKSolverStats>
    };

    template<typename KI> requires KinematicInterfaceConstraint<KI> && ForwardKinematicInterfaceConstraint<KI>
    struct ForwardKinematicsInterface : KI {
    };


    template<typename IK> requires KinematicInterfaceConstraint<IK> && ForwardKinematicInterfaceConstraint<IK> &&
                                   InverseInterfaceConstraint<IK>
    struct InverseKinematicsInterface : IK {
    };

}

          

#include "chrono"
#include "iostream"

#ifdef USE_FAST_KINEMATICS

#include "fast_kinematics.hpp"
using KI = fast_fk::JointData;
#else

#include "kdl_kinematics.hpp"
using KI = kdl_impl::JointData;
#endif

int main(int arc, char **argv) {
    constexpr int iterations = 128 * 128;
    constexpr int multiplier = 128 * MULTIPLIER;
    std::array<Eigen::Vector<float, KI::get_num_joints()>, iterations> rand_values;
    for (auto &rand_val: rand_values) {
        rand_val = Eigen::Vector<float, KI::get_num_joints()>::Random();
    }

    fk_interface::ForwardKinematicsInterface<KI> fk_interface;

    auto start = std::chrono::high_resolution_clock::now();
    for (int k = 0; k < multiplier; k++) {
        for (int i = 0; i < iterations; i++) {
            fk_interface.set_joints(rand_values[i]);
            fk_interface.forward_kinematics();
        }
    }
    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(stop - start);

    std::cout << "Time taken by function: " << (float) duration.count() << " nanoseconds" << std::endl;
    std::cout << "Average: " << ((float) duration.count()) / (iterations * multiplier) << " nanoseconds"
              << std::endl;

    return 0;
}

          

#include "chrono"
#include "iostream"

#ifdef USE_FAST_KINEMATICS

#include "fast_kinematics.hpp"
using IK = fast_fk::JointData;
#else

#include "kdl_kinematics.hpp"
using IK = kdl_impl::JointData;
#endif

int main(int arc, char **argv) {
    constexpr int iterations = 128 * 128 * 5 * MULTIPLIER;

    // get target pose
    Eigen::Matrix<float, 4, 4> tf;
    Eigen::VectorX<float> q_in = Eigen::VectorX<float>::Random(IK::get_num_joints());
    fk_interface::InverseKinematicsInterface<IK> fk_interface;
    fk_interface.set_joints(q_in);
    fk_interface.forward_kinematics();
    fk_interface.get_frame(IK::get_num_joints() - 1, tf);

    auto start = std::chrono::high_resolution_clock::now();
    fk_interface::IKSolverStats stats;
    size_t failed = 0;
    size_t succeeded = 0;
    Eigen::VectorX<float> q;
    for (int ind = 0; ind < iterations; ++ind) {
        q = Eigen::VectorX<float>::Random(IK::get_num_joints());
        stats = fk_interface.inverse_kinematics(tf, q);
        failed += stats.success == 0;
        succeeded += stats.success == 1;
    }
    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(stop - start);

    std::cout << "Time taken by function: " << (double) duration.count() << " microseconds" << std::endl;
    std::cout << "Average: " << ((double) duration.count()) / (iterations) << " microseconds"
              << std::endl;
    std::cout << "Percent success: " << ((double) (100.0 * succeeded) / (failed + succeeded)) << "%"
              << std::endl;

    return 0;
}

          

#include "fast_kinematics.hpp"

namespace fast_fk::internal {
    // sin(t) cos(t)  px py pz R11, R12, R13...
    constexpr int data_size = 2 + 3 + 9 + 2; // 16

    void forward_kinematics_internal(float *input_data, std::size_t /*size*/) {

        // row major
        constexpr std::array<float, 6 * 9> R_all = { -1.0, -0.0, 0.0,
                                                    0.0, -1.0, 0.0,
                                                    0.0, 0.0, 1.0,
                                                    0.9999999999920636, -0.0, 0.0,
                                                    -0.0, -0.0007013484622411433, 0.9999997540472005,
                                                    0.0, -0.999999754055137, -0.0007013484622467095,
                                                    0.9999993163649181, 0.0, -0.001169289425224364,
                                                    -0.0, 0.9999994508798468, 0.0010479601216713834,
                                                    0.001169294705420077, -0.00104795423012248, 0.9999987672701519,
                                                    0.9999999072555382, -1.0253697707014652e-05, 0.00043056216350321623,
                                                    0.0, 0.9999720394582239, 0.0074779845070412285,
                                                    -0.0004306268017444494, -0.007477980785065483, 0.9999719467884766,
                                                    0.9999999999999949, -0.0, -0.0,
                                                    0.0, 0.001812660647159262, 0.9999983571293345,
                                                    0.0, -0.9999983571293396, 0.0018126606471592713,
                                                    0.9999999999999998, 0.0, -0.0,
                                                    -0.0, 0.0011436177325240319, -0.9999993460690269,
                                                    -0.0, 0.9999993460690271, 0.001143617732524032,
        };

        constexpr std::array<float, 6 * 3> offset_all = { 0.0, 0.0, 0.1625702965797758,
                                                         0.000182214465989093, 0.0, 0.0,
                                                         -0.4249817627044961, 0.0, 0.0,
                                                         -0.3921666446509172, -0.0009975307642066673, 0.1333919956383524,
                                                         0.0, -0.09959821611958637, 0.0001805380634879481,
                                                         7.603964784130673e-05, 0.09950302422228456, 0.0001137934973534554,
        };

        float tmp1;
        float tmp2;
        float tmp3;


        {
            constexpr std::size_t ind = 0;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // apply revolute rotation to R_fixed and store in joint to base rotation R
            R11 = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12 = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13 = R13_fixed;
            R21 = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22 = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23 = R23_fixed;
            R31 = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32 = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33 = R33_fixed;

            // rotate offset to be in base frame
            input_data[ind * data_size + 2] = offset_x;
            input_data[ind * data_size + 3] = offset_y;
            input_data[ind * data_size + 4] = offset_z;




        }

        {
            constexpr std::size_t ind = 1;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 2;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 3;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 4;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

        {
            constexpr std::size_t ind = 5;

            // R_fixed is the rotation from joint_i+1 in frame joint_i (row major)
            const float &R11_fixed = R_all[ind * 9 + 0];
            const float &R21_fixed = R_all[ind * 9 + 1];
            const float &R31_fixed = R_all[ind * 9 + 2];
            const float &R12_fixed = R_all[ind * 9 + 3];
            const float &R22_fixed = R_all[ind * 9 + 4];
            const float &R32_fixed = R_all[ind * 9 + 5];
            const float &R13_fixed = R_all[ind * 9 + 6];
            const float &R23_fixed = R_all[ind * 9 + 7];
            const float &R33_fixed = R_all[ind * 9 + 8];

            // offset is the offset of joint_i+1 in frame joint_i
            const float &offset_x = offset_all[ind * 3 + 0];
            const float &offset_y = offset_all[ind * 3 + 1];
            const float &offset_z = offset_all[ind * 3 + 2];

            // R is the rotation from joint_i+1 in the base frame (row major)
            float &R11 = input_data[ind * data_size + 5];
            float &R12 = input_data[ind * data_size + 6];
            float &R13 = input_data[ind * data_size + 7];
            float &R21 = input_data[ind * data_size + 8];
            float &R22 = input_data[ind * data_size + 9];
            float &R23 = input_data[ind * data_size + 10];
            float &R31 = input_data[ind * data_size + 11];
            float &R32 = input_data[ind * data_size + 12];
            float &R33 = input_data[ind * data_size + 13];


            const float &sin_theta = input_data[ind * data_size + 0];
            const float &cos_theta = input_data[ind * data_size + 1];

            // R_old is the rotation from joint_i in base frame (column major)
            float &R11_old = input_data[(ind - 1) * data_size + 5];
            float &R12_old = input_data[(ind - 1) * data_size + 6];
            float &R13_old = input_data[(ind - 1) * data_size + 7];
            float &R21_old = input_data[(ind - 1) * data_size + 8];
            float &R22_old = input_data[(ind - 1) * data_size + 9];
            float &R23_old = input_data[(ind - 1) * data_size + 10];
            float &R31_old = input_data[(ind - 1) * data_size + 11];
            float &R32_old = input_data[(ind - 1) * data_size + 12];
            float &R33_old = input_data[(ind - 1) * data_size + 13];

            // R_tmp is the rotation from joint_i+1 in frame joint_i after rotation applied (column major)
            float &R11_tmp = input_data[ind * data_size + 5];
            float &R12_tmp = input_data[ind * data_size + 6];
            float &R13_tmp = input_data[ind * data_size + 7];
            float &R21_tmp = input_data[ind * data_size + 8];
            float &R22_tmp = input_data[ind * data_size + 9];
            float &R23_tmp = input_data[ind * data_size + 10];
            float &R31_tmp = input_data[ind * data_size + 11];
            float &R32_tmp = input_data[ind * data_size + 12];
            float &R33_tmp = input_data[ind * data_size + 13];

            // apply revolute rotation to R11_fixed and store in R11_tmp
            R11_tmp = R11_fixed * cos_theta + R12_fixed * sin_theta;
            R12_tmp = -R11_fixed * sin_theta + R12_fixed * cos_theta;
            R13_tmp = R13_fixed;
            R21_tmp = R21_fixed * cos_theta + R22_fixed * sin_theta;
            R22_tmp = -R21_fixed * sin_theta + R22_fixed * cos_theta;
            R23_tmp = R23_fixed;
            R31_tmp = R31_fixed * cos_theta + R32_fixed * sin_theta;
            R32_tmp = -R31_fixed * sin_theta + R32_fixed * cos_theta;
            R33_tmp = R33_fixed;

            // apply R11_old rotation R11_tmp, to make it in base frame
            tmp1 = R11_old * R11_tmp + R12_old * R21_tmp + R13_old * R31_tmp;
            tmp2 = R21_old * R11_tmp + R22_old * R21_tmp + R23_old * R31_tmp;
            tmp3 = R31_old * R11_tmp + R32_old * R21_tmp + R33_old * R31_tmp;
            R11 = tmp1;
            R21 = tmp2;
            R31 = tmp3;
            tmp1 = R11_old * R12_tmp + R12_old * R22_tmp + R13_old * R32_tmp;
            tmp2 = R21_old * R12_tmp + R22_old * R22_tmp + R23_old * R32_tmp;
            tmp3 = R31_old * R12_tmp + R32_old * R22_tmp + R33_old * R32_tmp;
            R12 = tmp1;
            R22 = tmp2;
            R32 = tmp3;
            tmp1 = R11_old * R13_tmp + R12_old * R23_tmp + R13_old * R33_tmp;
            tmp2 = R21_old * R13_tmp + R22_old * R23_tmp + R23_old * R33_tmp;
            tmp3 = R31_old * R13_tmp + R32_old * R23_tmp + R33_old * R33_tmp;
            R13 = tmp1;
            R23 = tmp2;
            R33 = tmp3;

            // p_old is the position of joint_i in base frame
            const float &px_old = input_data[(ind - 1) * data_size + 2];
            const float &py_old = input_data[(ind - 1) * data_size + 3];
            const float &pz_old = input_data[(ind - 1) * data_size + 4];

            // rotate offset to be in base frame and add p_old
            tmp1 = R11_old * offset_x + R12_old * offset_y + R13_old * offset_z + px_old;
            tmp2 = R21_old * offset_x + R22_old * offset_y + R23_old * offset_z + py_old;
            tmp3 = R31_old * offset_x + R32_old * offset_y + R33_old * offset_z + pz_old;
            input_data[ind * data_size + 2] = tmp1;
            input_data[ind * data_size + 3] = tmp2;
            input_data[ind * data_size + 4] = tmp3;



        }

    }

}


namespace fast_fk::internal {
    constexpr float axis_scale = 1;

    float InverseKinematics::operator()(const Eigen::VectorX<float> &q, Eigen::VectorX<float> &grad) {
        // construct input_data from q
        float tmp1;
        float tmp2;

        tmp1 = q[0];
        tmp2 = q[0];
        joint_data[0][0] = std::sin(tmp1);
        joint_data[0][1] = std::cos(tmp2);
        tmp1 = q[1];
        tmp2 = q[1];
        joint_data[1][0] = std::sin(tmp1);
        joint_data[1][1] = std::cos(tmp2);
        tmp1 = q[2];
        tmp2 = q[2];
        joint_data[2][0] = std::sin(tmp1);
        joint_data[2][1] = std::cos(tmp2);
        tmp1 = q[3];
        tmp2 = q[3];
        joint_data[3][0] = std::sin(tmp1);
        joint_data[3][1] = std::cos(tmp2);
        tmp1 = q[4];
        tmp2 = q[4];
        joint_data[4][0] = std::sin(tmp1);
        joint_data[4][1] = std::cos(tmp2);
        tmp1 = q[5];
        tmp2 = q[5];
        joint_data[5][0] = std::sin(tmp1);
        joint_data[5][1] = std::cos(tmp2);

        float *input_data = joint_data.data()->data();
        internal::forward_kinematics_internal(input_data, data_size * 6);

        Eigen::Vector3<float> target_x_axis_ = target_rot_.block<3, 1>(0, 0);
        Eigen::Vector3<float> target_y_axis_ = target_rot_.block<3, 1>(0, 1);
        Eigen::Vector3<float> target_z_axis_ = target_rot_.block<3, 1>(0, 2);
        Eigen::Vector3<float> target_x_ = target_pose_ + axis_scale * target_x_axis_;
        Eigen::Vector3<float> target_y_ = target_pose_ + axis_scale * target_y_axis_;
        Eigen::Vector3<float> target_z_ = target_pose_ + axis_scale * target_z_axis_;


        Eigen::Vector<float, 3> ee_pose;
        ee_pose << input_data[(6 - 1) *data_size + 2], input_data[(6 - 1)
                                                                  *data_size + 3], input_data[(6 - 1) *data_size + 4];

        // x point
        Eigen::Vector<float, 3> x_axis;
        x_axis << input_data[(6 - 1) *data_size + 5], input_data[(6 - 1)
                                                                 *data_size + 8], input_data[(6 - 1) *data_size + 11];
        Eigen::Vector<float, 3> current_x = ee_pose + x_axis;

        // y point
        Eigen::Vector<float, 3> y_axis;
        y_axis << input_data[(6 - 1) *data_size + 6], input_data[(6 - 1)
                                                                 *data_size + 9], input_data[(6 - 1) *data_size + 12];
        Eigen::Vector<float, 3> current_y = ee_pose + y_axis;

        // z point
        Eigen::Vector<float, 3> z_axis;
        z_axis << input_data[(6 - 1) *data_size + 7], input_data[(6 - 1)
                                                                 *data_size + 10], input_data[(6 - 1) *data_size + 13];
        Eigen::Vector<float, 3> current_z = ee_pose + z_axis;

        float fx = 0.0;
        Eigen::Vector3<float> delta_x = target_x_.array() - current_x.array();
        Eigen::Vector3<float> delta_y = target_y_.array() - current_y.array();
        Eigen::Vector3<float> delta_z = target_z_.array() - current_z.array();
        fx += (delta_x.array() * delta_x.array()).sum();
        fx += (delta_y.array() * delta_y.array()).sum();
        fx += (delta_z.array() * delta_z.array()).sum();

        // J =  (delta_x)^2
        //    + (delta_y)^2
        //    + (delta_z)^2

        // dJ/dq = dJ1/dq + dJ2/dq + dJ3/dq

        // dJ1/dq = dJ/ddelta_x * ddelta_x/dq
        // dJ2/dq = dJ/ddelta_y * ddelta_y/dq
        // dJ3/dq = dJ/ddelta_z * ddelta_z/dq

        // dJ/ddelta_x : 1x3 2*(x_target - x_current)
        // dJ/ddelta_y : 1x3 2*(y_target - y_current)
        // dJ/ddelta_z : 1x3 2*(z_target - z_current)

        // ddelta_x/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x x_axis
        // ddelta_y/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x y_axis
        // ddelta_z/dq: 3 x dofs J_{:3,:} + dofs J_{3:,:} x z_axis


        // gradient
        grad *= 0;
        Eigen::Vector3<float> dJ_ddelta_x = 2 * delta_x;
        Eigen::Vector3<float> dJ_ddelta_y = 2 * delta_y;
        Eigen::Vector3<float> dJ_ddelta_z = 2 * delta_z;

        Eigen::Vector<float, 3> joint_pose;
        Eigen::Vector<float, 3> delta;
        Eigen::Vector<float, 3> joint_axis;
        Eigen::Vector<float, 3> jac_linear;
        Eigen::Vector<float, 3> jac_angular;
        Eigen::Vector3<float> ddelta_x_dq_ind;
        Eigen::Vector3<float> ddelta_y_dq_ind;
        Eigen::Vector3<float> ddelta_z_dq_ind;


        {
            constexpr std::size_t ind = 0;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 1;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 2;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 3;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 4;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }

        {
            constexpr std::size_t ind = 5;

            joint_pose[0] = input_data[ind * data_size + 2];
            joint_pose[1] = input_data[ind * data_size + 3];
            joint_pose[2] = input_data[ind * data_size + 4];
            delta = ee_pose - joint_pose;
            joint_axis[0] = input_data[ind * data_size + 7];
            joint_axis[1] = input_data[ind * data_size + 10];
            joint_axis[2] = input_data[ind * data_size + 13];
            jac_linear = joint_axis.cross(delta);
            jac_angular = joint_axis;

            // x point
            ddelta_x_dq_ind = jac_linear + jac_angular.cross(axis_scale * x_axis);
            grad[ind] -= dJ_ddelta_x.dot(ddelta_x_dq_ind);

            // y point
            ddelta_y_dq_ind = jac_linear + jac_angular.cross(axis_scale * y_axis);
            grad[ind] -= dJ_ddelta_y.dot(ddelta_y_dq_ind);

            // z point
            ddelta_z_dq_ind = jac_linear + jac_angular.cross(axis_scale * z_axis);
            grad[ind] -= dJ_ddelta_z.dot(ddelta_z_dq_ind);


        }


        return fx;
    }


}

The Results

Summary: lessons learned

  • Leverage compile time optimization when needed in common robotics subroutines
  • Leverage code generation to squeeze out as much performance as possible from both the CPU optimizations and known hardware kinematic description
  • Try to optimize memory layout and algorithmic improvements

Thank you!

Questions?