add at_time overload which accepts lists

pantor · Dec 26, 2024 · 270d08e · 270d08e
1 parent beb713a
commit 270d08e
Show file tree

Hide file tree

Showing 2 changed files with 24 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,24 @@ 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<double> new_position(traj.degrees_of_freedom), new_velocity(traj.degrees_of_freedom), new_acceleration(traj.degrees_of_freedom), new_jerk(traj.degrees_of_freedom);
+            py::list result_position(times.size()), result_velocity(times.size()), result_acceleration(times.size()), result_section(times.size());
+            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, new_velocity, new_acceleration, new_jerk, new_section);
+                result_position[i] = new_position;
+                result_velocity[i] = new_velocity;
+                result_acceleration[i] = new_acceleration;
+                result_section[i] = 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)) {