feat: export compute_total_inertia_matrix to python binding

HiroIshida · Jun 2, 2024 · 2cf4e1d · 2cf4e1d
1 parent 298a0d9
commit 2cf4e1d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/python/tinyfk/__init__.py b/python/tinyfk/__init__.py
@@ -188,6 +188,9 @@ def solve_com_fk(
         P, J = self._robot.solve_com_forward_kinematics(qs, joint_ids, with_base, with_jacobian)
         return P, J
 
+    def compute_total_inertia_matrix(self, q: np.ndarray, joint_ids: List[int]) -> np.ndarray:
+        return self._robot.compute_total_inertia_matrix(q, joint_ids)
+
     def get_joint_names(self) -> List[str]:
         return self._robot.get_joint_names()
 

diff --git a/src/wrapper.cpp b/src/wrapper.cpp
@@ -152,6 +152,14 @@ class KinematicsModelPyWrapper : public KinematicModel {
     return std::array<Eigen::MatrixXd, 2>{coms.transpose(), J};
   }
 
+  Eigen::Matrix3d
+  compute_total_inertia_matrix(const std::vector<double> &q,
+                               const std::vector<size_t> &joint_ids) {
+    // should we take vector of q as input...?
+    this->_set_joint_angles(joint_ids, q);
+    return this->get_total_inertia_matrix();
+  }
+
   std::pair<Eigen::VectorXd, Eigen::MatrixXd>
   compute_inter_link_squared_dists(const std::vector<std::vector<double>> &qs,
                                    const std::vector<size_t> &link_ids1,
@@ -226,6 +234,8 @@ PYBIND11_MODULE(_tinyfk, m) {
       .def("set_joint_angles", &KinematicsModelPyWrapper::set_joint_angles)
       .def("solve_com_forward_kinematics",
            &KinematicsModelPyWrapper::solve_com_forward_kinematics)
+      .def("compute_total_inertia_matrix",
+           &KinematicsModelPyWrapper::compute_total_inertia_matrix)
       .def("get_joint_names", &KinematicsModelPyWrapper::get_joint_names)
       .def("get_joint_ids", &KinematicsModelPyWrapper::get_joint_ids)
       .def("get_joint_limits", &KinematicsModelPyWrapper::get_joint_limits)