trajectory generation update

.gitignore

@@ -3,6 +3,8 @@ TODOs.md
 examples/debug.py
 examples/test.py
 
+build_vscode/
+
 # NumPy saves and videos
 results/
 tmp/

pyproject.toml

@@ -12,6 +12,7 @@ packages = [
 
 [tool.poetry.dependencies]
 python = "^3.10"
+trajectory_cpp = { path = "./trajectories/trajectory_generation" }
 numpy = "^1.24"
 scipy = "^1.10"
 transforms3d = "^0.4.1"

tests/test_trajectories.py

@@ -1,9 +1,10 @@
+
 def test_imports():
     from trajectories import TrajectoryFactory
 
 def test_square_trajectory():
     from trajectories import TrajectoryFactory 
-    from trajectories.trajectories import SquareLinearTrajectory
+    from trajectories.square_linear_trajectory import SquareLinearTrajectory
     import numpy as np
     trajectory: SquareLinearTrajectory = TrajectoryFactory.get_linear_square_trajectory()
     p0 = trajectory.get_waypoint(0) 
@@ -20,3 +21,31 @@ def test_square_trajectory():
     # assert interpolation
     assert(trajectory.get_waypoint(0.1).coordinate == np.asarray([0.4, 0., 1.], dtype=np.float32)).all()
 
+def test_pol_trajectory():
+    from trajectories import TrajectoryFactory, Waypoint
+    import numpy as np
+    t_waypoints = [
+        Waypoint(
+            np.array([0,0,0], dtype=np.float64),
+            0
+        ),
+        Waypoint(
+            np.array([1,1,1], dtype=np.float64),
+            1
+        ),
+        Waypoint(
+            np.array([2,2,2], dtype=np.float64),
+            2
+        ),
+        Waypoint(
+            np.array([3,3,3], dtype=np.float64),
+            3
+        )
+    ]
+    traj = TrajectoryFactory.get_pol_discretized_trajectory(
+        t_waypoints,
+        10
+    )
+    threshhold = 0.001
+    assert(np.sum(np.square(traj[0].coordinate - t_waypoints[0].coordinate)) < threshhold)
+    assert(np.sum(np.square(traj[9].coordinate - t_waypoints[3].coordinate)) < threshhold)

trajectories/__init__.py

@@ -1 +1,2 @@
-from .trajectory_factory import TrajectoryFactory
+from .trajectory_factory import TrajectoryFactory
+from .waypoint import Waypoint

trajectories/discretized_trajectory.py

@@ -0,0 +1,20 @@
+from .waypoint import Waypoint
+from typing import List
+
+class DiscretizedTrajectory:
+    waypoints: List[Waypoint]
+
+    def __init__(self) -> None:
+        pass
+
+    def __len__(self) -> int: 
+        """
+        Must be implemented by child class.
+        """
+        raise NotImplementedError("Must be implemented by child class.")
+
+    def __getitem__(self, idx: int) -> Waypoint:
+        """
+        Must be implemented by child class.
+        """
+        raise NotImplementedError("Must be implemented by child class.")

trajectories/polynomial_discretized_trajectory.py

@@ -0,0 +1,57 @@
+from .discretized_trajectory import DiscretizedTrajectory
+from .waypoint import Waypoint
+from .traj_gen_cpp_wrapper import calculate_trajectory
+
+from typing import List
+import numpy as np
+
+def calc_target_durations(waypoints: List[Waypoint]) -> np.ndarray:
+    wps = iter(waypoints)
+    cur_time = next(wps).timestamp
+    assert(cur_time == 0)
+    durations = []
+    for wp in wps:
+        durations.append(wp.timestamp - cur_time)
+        cur_time = wp.timestamp
+
+    return np.array(durations)
+
+def convert_to_np_waypoints(waypoints: List[Waypoint]) -> np.ndarray:
+    t_waypoints_np = np.zeros((3, len(waypoints)), dtype=np.float64)
+
+    for idx, wp in enumerate(waypoints):
+        t_waypoints_np[:,idx] = wp.coordinate
+
+    return t_waypoints_np
+
+class PolynomialDiscretizedTrajectory(DiscretizedTrajectory):
+    waypoints: np.ndarray
+    timestamps: np.ndarray
+
+    def __init__(self, t_waypoints: List[Waypoint], n_points_discretization_level: int) -> None:
+
+        # target waypoints with durations
+        t_waypoints_np = convert_to_np_waypoints(t_waypoints)
+        t_durations_np = calc_target_durations(t_waypoints)
+
+        # call cpp backbone
+        self.waypoints, self.timestamps = calculate_trajectory(
+            t_waypoints_np,
+            t_durations_np,
+            n_points_discretization_level
+        )
+
+        # assert n_points
+        assert(self.waypoints.shape[1] == n_points_discretization_level) 
+        assert(self.waypoints.shape[1] == len(self.timestamps))
+
+    def __len__(self) -> int:
+        return len(self.timestamps)        
+
+    def __getitem__(self, idx: int) -> Waypoint:
+        wp = Waypoint(
+            self.waypoints[:, idx],
+            self.timestamps[idx]
+        )
+        return wp
+

trajectories/trajectories.py → trajectories/square_linear_trajectory.py

@@ -1,17 +1,7 @@
+from .trajectory import Trajectory
+import numpy as np
 from .waypoint import Waypoint
-import numpy as np 
-import math 
-
-class Trajectory:
-
-    def __init__(self) -> None:
-        pass
-
-    def get_waypoint(self, time: float) -> Waypoint:
-        """
-        Needs to be implemented by subclass. time \in [0,1]
-        """
-        raise NotImplementedError()
+import math
 
 class SquareLinearTrajectory(Trajectory):
     """
@@ -53,6 +43,3 @@ def get_waypoint(self, time: float):
         )
 
         return target_wp
-
-
-

trajectories/traj_gen_cpp_wrapper.py

@@ -0,0 +1,19 @@
+from trajectory_cpp import calc_trajectory
+import numpy as np
+from typing import Tuple
+
+def calculate_trajectory(
+    t_waypoints: np.ndarray,
+    t_durations: np.ndarray,
+    n_discretization_level: int
+) -> Tuple[np.ndarray, np.ndarray]:
+    r_waypoints = np.zeros((3, n_discretization_level), dtype=np.float64)
+    r_timestamps = np.zeros(n_discretization_level, dtype=np.float64)
+    calc_trajectory(
+        t_waypoints,
+        t_durations,
+        r_waypoints,
+        r_timestamps
+    )
+    return r_waypoints, r_timestamps
+

trajectories/trajectory.py

@@ -0,0 +1,15 @@
+from .waypoint import Waypoint
+
+class Trajectory:
+
+    def __init__(self) -> None:
+        pass
+
+    def get_waypoint(self, time: float) -> Waypoint:
+        """
+        Needs to be implemented by subclass. time \in [0,1]
+        """
+        raise NotImplementedError()
+
+
+

trajectories/trajectory_factory.py

@@ -1,14 +1,24 @@
-from .trajectories import Trajectory, SquareLinearTrajectory
+from .trajectory import Trajectory
+from .square_linear_trajectory import SquareLinearTrajectory
+from .polynomial_discretized_trajectory import PolynomialDiscretizedTrajectory
+from .waypoint import Waypoint
+from typing import List
 
 class TrajectoryFactory:
     """
     Wrapper class for instantiating target trajectories.
     """
 
     @classmethod
-    def get_linear_square_trajectory(cls, square_scale: float=1, time_scale: float=1) -> Trajectory:
+    def get_linear_square_trajectory(cls, square_scale: float=1, time_scale: float=1) -> SquareLinearTrajectory:
         return SquareLinearTrajectory(
             square_scale, time_scale
         )
 
+    @classmethod
+    def get_pol_discretized_trajectory(cls, t_waypoints: List[Waypoint], n_points_discretization_level: int) -> PolynomialDiscretizedTrajectory:
+        return PolynomialDiscretizedTrajectory(
+            t_waypoints,
+            n_points_discretization_level
+        )