move nearest_point util to track.utils,

refactor type hint with optional instead of Union[type, None]

f1tenth_gym/envs/track/cubic_spline.py

@@ -8,62 +8,9 @@
 import numpy as np
 import scipy.optimize as so
 from scipy import interpolate
-from numba import njit
-from typing import Union
+from typing import Union, Optional
 
-
-@njit(fastmath=False, cache=True)
-def nearest_point_on_trajectory(point: np.ndarray, trajectory: np.ndarray) -> tuple:
-    """
-    Return the nearest point along the given piecewise linear trajectory.
-
-    Same as nearest_point_on_line_segment, but vectorized. This method is quite fast, time constraints should
-    not be an issue so long as trajectories are not insanely long.
-
-        Order of magnitude: trajectory length: 1000 --> 0.0002 second computation (5000fps)
-
-    Parameters
-    ----------
-    point: np.ndarray
-        The 2d point to project onto the trajectory
-    trajectory: np.ndarray
-        The trajectory to project the point onto, shape (N, 2)
-        The points must be unique. If they are not unique, a divide by 0 error will destroy the world
-
-    Returns
-    -------
-    nearest_point: np.ndarray
-        The nearest point on the trajectory
-    distance: float
-        The distance from the point to the nearest point on the trajectory
-    t: float
-    min_dist_segment: int
-        The index of the nearest point on the trajectory
-    """
-    diffs = trajectory[1:, :] - trajectory[:-1, :]
-    l2s = diffs[:, 0] ** 2 + diffs[:, 1] ** 2
-    # this is equivalent to the elementwise dot product
-    # dots = np.sum((point - trajectory[:-1,:]) * diffs[:,:], axis=1)
-    dots = np.empty((trajectory.shape[0] - 1,))
-    for i in range(dots.shape[0]):
-        dots[i] = np.dot((point - trajectory[i, :]), diffs[i, :])
-    t = dots / l2s
-    t[t < 0.0] = 0.0
-    t[t > 1.0] = 1.0
-    # t = np.clip(dots / l2s, 0.0, 1.0)
-    projections = trajectory[:-1, :] + (t * diffs.T).T
-    # dists = np.linalg.norm(point - projections, axis=1)
-    dists = np.empty((projections.shape[0],))
-    for i in range(dists.shape[0]):
-        temp = point - projections[i]
-        dists[i] = np.sqrt(np.sum(temp * temp))
-    min_dist_segment = np.argmin(dists)
-    return (
-        projections[min_dist_segment],
-        dists[min_dist_segment],
-        t[min_dist_segment],
-        min_dist_segment,
-    )
+from f1tenth_gym.envs.track.utils import nearest_point_on_trajectory
 
 
 class CubicSpline2D:
@@ -114,7 +61,7 @@ def __calc_s(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
         s.extend(np.cumsum(self.ds))
         return np.array(s)
 
-    def calc_position(self, s: float) -> tuple[Union[float, None], Union[float, None]]:
+    def calc_position(self, s: float) -> np.ndarray:
         """
         Calc position at the given s.
 
@@ -133,7 +80,7 @@ def calc_position(self, s: float) -> tuple[Union[float, None], Union[float, None
         """
         return self.spline(s)
 
-    def calc_curvature(self, s: float) -> Union[float, None]:
+    def calc_curvature(self, s: float) -> Optional[float]:
         """
         Calc curvature at the given s.
 
@@ -153,7 +100,7 @@ def calc_curvature(self, s: float) -> Union[float, None]:
         k = (ddy * dx - ddx * dy) / ((dx**2 + dy**2) ** (3 / 2))
         return k
 
-    def calc_yaw(self, s: float) -> Union[float, None]:
+    def calc_yaw(self, s: float) -> Optional[float]:
         """
         Calc yaw angle at the given s.
 

f1tenth_gym/envs/track/utils.py

@@ -2,7 +2,9 @@
 import tarfile
 import tempfile
 
+import numpy as np
 import requests
+from numba import njit
 
 
 def find_track_dir(track_name: str) -> pathlib.Path:
@@ -49,3 +51,57 @@ def find_track_dir(track_name: str) -> pathlib.Path:
             return subdir
 
     raise FileNotFoundError(f"no mapdir matching {track_name} in {[map_dir]}")
+
+
+@njit(fastmath=False, cache=True)
+def nearest_point_on_trajectory(point: np.ndarray, trajectory: np.ndarray) -> tuple:
+    """
+    Return the nearest point along the given piecewise linear trajectory.
+
+    Same as nearest_point_on_line_segment, but vectorized. This method is quite fast, time constraints should
+    not be an issue so long as trajectories are not insanely long.
+
+        Order of magnitude: trajectory length: 1000 --> 0.0002 second computation (5000fps)
+
+    Parameters
+    ----------
+    point: np.ndarray
+        The 2d point to project onto the trajectory
+    trajectory: np.ndarray
+        The trajectory to project the point onto, shape (N, 2)
+        The points must be unique. If they are not unique, a divide by 0 error will destroy the world
+
+    Returns
+    -------
+    nearest_point: np.ndarray
+        The nearest point on the trajectory
+    distance: float
+        The distance from the point to the nearest point on the trajectory
+    t: float
+    min_dist_segment: int
+        The index of the nearest point on the trajectory
+    """
+    diffs = trajectory[1:, :] - trajectory[:-1, :]
+    l2s = diffs[:, 0] ** 2 + diffs[:, 1] ** 2
+    # this is equivalent to the elementwise dot product
+    # dots = np.sum((point - trajectory[:-1,:]) * diffs[:,:], axis=1)
+    dots = np.empty((trajectory.shape[0] - 1,))
+    for i in range(dots.shape[0]):
+        dots[i] = np.dot((point - trajectory[i, :]), diffs[i, :])
+    t = dots / l2s
+    t[t < 0.0] = 0.0
+    t[t > 1.0] = 1.0
+    # t = np.clip(dots / l2s, 0.0, 1.0)
+    projections = trajectory[:-1, :] + (t * diffs.T).T
+    # dists = np.linalg.norm(point - projections, axis=1)
+    dists = np.empty((projections.shape[0],))
+    for i in range(dists.shape[0]):
+        temp = point - projections[i]
+        dists[i] = np.sqrt(np.sum(temp * temp))
+    min_dist_segment = np.argmin(dists)
+    return (
+        projections[min_dist_segment],
+        dists[min_dist_segment],
+        t[min_dist_segment],
+        min_dist_segment,
+    )