add at_time overload which accepts lists

pantor · Dec 25, 2024 · 93e0858 · 93e0858
1 parent beb713a
commit 93e0858
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/examples/02_position_offline.py b/examples/02_position_offline.py
@@ -38,5 +38,11 @@
 
     print(f'Position at time {new_time:0.4f} [s]: {new_position}')
 
+    # Or we can pass a list of time values to get the corresponding kinematic states
+    new_times = [0.5, 1.0, 1.5]
+    new_positions, new_velocities, new_accelerations = trajectory.at_time(new_times)
+
+    print(f'Positions at times {new_times[0]:0.4f} {new_times[1]:0.4f} {new_times[2]:0.4f} [s]: {new_positions}')
+
     # Get some info about the position extrema of the trajectory
     print(f'Position extremas are {trajectory.position_extrema}')
diff --git a/src/ruckig/python.cpp b/src/ruckig/python.cpp
@@ -77,6 +77,27 @@ limited by velocity, acceleration, and jerk constraints.";
             }
             return py::make_tuple(new_position, new_velocity, new_acceleration);
         }, "time"_a, "return_section"_a=false)
+        .def("at_time", [](const Trajectory<DynamicDOFs>& traj, const py::list& times, bool return_section=false) {
+            std::vector<std::vector<double>> new_position(times.size(), std::vector<double>(traj.degrees_of_freedom));
+            std::vector<std::vector<double>> new_velocity(times.size(), std::vector<double>(traj.degrees_of_freedom));
+            std::vector<std::vector<double>> new_acceleration(times.size(), std::vector<double>(traj.degrees_of_freedom));
+            std::vector<std::vector<double>> new_jerk(times.size(), std::vector<double>(traj.degrees_of_freedom));
+            py::list result_position, result_velocity, result_acceleration, result_section;
+            for (std::size_t i = 0; i < times.size(); i++) {
+                const auto time = times[i].cast<double>();
+                size_t new_section;
+                traj.at_time(time, new_position[i], new_velocity[i], new_acceleration[i], new_jerk[i], new_section);
+                result_position.append(new_position[i]);
+                result_velocity.append(new_velocity[i]);
+                result_acceleration.append(new_acceleration[i]);
+                result_section.append(new_section);
+            }
+
+            if (return_section) {
+                return py::make_tuple(result_position, result_velocity, result_acceleration, result_section);
+            }
+            return py::make_tuple(result_position, result_velocity, result_acceleration);
+        }, "times"_a, "return_section"_a=false)
         .def("get_first_time_at_position", [](const Trajectory<DynamicDOFs>& traj, size_t dof, double position) -> py::object {
             double time;
             if (traj.get_first_time_at_position(dof, position, time)) {