From 7528d560e00ecb24d954c73abd91814c249c3a6d Mon Sep 17 00:00:00 2001 From: "Berducci, Luigi" Date: Fri, 14 Jun 2024 16:42:08 +0200 Subject: [PATCH] linting, docstrings and minor changes --- f1tenth_gym/envs/track/cubic_spline.py | 106 +++++++++++++++++-------- tests/test_cubic_spline.py | 43 +++++++--- 2 files changed, 103 insertions(+), 46 deletions(-) diff --git a/f1tenth_gym/envs/track/cubic_spline.py b/f1tenth_gym/envs/track/cubic_spline.py index fd61e644..fa87d9a6 100644 --- a/f1tenth_gym/envs/track/cubic_spline.py +++ b/f1tenth_gym/envs/track/cubic_spline.py @@ -2,6 +2,7 @@ Cubic Spline interpolation using scipy.interpolate Provides utilities for position, curvature, yaw, and arclength calculation """ + import math import numpy as np @@ -10,8 +11,9 @@ from numba import njit from typing import Union + @njit(fastmath=False, cache=True) -def nearest_point_on_trajectory(point, trajectory): +def nearest_point_on_trajectory(point: np.ndarray, trajectory: np.ndarray) -> tuple: """ Return the nearest point along the given piecewise linear trajectory. @@ -20,9 +22,23 @@ def nearest_point_on_trajectory(point, trajectory): Order of magnitude: trajectory length: 1000 --> 0.0002 second computation (5000fps) - point: size 2 numpy array - trajectory: Nx2 matrix of (x,y) trajectory waypoints - - these must be unique. If they are not unique, a divide by 0 error will destroy the world + 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 @@ -49,6 +65,7 @@ def nearest_point_on_trajectory(point, trajectory): min_dist_segment, ) + class CubicSpline2D: """ Cubic CubicSpline2D class. @@ -66,11 +83,13 @@ class CubicSpline2D: def __init__(self, x, y): self.points = np.c_[x, y] if not np.all(self.points[-1] == self.points[0]): - self.points = np.vstack((self.points, self.points[0])) # Ensure the path is closed + self.points = np.vstack( + (self.points, self.points[0]) + ) # Ensure the path is closed self.s = self.__calc_s(self.points[:, 0], self.points[:, 1]) # Use scipy CubicSpline to interpolate the points with periodic boundary conditions # This is necessary to ensure the path is continuous - self.spline = interpolate.CubicSpline(self.s, self.points, bc_type='periodic') + self.spline = interpolate.CubicSpline(self.s, self.points, bc_type="periodic") def __calc_s(self, x: np.ndarray, y: np.ndarray) -> np.ndarray: """ @@ -114,7 +133,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) -> Union[float, None]: """ Calc curvature at the given s. @@ -152,18 +171,24 @@ def calc_yaw(self, s: float) -> Union[float, None]: yaw = math.atan2(dy, dx) # Convert yaw to [0, 2pi] yaw = yaw % (2 * math.pi) - + return yaw - def calc_arclength(self, x, y, s_guess=0.0): + def calc_arclength( + self, x: float, y: float, s_guess: float = 0.0 + ) -> tuple[float, float]: """ - calc arclength + Calculate arclength for a given point (x, y) on the trajectory. + Parameters ---------- x : float x position. y : float y position. + s_guess : float + initial guess for s. + Returns ------- s : float @@ -174,24 +199,26 @@ def calc_arclength(self, x, y, s_guess=0.0): def distance_to_spline(s): x_eval, y_eval = self.spline(s)[0] - return np.sqrt((x - x_eval)**2 + (y - y_eval)**2) - + return np.sqrt((x - x_eval) ** 2 + (y - y_eval) ** 2) + output = so.fmin(distance_to_spline, s_guess, full_output=True, disp=False) - closest_s = output[0][0] + closest_s = float(output[0][0]) absolute_distance = output[1] return closest_s, absolute_distance - - def calc_arclength_inaccurate(self, x, y, s_guess=0.0): + + def calc_arclength_inaccurate(self, x: float, y: float) -> tuple[float, float]: """ - calc arclength, use nearest_point_on_trajectory - Less accuarate and less smooth than calc_arclength but - much faster - suitable for lap counting + Fast calculation of arclength for a given point (x, y) on the trajectory. + Less accuarate and less smooth than calc_arclength but much faster. + Suitable for lap counting. + Parameters ---------- x : float x position. y : float y position. + Returns ------- s : float @@ -199,42 +226,51 @@ def calc_arclength_inaccurate(self, x, y, s_guess=0.0): ey : float lateral deviation for given x, y. """ - _, ey, t, min_dist_segment = nearest_point_on_trajectory(np.array([x, y]), self.points) + _, ey, t, min_dist_segment = nearest_point_on_trajectory( + np.array([x, y]), self.points + ) # s = s at closest_point + t - s = self.s[min_dist_segment] + t * (self.s[min_dist_segment + 1] - self.s[min_dist_segment]) + s = float( + self.s[min_dist_segment] + + t * (self.s[min_dist_segment + 1] - self.s[min_dist_segment]) + ) - return s, 0 + return s, 0.0 + + def _calc_tangent(self, s: float) -> np.ndarray: + """ + Calculates the tangent to the curve at a given point. - def _calc_tangent(self, s): - ''' - calculates the tangent to the curve at a given point Parameters ---------- s : float - distance from the start point. if `s` is outside the data point's - range, return None. + distance from the start point. + If `s` is outside the data point's range, return None. + Returns ------- tangent : float tangent vector for given s. - ''' + """ dx, dy = self.spline(s, 1) tangent = np.array([dx, dy]) return tangent - - def _calc_normal(self, s): - ''' - calculates the normal to the curve at a given point + + def _calc_normal(self, s: float) -> np.ndarray: + """ + Calculate the normal to the curve at a given point. + Parameters ---------- s : float - distance from the start point. if `s` is outside the data point's - range, return None. + distance from the start point. + If `s` is outside the data point's range, return None. + Returns ------- normal : float normal vector for given s. - ''' + """ dx, dy = self.spline(s, 1) normal = np.array([-dy, dx]) - return normal \ No newline at end of file + return normal diff --git a/tests/test_cubic_spline.py b/tests/test_cubic_spline.py index 4867bfbf..a9b3ca94 100644 --- a/tests/test_cubic_spline.py +++ b/tests/test_cubic_spline.py @@ -3,6 +3,7 @@ import numpy as np from f1tenth_gym.envs.track import cubic_spline + class TestCubicSpline(unittest.TestCase): def test_calc_curvature(self): circle_x = np.cos(np.linspace(0, 2 * np.pi, 100))[:-1] @@ -25,17 +26,27 @@ def test_calc_yaw(self): self.assertAlmostEqual(track.calc_yaw(np.pi / 2), np.pi, places=2) self.assertAlmostEqual(track.calc_yaw(np.pi), 3 * np.pi / 2, places=2) self.assertAlmostEqual(track.calc_yaw(3 * np.pi / 2), 0, places=2) - + def test_calc_position(self): circle_x = np.cos(np.linspace(0, 2 * np.pi, 100))[:-1] circle_y = np.sin(np.linspace(0, 2 * np.pi, 100))[:-1] track = cubic_spline.CubicSpline2D(circle_x, circle_y) # Test the position at the four corners of the circle # The position of a circle is (x, y) = (cos(s), sin(s)) - self.assertTrue(np.allclose(track.calc_position(0), np.array([1, 0]), atol=1e-3)) - self.assertTrue(np.allclose(track.calc_position(np.pi / 2), np.array([0, 1]), atol=1e-3)) - self.assertTrue(np.allclose(track.calc_position(np.pi), np.array([-1, 0]), atol=1e-3)) - self.assertTrue(np.allclose(track.calc_position(3 * np.pi / 2), np.array([0, -1]), atol=1e-3)) + self.assertTrue( + np.allclose(track.calc_position(0), np.array([1, 0]), atol=1e-3) + ) + self.assertTrue( + np.allclose(track.calc_position(np.pi / 2), np.array([0, 1]), atol=1e-3) + ) + self.assertTrue( + np.allclose(track.calc_position(np.pi), np.array([-1, 0]), atol=1e-3) + ) + self.assertTrue( + np.allclose( + track.calc_position(3 * np.pi / 2), np.array([0, -1]), atol=1e-3 + ) + ) def test_calc_arclength(self): circle_x = np.cos(np.linspace(0, 2 * np.pi, 100))[:-1] @@ -43,9 +54,13 @@ def test_calc_arclength(self): track = cubic_spline.CubicSpline2D(circle_x, circle_y) # Test the arclength at the four corners of the circle self.assertAlmostEqual(track.calc_arclength(1, 0, 0)[0], 0, places=2) - self.assertAlmostEqual(track.calc_arclength(0, 1, 0)[0], np.pi/2, places=2) - self.assertAlmostEqual(track.calc_arclength(-1, 0, np.pi/2)[0], np.pi, places=2) - self.assertAlmostEqual(track.calc_arclength(0, -1, np.pi)[0], 3*np.pi/2, places=2) + self.assertAlmostEqual(track.calc_arclength(0, 1, 0)[0], np.pi / 2, places=2) + self.assertAlmostEqual( + track.calc_arclength(-1, 0, np.pi / 2)[0], np.pi, places=2 + ) + self.assertAlmostEqual( + track.calc_arclength(0, -1, np.pi)[0], 3 * np.pi / 2, places=2 + ) def test_calc_arclength_inaccurate(self): circle_x = np.cos(np.linspace(0, 2 * np.pi, 100))[:-1] @@ -53,6 +68,12 @@ def test_calc_arclength_inaccurate(self): track = cubic_spline.CubicSpline2D(circle_x, circle_y) # Test the arclength at the four corners of the circle self.assertAlmostEqual(track.calc_arclength_inaccurate(1, 0)[0], 0, places=2) - self.assertAlmostEqual(track.calc_arclength_inaccurate(0, 1)[0], np.pi/2, places=2) - self.assertAlmostEqual(track.calc_arclength_inaccurate(-1, 0)[0], np.pi, places=2) - self.assertAlmostEqual(track.calc_arclength_inaccurate(0, -1)[0], 3*np.pi/2, places=2) + self.assertAlmostEqual( + track.calc_arclength_inaccurate(0, 1)[0], np.pi / 2, places=2 + ) + self.assertAlmostEqual( + track.calc_arclength_inaccurate(-1, 0)[0], np.pi, places=2 + ) + self.assertAlmostEqual( + track.calc_arclength_inaccurate(0, -1)[0], 3 * np.pi / 2, places=2 + )