feat: python interface

- python interface
- extended readme with example
- added example to tests
- made q0 optional in cpp code
- bind home pose

Author: juelg

README.md

@@ -4,6 +4,26 @@
 
 `frankik` is a standalone Python library that implements the analytical geometric IK solver proposed by **He & Liu (2021)**. It is designed for researchers and developers who need high-performance kinematics without the overhead of ROS, MoveIt, or hardware drivers.
 
+## Example
+```python
+import numpy as np
+from frankik import FrankaKinematics, RobotType
+kinematics = FrankaKinematics(robot_type=RobotType.FR3) # Robot.Type.PANDA also supported
+q_home = np.array([
+    0.0, 
+    -np.pi / 4, 
+    0.0, 
+    -3 * np.pi / 4, 
+    0.0, 
+    np.pi / 2, 
+    np.pi / 4
+])
+pose_home = kinematics.forward(q_home)
+q = kinematics.inverse(pose_home)
+assert np.allclose(q, q_home)
+print(q)
+```
+
 ## Installation from source
 ```shell
 sudo apt install libeigen3-dev
@@ -20,7 +40,7 @@ Coming soon...
 ## Feature List
 - [x] basic bindings from he
 - [x] support for fr3
-- [ ] nice python interface
-- [ ] tests
+- [x] nice python interface
+- [x] tests
 - [ ] performance benchmarks
 

src/frankik/__init__.py

@@ -1,12 +1,121 @@
+import typing
+
+import numpy as np
+
 from frankik._core import (
     __version__,
     fk,
     ik,
     ik_full,
+    kQDefault,
     q_max_fr3,
     q_max_panda,
     q_min_fr3,
     q_min_panda,
 )
 
-__all__ = ["__version__", "ik", "ik_full", "fk", "q_max_fr3", "q_max_panda", "q_min_fr3", "q_min_panda"]
+
+class RobotType:
+    PANDA = "panda"
+    FR3 = "fr3"
+
+
+class FrankaKinematics:
+
+    def __init__(self, robot_type: str = RobotType.FR3):
+        """Initialize Franka Kinematics for the specified robot type.
+        Args:
+            robot_type (str): Type of the robot, either 'panda' or 'fr3'.
+        Raises:
+            ValueError: If an unsupported robot type is provided.
+        """
+        if robot_type not in (RobotType.PANDA, RobotType.FR3):
+            raise ValueError(f"Unsupported robot type: {robot_type}. Choose 'panda' or 'fr3'.")
+        self.robot_type = robot_type
+        self.q_min = q_min_fr3 if robot_type == RobotType.FR3 else q_min_panda
+        self.q_max = q_max_fr3 if robot_type == RobotType.FR3 else q_max_panda
+        self.q_home = kQDefault
+
+    @staticmethod
+    def pose_inverse(
+        T: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]]
+    ) -> np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]]:
+        """Compute the inverse of a homogeneous transformation matrix.
+        Args:
+            T (np.ndarray): A 4x4 homogeneous transformation matrix.
+        Returns:
+            np.ndarray: The inverse of the input transformation matrix.
+        """
+        R = T[:3, :3]
+        t = T[:3, 3]
+        T_inv = np.eye(4)
+        T_inv[:3, :3] = R.T
+        T_inv[:3, 3] = -R.T @ t
+        return T_inv
+
+    def forward(
+        self,
+        q0: np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]],
+        tcp_offset: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]] | None = None,
+    ) -> np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]]:
+        """Compute the forward kinematics for the given joint configuration.
+        Args:
+            q0 (np.ndarray): A 7-element array representing joint angles.
+            tcp_offset (np.ndarray, optional): A 4x4 homogeneous transformation matrix representing
+                the tool center point offset. Defaults to None.
+        Returns:
+            np.ndarray: A 4x4 homogeneous transformation matrix representing the end-effector pose.
+        """
+        pose = fk(q0)
+        return pose @ self.pose_inverse(tcp_offset) if tcp_offset is not None else pose
+
+    def inverse(
+        self,
+        pose: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]],
+        q0: np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]] | None = None,
+        tcp_offset: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]] | None = None,
+        allow_elbow_flips: bool = False,
+    ) -> np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]] | None:
+        """Compute the inverse kinematics for the given end-effector pose.
+        Args:
+            pose (np.ndarray): A 4x4 homogeneous transformation matrix representing the desired end-effector pose.
+            q0 (np.ndarray, optional): A 7-element array representing the current joint angles. Defaults to None.
+            tcp_offset (np.ndarray, optional): A 4x4 homogeneous transformation matrix representing
+                the tool center point offset. Defaults to None.
+            allow_elbow_flips (bool): Whether to consider multiple IK solutions (elbow flips). Defaults to False.
+        Returns:
+            np.ndarray | None: A 7-element array representing the joint angles if a solution is found; otherwise, None.
+        """
+        new_pose = pose @ self.pose_inverse(tcp_offset) if tcp_offset is not None else pose
+        q = ik(new_pose, q0, is_fr3=(self.robot_type == RobotType.FR3))
+
+        if not np.isnan(q).any():
+            return q
+        if not allow_elbow_flips:
+            return None
+        qs = ik_full(new_pose, q0, is_fr3=(self.robot_type == RobotType.FR3))
+        qs_list = list(qs)
+        # drop nan solutions
+        qs_list = [q for q in qs_list if not np.isnan(q).any()]
+        if len(qs_list) == 0:
+            return None
+        if q0 is None:
+            q0 = self.q_home
+
+        # find the closest solution
+        best_q = min(qs_list, key=lambda q_sol: np.linalg.norm(q_sol - q0))
+        return best_q
+
+
+__all__ = [
+    "__version__",
+    "ik",
+    "ik_full",
+    "fk",
+    "q_max_fr3",
+    "q_max_panda",
+    "q_min_fr3",
+    "q_min_panda",
+    "FrankaKinematics",
+    "RobotType",
+]

src/frankik/_core.pyi

@@ -8,7 +8,7 @@ import typing
 
 import numpy
 
-__all__: list[str] = ["fk", "ik", "ik_full", "q_max_fr3", "q_max_panda", "q_min_fr3", "q_min_panda"]
+__all__: list[str] = ["fk", "ik", "ik_full", "kQDefault", "q_max_fr3", "q_max_panda", "q_min_fr3", "q_min_panda"]
 
 def fk(
     q: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]]
@@ -25,7 +25,7 @@ def fk(
 
 def ik(
     O_T_EE: numpy.ndarray[tuple[typing.Literal[4], typing.Literal[4]], numpy.dtype[numpy.float64]],
-    q_actual_array: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] = ...,
+    q_actual: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] | None = None,
     q7: float = 0.7853981633974483,
     is_fr3: bool = False,
 ) -> numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]]:
@@ -34,7 +34,7 @@ def ik(
 
     Args:
         O_T_EE (Eigen::Matrix<double, 4, 4>): Desired end-effector pose.
-        q_actual_array (Vector7d, optional): Current joint angles. Defaults to kQDefault.
+        q_actual (Vector7d, optional): Current joint angles. Defaults to kQDefault.
         q7 (double, optional): Joint 7 angle. Defaults to M_PI_4.
         is_fr3 (bool, optional): Whether to use FR3 joint limits. Defaults to false.
 
@@ -44,7 +44,7 @@ def ik(
 
 def ik_full(
     O_T_EE: numpy.ndarray[tuple[typing.Literal[4], typing.Literal[4]], numpy.dtype[numpy.float64]],
-    q_actual_array: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] = ...,
+    q_actual: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] | None = None,
     q7: float = 0.7853981633974483,
     is_fr3: bool = False,
 ) -> numpy.ndarray[tuple[typing.Literal[4], typing.Literal[7]], numpy.dtype[numpy.float64]]:
@@ -53,7 +53,7 @@ def ik_full(
 
     Args:
         O_T_EE (Eigen::Matrix<double, 4, 4>): Desired end-effector pose.
-        q_actual_array (Vector7d, optional): Current joint angles. Defaults to kQDefault.
+        q_actual (Vector7d, optional): Current joint angles. Defaults to kQDefault.
         q7 (double, optional): Joint 7 angle. Defaults to M_PI_4.
         is_fr3 (bool, optional): Whether to use FR3 joint limits. Defaults to false.
 
@@ -62,6 +62,7 @@ def ik_full(
     """
 
 __version__: str = "0.1"
+kQDefault: numpy.ndarray  # value = array([ 0.        , -0.78539816,  0.        , -2.35619449,  0.        ,...
 q_max_fr3: list = [2.3093, 1.5133, 2.4937, -0.4461, 2.48, 4.2094, 2.6895]
 q_max_panda: list = [2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973]
 q_min_fr3: list = [-2.3093, -1.5133, -2.4937, -2.7478, -2.48, 0.8521, -2.6895]

src/pybind/bindings.cpp

@@ -27,14 +27,15 @@ PYBIND11_MODULE(_core, m) {
   m.attr("q_max_panda") = frankik::q_max_panda;
   m.attr("q_min_fr3") = frankik::q_min_fr3;
   m.attr("q_max_fr3") = frankik::q_max_fr3;
+  m.attr("kQDefault") = frankik::kQDefault;
 
   m.def("ik_full", &frankik::ik_full, py::arg("O_T_EE"),
-        py::arg("q_actual_array") = frankik::kQDefault, py::arg("q7") = M_PI_4,
+        py::arg("q_actual") = std::nullopt, py::arg("q7") = M_PI_4,
         py::arg("is_fr3") = false,
         "Compute full inverse kinematics for Franka Emika Panda robot.\n\n"
         "Args:\n"
         "    O_T_EE (Eigen::Matrix<double, 4, 4>): Desired end-effector pose.\n"
-        "    q_actual_array (Vector7d, optional): Current joint angles. "
+        "    q_actual (Vector7d, optional): Current joint angles. "
         "Defaults to kQDefault.\n"
         "    q7 (double, optional): Joint 7 angle. Defaults to M_PI_4.\n"
         "    is_fr3 (bool, optional): Whether to use FR3 joint limits. "
@@ -43,13 +44,13 @@ PYBIND11_MODULE(_core, m) {
         "    Eigen::Matrix<double, 4, 7>: All possible IK solutions (up to 4). "
         "NaN if no solution.");
   m.def("ik", &frankik::ik, py::arg("O_T_EE"),
-        py::arg("q_actual_array") = frankik::kQDefault, py::arg("q7") = M_PI_4,
+        py::arg("q_actual") = std::nullopt, py::arg("q7") = M_PI_4,
         py::arg("is_fr3") = false,
         "Compute one inverse kinematics solution for Franka Emika Panda "
         "robot.\n\n"
         "Args:\n"
         "    O_T_EE (Eigen::Matrix<double, 4, 4>): Desired end-effector pose.\n"
-        "    q_actual_array (Vector7d, optional): Current joint angles. "
+        "    q_actual (Vector7d, optional): Current joint angles. "
         "Defaults to kQDefault.\n"
         "    q7 (double, optional): Joint 7 angle. Defaults to M_PI_4.\n"
         "    is_fr3 (bool, optional): Whether to use FR3 joint limits. "

src/pybind/frankik.h

@@ -41,10 +41,11 @@ const Vector7d kQDefault =
     (Vector7d() << 0.0, -M_PI_4, 0.0, -3 * M_PI_4, 0.0, M_PI_2, M_PI_4)
         .finished();
 
-Eigen::Matrix<double, 4, 7> ik_full(Eigen::Matrix<double, 4, 4> O_T_EE,
-                                    Vector7d q_actual_array = kQDefault,
-                                    double q7 = M_PI_4, bool is_fr3 = false) {
-
+Eigen::Matrix<double, 4, 7>
+ik_full(Eigen::Matrix<double, 4, 4> O_T_EE,
+        std::optional<Vector7d> q_actual = std::nullopt, double q7 = M_PI_4,
+        bool is_fr3 = false) {
+  Vector7d q_actual_array = q_actual.value_or(kQDefault);
   const Eigen::Matrix<double, 4, 7> q_all_NAN =
       Eigen::Matrix<double, 4, 7>::Constant(
           std::numeric_limits<double>::quiet_NaN());
@@ -221,8 +222,9 @@ Eigen::Matrix<double, 4, 7> ik_full(Eigen::Matrix<double, 4, 4> O_T_EE,
 }
 
 Vector7d ik(Eigen::Matrix<double, 4, 4> O_T_EE,
-            Vector7d q_actual_array = frankik::kQDefault, double q7 = M_PI_4,
+            std::optional<Vector7d> q_actual = std::nullopt, double q7 = M_PI_4,
             bool is_fr3 = false) {
+  Vector7d q_actual_array = q_actual.value_or(kQDefault);
   const Vector7d q_NAN =
       Vector7d::Constant(std::numeric_limits<double>::quiet_NaN());
   Vector7d q;

src/tests/test_example.py

@@ -0,0 +1,13 @@
+import numpy as np
+import pytest
+
+from frankik import FrankaKinematics, RobotType
+
+
+@pytest.mark.parametrize("robot_type", [RobotType.PANDA, RobotType.FR3])
+def test_example(robot_type):
+    kinematics = FrankaKinematics(robot_type=robot_type)
+    q_home = np.array([0.0, -np.pi / 4, 0.0, -3 * np.pi / 4, 0.0, np.pi / 2, np.pi / 4])
+    pose_home = kinematics.forward(q_home)
+    q = kinematics.inverse(pose_home)
+    assert np.allclose(q, q_home)