diff --git a/Localization/extended_kalman_filter/extended_kalman_filter.py b/Localization/extended_kalman_filter/extended_kalman_filter.py index c2fe8e3a..d9ece6c6 100644 --- a/Localization/extended_kalman_filter/extended_kalman_filter.py +++ b/Localization/extended_kalman_filter/extended_kalman_filter.py @@ -7,13 +7,14 @@ author: Atsushi Sakai (@Atsushi_twi) """ import sys import pathlib + sys.path.append(str(pathlib.Path(__file__).parent.parent.parent)) import math import matplotlib.pyplot as plt import numpy as np -from utils.angle import rot_mat_2d +from utils.plot import plot_covariance_ellipse # Covariance for EKF simulation Q = np.diag([ @@ -135,29 +136,6 @@ def ekf_estimation(xEst, PEst, z, u): return xEst, PEst -def plot_covariance_ellipse(xEst, PEst): # pragma: no cover - Pxy = PEst[0:2, 0:2] - eigval, eigvec = np.linalg.eig(Pxy) - - if eigval[0] >= eigval[1]: - bigind = 0 - smallind = 1 - else: - bigind = 1 - smallind = 0 - - t = np.arange(0, 2 * math.pi + 0.1, 0.1) - a = math.sqrt(eigval[bigind]) - b = math.sqrt(eigval[smallind]) - x = [a * math.cos(it) for it in t] - y = [b * math.sin(it) for it in t] - angle = math.atan2(eigvec[1, bigind], eigvec[0, bigind]) - fx = rot_mat_2d(angle) @ (np.array([x, y])) - px = np.array(fx[0, :] + xEst[0, 0]).flatten() - py = np.array(fx[1, :] + xEst[1, 0]).flatten() - plt.plot(px, py, "--r") - - def main(): print(__file__ + " start!!") @@ -202,7 +180,7 @@ def main(): hxDR[1, :].flatten(), "-k") plt.plot(hxEst[0, :].flatten(), hxEst[1, :].flatten(), "-r") - plot_covariance_ellipse(xEst, PEst) + plot_covariance_ellipse(xEst[0, 0], xEst[1, 0], PEst) plt.axis("equal") plt.grid(True) plt.pause(0.001) diff --git a/Mapping/grid_map_lib/grid_map_lib.py b/Mapping/grid_map_lib/grid_map_lib.py index 85dd76e7..d08d8ec5 100644 --- a/Mapping/grid_map_lib/grid_map_lib.py +++ b/Mapping/grid_map_lib/grid_map_lib.py @@ -5,22 +5,42 @@ Grid map library in python author: Atsushi Sakai """ - +from functools import total_ordering import matplotlib.pyplot as plt import numpy as np +@total_ordering +class FloatGrid: + + def __init__(self, init_val=0.0): + self.data = init_val + + def get_float_data(self): + return self.data + + def __eq__(self, other): + if not isinstance(other, FloatGrid): + return NotImplemented + return self.get_float_data() == other.get_float_data() + + def __lt__(self, other): + if not isinstance(other, FloatGrid): + return NotImplemented + return self.get_float_data() < other.get_float_data() + + class GridMap: """ GridMap class """ def __init__(self, width, height, resolution, - center_x, center_y, init_val=0.0): + center_x, center_y, init_val=FloatGrid(0.0)): """__init__ :param width: number of grid for width - :param height: number of grid for heigt + :param height: number of grid for height :param resolution: grid resolution [m] :param center_x: center x position [m] :param center_y: center y position [m] @@ -35,8 +55,9 @@ class GridMap: self.left_lower_x = self.center_x - self.width / 2.0 * self.resolution self.left_lower_y = self.center_y - self.height / 2.0 * self.resolution - self.ndata = self.width * self.height - self.data = [init_val] * self.ndata + self.n_data = self.width * self.height + self.data = [init_val] * self.n_data + self.data_type = type(init_val) def get_value_from_xy_index(self, x_ind, y_ind): """get_value_from_xy_index @@ -49,7 +70,7 @@ class GridMap: grid_ind = self.calc_grid_index_from_xy_index(x_ind, y_ind) - if 0 <= grid_ind < self.ndata: + if 0 <= grid_ind < self.n_data: return self.data[grid_ind] else: return None @@ -101,7 +122,7 @@ class GridMap: grid_ind = int(y_ind * self.width + x_ind) - if 0 <= grid_ind < self.ndata: + if 0 <= grid_ind < self.n_data and isinstance(val, self.data_type): self.data[grid_ind] = val return True # OK else: @@ -138,6 +159,27 @@ class GridMap: grid_ind = int(y_ind * self.width + x_ind) return grid_ind + def calc_xy_index_from_grid_index(self, grid_ind): + y_ind, x_ind = divmod(grid_ind, self.width) + return x_ind, y_ind + + def calc_grid_index_from_xy_pos(self, x_pos, y_pos): + """get_xy_index_from_xy_pos + + :param x_pos: x position [m] + :param y_pos: y position [m] + """ + x_ind = self.calc_xy_index_from_position( + x_pos, self.left_lower_x, self.width) + y_ind = self.calc_xy_index_from_position( + y_pos, self.left_lower_y, self.height) + + return self.calc_grid_index_from_xy_index(x_ind, y_ind) + + def calc_grid_central_xy_position_from_grid_index(self, grid_ind): + x_ind, y_ind = self.calc_xy_index_from_grid_index(grid_ind) + return self.calc_grid_central_xy_position_from_xy_index(x_ind, y_ind) + def calc_grid_central_xy_position_from_xy_index(self, x_ind, y_ind): x_pos = self.calc_grid_central_xy_position_from_index( x_ind, self.left_lower_x) @@ -156,39 +198,40 @@ class GridMap: else: return None - def check_occupied_from_xy_index(self, xind, yind, occupied_val=1.0): + def check_occupied_from_xy_index(self, x_ind, y_ind, occupied_val): - val = self.get_value_from_xy_index(xind, yind) + val = self.get_value_from_xy_index(x_ind, y_ind) if val is None or val >= occupied_val: return True else: return False - def expand_grid(self): - xinds, yinds = [], [] + def expand_grid(self, occupied_val=FloatGrid(1.0)): + x_inds, y_inds, values = [], [], [] for ix in range(self.width): for iy in range(self.height): - if self.check_occupied_from_xy_index(ix, iy): - xinds.append(ix) - yinds.append(iy) + if self.check_occupied_from_xy_index(ix, iy, occupied_val): + x_inds.append(ix) + y_inds.append(iy) + values.append(self.get_value_from_xy_index(ix, iy)) - for (ix, iy) in zip(xinds, yinds): - self.set_value_from_xy_index(ix + 1, iy, val=1.0) - self.set_value_from_xy_index(ix, iy + 1, val=1.0) - self.set_value_from_xy_index(ix + 1, iy + 1, val=1.0) - self.set_value_from_xy_index(ix - 1, iy, val=1.0) - self.set_value_from_xy_index(ix, iy - 1, val=1.0) - self.set_value_from_xy_index(ix - 1, iy - 1, val=1.0) + for (ix, iy, value) in zip(x_inds, y_inds, values): + self.set_value_from_xy_index(ix + 1, iy, val=value) + self.set_value_from_xy_index(ix, iy + 1, val=value) + self.set_value_from_xy_index(ix + 1, iy + 1, val=value) + self.set_value_from_xy_index(ix - 1, iy, val=value) + self.set_value_from_xy_index(ix, iy - 1, val=value) + self.set_value_from_xy_index(ix - 1, iy - 1, val=value) @staticmethod def check_inside_polygon(iox, ioy, x, y): - npoint = len(x) - 1 + n_point = len(x) - 1 inside = False - for i1 in range(npoint): - i2 = (i1 + 1) % (npoint + 1) + for i1 in range(n_point): + i2 = (i1 + 1) % (n_point + 1) if x[i1] >= x[i2]: min_x, max_x = x[i2], x[i1] @@ -211,27 +254,26 @@ class GridMap: print("center_y:", self.center_y) print("left_lower_x:", self.left_lower_x) print("left_lower_y:", self.left_lower_y) - print("ndata:", self.ndata) + print("n_data:", self.n_data) def plot_grid_map(self, ax=None): - - grid_data = np.reshape(np.array(self.data), (self.height, self.width)) + float_data_array = np.array([d.get_float_data() for d in self.data]) + grid_data = np.reshape(float_data_array, (self.height, self.width)) if not ax: fig, ax = plt.subplots() heat_map = ax.pcolor(grid_data, cmap="Blues", vmin=0.0, vmax=1.0) plt.axis("equal") - # plt.show() return heat_map -def test_polygon_set(): +def polygon_set_demo(): ox = [0.0, 4.35, 20.0, 50.0, 100.0, 130.0, 40.0] oy = [0.0, -4.15, -20.0, 0.0, 30.0, 60.0, 80.0] grid_map = GridMap(600, 290, 0.7, 60.0, 30.5) - grid_map.set_value_from_polygon(ox, oy, 1.0, inside=False) + grid_map.set_value_from_polygon(ox, oy, FloatGrid(1.0), inside=False) grid_map.plot_grid_map() @@ -239,24 +281,27 @@ def test_polygon_set(): plt.grid(True) -def test_position_set(): +def position_set_demo(): grid_map = GridMap(100, 120, 0.5, 10.0, -0.5) - grid_map.set_value_from_xy_pos(10.1, -1.1, 1.0) - grid_map.set_value_from_xy_pos(10.1, -0.1, 1.0) - grid_map.set_value_from_xy_pos(10.1, 1.1, 1.0) - grid_map.set_value_from_xy_pos(11.1, 0.1, 1.0) - grid_map.set_value_from_xy_pos(10.1, 0.1, 1.0) - grid_map.set_value_from_xy_pos(9.1, 0.1, 1.0) + grid_map.set_value_from_xy_pos(10.1, -1.1, FloatGrid(1.0)) + grid_map.set_value_from_xy_pos(10.1, -0.1, FloatGrid(1.0)) + grid_map.set_value_from_xy_pos(10.1, 1.1, FloatGrid(1.0)) + grid_map.set_value_from_xy_pos(11.1, 0.1, FloatGrid(1.0)) + grid_map.set_value_from_xy_pos(10.1, 0.1, FloatGrid(1.0)) + grid_map.set_value_from_xy_pos(9.1, 0.1, FloatGrid(1.0)) grid_map.plot_grid_map() + plt.axis("equal") + plt.grid(True) + def main(): print("start!!") - test_position_set() - test_polygon_set() + position_set_demo() + polygon_set_demo() plt.show() diff --git a/Mapping/ndt_map/ndt_map.py b/Mapping/ndt_map/ndt_map.py new file mode 100644 index 00000000..f4f32996 --- /dev/null +++ b/Mapping/ndt_map/ndt_map.py @@ -0,0 +1,135 @@ +""" +Normal Distribution Transform (NDTGrid) mapping sample +""" +import matplotlib.pyplot as plt +import numpy as np +from collections import defaultdict + +from Mapping.grid_map_lib.grid_map_lib import GridMap +from utils.plot import plot_covariance_ellipse + + +class NDTMap: + """ + Normal Distribution Transform (NDT) map class + + :param ox: obstacle x position list + :param oy: obstacle y position list + :param resolution: grid resolution [m] + """ + + class NDTGrid: + """ + NDT grid + """ + + def __init__(self): + #: Number of points in the NDTGrid grid + self.n_points = 0 + #: Mean x position of points in the NDTGrid cell + self.mean_x = None + #: Mean y position of points in the NDTGrid cell + self.mean_y = None + #: Center x position of the NDT grid + self.center_grid_x = None + #: Center y position of the NDT grid + self.center_grid_y = None + #: Covariance matrix of the NDT grid + self.covariance = None + #: Eigen vectors of the NDT grid + self.eig_vec = None + #: Eigen values of the NDT grid + self.eig_values = None + + def __init__(self, ox, oy, resolution): + #: Minimum number of points in the NDT grid + self.min_n_points = 3 + #: Resolution of the NDT grid [m] + self.resolution = resolution + width = int((max(ox) - min(ox))/resolution) + 3 # rounding up + right and left margin + height = int((max(oy) - min(oy))/resolution) + 3 + center_x = np.mean(ox) + center_y = np.mean(oy) + self.ox = ox + self.oy = oy + #: NDT grid index map + self.grid_index_map = self._create_grid_index_map(ox, oy) + + #: NDT grid map. Each grid contains NDTGrid object + self._construct_grid_map(center_x, center_y, height, ox, oy, resolution, width) + + def _construct_grid_map(self, center_x, center_y, height, ox, oy, resolution, width): + self.grid_map = GridMap(width, height, resolution, center_x, center_y, self.NDTGrid()) + for grid_index, inds in self.grid_index_map.items(): + ndt = self.NDTGrid() + ndt.n_points = len(inds) + if ndt.n_points >= self.min_n_points: + ndt.mean_x = np.mean(ox[inds]) + ndt.mean_y = np.mean(oy[inds]) + ndt.center_grid_x, ndt.center_grid_y = \ + self.grid_map.calc_grid_central_xy_position_from_grid_index(grid_index) + ndt.covariance = np.cov(ox[inds], oy[inds]) + ndt.eig_values, ndt.eig_vec = np.linalg.eig(ndt.covariance) + self.grid_map.data[grid_index] = ndt + + def _create_grid_index_map(self, ox, oy): + grid_index_map = defaultdict(list) + for i in range(len(ox)): + grid_index = self.grid_map.calc_grid_index_from_xy_pos(ox[i], oy[i]) + grid_index_map[grid_index].append(i) + return grid_index_map + + +def create_dummy_observation_data(): + ox = [] + oy = [] + # left corridor + for y in range(-50, 50): + ox.append(-20.0) + oy.append(y) + # right corridor 1 + for y in range(-50, 0): + ox.append(20.0) + oy.append(y) + # right corridor 2 + for x in range(20, 50): + ox.append(x) + oy.append(0) + # right corridor 3 + for x in range(20, 50): + ox.append(x) + oy.append(x/2.0+10) + # right corridor 4 + for y in range(20, 50): + ox.append(20) + oy.append(y) + ox = np.array(ox) + oy = np.array(oy) + # Adding random noize + ox += np.random.rand(len(ox)) * 1.0 + oy += np.random.rand(len(ox)) * 1.0 + return ox, oy + + +def main(): + print(__file__ + " start!!") + + ox, oy = create_dummy_observation_data() + grid_resolution = 10.0 + ndt_map = NDTMap(ox, oy, grid_resolution) + + # plot raw observation + plt.plot(ox, oy, ".r") + + # plot grid clustering + [plt.plot(ox[inds], oy[inds], "x") for inds in ndt_map.grid_index_map.values()] + + # plot ndt grid map + [plot_covariance_ellipse(ndt.mean_x, ndt.mean_y, ndt.covariance, color="-k") for ndt in ndt_map.grid_map.data if ndt.n_points > 0] + + plt.axis("equal") + plt.show() + + +if __name__ == '__main__': + main() diff --git a/PathPlanning/GridBasedSweepCPP/grid_based_sweep_coverage_path_planner.py b/PathPlanning/GridBasedSweepCPP/grid_based_sweep_coverage_path_planner.py index 31dc9175..ee192e92 100644 --- a/PathPlanning/GridBasedSweepCPP/grid_based_sweep_coverage_path_planner.py +++ b/PathPlanning/GridBasedSweepCPP/grid_based_sweep_coverage_path_planner.py @@ -13,7 +13,7 @@ import pathlib sys.path.append(str(pathlib.Path(__file__).parent.parent.parent)) from utils.angle import rot_mat_2d -from Mapping.grid_map_lib.grid_map_lib import GridMap +from Mapping.grid_map_lib.grid_map_lib import GridMap, FloatGrid do_animation = True @@ -41,8 +41,7 @@ class SweepSearcher: n_y_index = c_y_index # found safe grid - if not grid_map.check_occupied_from_xy_index(n_x_index, n_y_index, - occupied_val=0.5): + if not self.check_occupied(n_x_index, n_y_index, grid_map): return n_x_index, n_y_index else: # occupied next_c_x_index, next_c_y_index = self.find_safe_turning_grid( @@ -51,19 +50,20 @@ class SweepSearcher: # moving backward next_c_x_index = -self.moving_direction + c_x_index next_c_y_index = c_y_index - if grid_map.check_occupied_from_xy_index(next_c_x_index, - next_c_y_index): + if self.check_occupied(next_c_x_index, next_c_y_index, grid_map): # moved backward, but the grid is occupied by obstacle return None, None else: # keep moving until end - while not grid_map.check_occupied_from_xy_index( - next_c_x_index + self.moving_direction, - next_c_y_index, occupied_val=0.5): + while not self.check_occupied(next_c_x_index + self.moving_direction, next_c_y_index, grid_map): next_c_x_index += self.moving_direction self.swap_moving_direction() return next_c_x_index, next_c_y_index + @staticmethod + def check_occupied(c_x_index, c_y_index, grid_map): + return grid_map.check_occupied_from_xy_index(c_x_index, c_y_index, FloatGrid(0.5)) + def find_safe_turning_grid(self, c_x_index, c_y_index, grid_map): for (d_x_ind, d_y_ind) in self.turing_window: @@ -72,17 +72,14 @@ class SweepSearcher: next_y_ind = d_y_ind + c_y_index # found safe grid - if not grid_map.check_occupied_from_xy_index(next_x_ind, - next_y_ind, - occupied_val=0.5): + if not self.check_occupied(next_x_ind, next_y_ind, grid_map): return next_x_ind, next_y_ind return None, None def is_search_done(self, grid_map): for ix in self.x_indexes_goal_y: - if not grid_map.check_occupied_from_xy_index(ix, self.goal_y, - occupied_val=0.5): + if not self.check_occupied(ix, self.goal_y, grid_map): return False # all lower grid is occupied @@ -168,7 +165,7 @@ def search_free_grid_index_at_edge_y(grid_map, from_upper=False): for iy in x_range: for ix in y_range: - if not grid_map.check_occupied_from_xy_index(ix, iy): + if not SweepSearcher.check_occupied(ix, iy, grid_map): y_index = iy x_indexes.append(ix) if y_index: @@ -185,7 +182,7 @@ def setup_grid_map(ox, oy, resolution, sweep_direction, offset_grid=10): grid_map = GridMap(width, height, resolution, center_x, center_y) grid_map.print_grid_map_info() - grid_map.set_value_from_polygon(ox, oy, 1.0, inside=False) + grid_map.set_value_from_polygon(ox, oy, FloatGrid(1.0), inside=False) grid_map.expand_grid() x_inds_goal_y = [] @@ -203,7 +200,7 @@ def setup_grid_map(ox, oy, resolution, sweep_direction, offset_grid=10): def sweep_path_search(sweep_searcher, grid_map, grid_search_animation=False): # search start grid c_x_index, c_y_index = sweep_searcher.search_start_grid(grid_map) - if not grid_map.set_value_from_xy_index(c_x_index, c_y_index, 0.5): + if not grid_map.set_value_from_xy_index(c_x_index, c_y_index, FloatGrid(0.5)): print("Cannot find start grid") return [], [] @@ -235,7 +232,7 @@ def sweep_path_search(sweep_searcher, grid_map, grid_search_animation=False): px.append(x) py.append(y) - grid_map.set_value_from_xy_index(c_x_index, c_y_index, 0.5) + grid_map.set_value_from_xy_index(c_x_index, c_y_index, FloatGrid(0.5)) if grid_search_animation: grid_map.plot_grid_map(ax=ax) diff --git a/docs/modules/mapping/mapping_main.rst b/docs/modules/mapping/mapping_main.rst index 1c02c75c..a98acaaf 100644 --- a/docs/modules/mapping/mapping_main.rst +++ b/docs/modules/mapping/mapping_main.rst @@ -7,6 +7,7 @@ Mapping :caption: Contents gaussian_grid_map/gaussian_grid_map + ndt_map/ndt_map ray_casting_grid_map/ray_casting_grid_map lidar_to_grid_map_tutorial/lidar_to_grid_map_tutorial point_cloud_sampling/point_cloud_sampling diff --git a/docs/modules/mapping/ndt_map/grid_clustering.png b/docs/modules/mapping/ndt_map/grid_clustering.png new file mode 100644 index 00000000..c5359028 Binary files /dev/null and b/docs/modules/mapping/ndt_map/grid_clustering.png differ diff --git a/docs/modules/mapping/ndt_map/ndt_map1.png b/docs/modules/mapping/ndt_map/ndt_map1.png new file mode 100644 index 00000000..1f34b9a6 Binary files /dev/null and b/docs/modules/mapping/ndt_map/ndt_map1.png differ diff --git a/docs/modules/mapping/ndt_map/ndt_map2.png b/docs/modules/mapping/ndt_map/ndt_map2.png new file mode 100644 index 00000000..b9970629 Binary files /dev/null and b/docs/modules/mapping/ndt_map/ndt_map2.png differ diff --git a/docs/modules/mapping/ndt_map/ndt_map_main.rst b/docs/modules/mapping/ndt_map/ndt_map_main.rst new file mode 100644 index 00000000..bfc39363 --- /dev/null +++ b/docs/modules/mapping/ndt_map/ndt_map_main.rst @@ -0,0 +1,53 @@ +.. _ndt_map: + +Normal Distance Transform (NDT) map +------------------------------------ + +This is a NDT mapping example. + +Normal Distribution Transform (NDT) is a map representation that uses normal distribution for observation point modeling. + +Normal Distribution +~~~~~~~~~~~~~~~~~~~~~ + +Normal distribution consists of two parameters: mean :math:`\mu` and covariance :math:`\Sigma`. + +:math:`\mathbf{X} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})` + +In the 2D case, :math:`\boldsymbol{\mu}` is a 2D vector and :math:`\boldsymbol{\Sigma}` is a 2x2 matrix. + +In the matrix form, the probability density function of thr normal distribution is: + +:math:`X=\frac{1}{\sqrt{(2 \pi)^2|\Sigma|}} \exp \left\{-\frac{1}{2}^t(x-\mu) \sum^{-1}(x-\mu)\right\}` + +Normal Distance Transform mapping steps +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  + +NDT mapping consists of two steps: + +When we have a new observation like this: + +.. figure:: raw_observations.png + +First, we need to cluster the observation points. +This is done by using a grid based clustering algorithm. + +The result is: + +.. figure:: grid_clustering.png + +Then, we need to fit a normal distribution to each grid cluster. + +Black ellipse shows each NDT grid like this: + +.. figure:: ndt_map1.png + +.. figure:: ndt_map2.png + +API +~~~~~ + +.. autoclass:: Mapping.ndt_map.ndt_map.NDTMap + :members: + :class-doc-from: class + diff --git a/docs/modules/mapping/ndt_map/raw_observations.png b/docs/modules/mapping/ndt_map/raw_observations.png new file mode 100644 index 00000000..c1a0dd47 Binary files /dev/null and b/docs/modules/mapping/ndt_map/raw_observations.png differ diff --git a/tests/test_grid_map_lib.py b/tests/test_grid_map_lib.py index 92ca67e2..670b357a 100644 --- a/tests/test_grid_map_lib.py +++ b/tests/test_grid_map_lib.py @@ -25,5 +25,16 @@ def test_polygon_set(): 1.0, inside=False) +def test_xy_and_grid_index_conversion(): + grid_map = GridMap(100, 120, 0.5, 10.0, -0.5) + + for x_ind in range(grid_map.width): + for y_ind in range(grid_map.height): + grid_ind = grid_map.calc_grid_index_from_xy_index(x_ind, y_ind) + x_ind_2, y_ind_2 = grid_map.calc_xy_index_from_grid_index(grid_ind) + assert x_ind == x_ind_2 + assert y_ind == y_ind_2 + + if __name__ == '__main__': conftest.run_this_test(__file__) diff --git a/utils/plot.py b/utils/plot.py index d4bbe29f..eb5aa7ad 100644 --- a/utils/plot.py +++ b/utils/plot.py @@ -9,6 +9,69 @@ from matplotlib.patches import FancyArrowPatch from mpl_toolkits.mplot3d.proj3d import proj_transform from mpl_toolkits.mplot3d import Axes3D +from utils.angle import rot_mat_2d + + +def plot_covariance_ellipse(x, y, cov, chi2=3.0, color="-r", ax=None): + """ + This function plots an ellipse that represents a covariance matrix. The ellipse is centered at (x, y) and its shape, size and rotation are determined by the covariance matrix. + + Parameters: + x : (float) The x-coordinate of the center of the ellipse. + y : (float) The y-coordinate of the center of the ellipse. + cov : (numpy.ndarray) A 2x2 covariance matrix that determines the shape, size, and rotation of the ellipse. + chi2 : (float, optional) A scalar value that scales the ellipse size. This value is typically set based on chi-squared distribution quantiles to achieve certain confidence levels (e.g., 3.0 corresponds to ~95% confidence for a 2D Gaussian). Defaults to 3.0. + color : (str, optional) The color and line style of the ellipse plot, following matplotlib conventions. Defaults to "-r" (a red solid line). + ax : (matplotlib.axes.Axes, optional) The Axes object to draw the ellipse on. If None (default), a new figure and axes are created. + + Returns: + None. This function plots the covariance ellipse on the specified axes. + """ + eig_val, eig_vec = np.linalg.eig(cov) + + if eig_val[0] >= eig_val[1]: + big_ind = 0 + small_ind = 1 + else: + big_ind = 1 + small_ind = 0 + a = math.sqrt(chi2 * eig_val[big_ind]) + b = math.sqrt(chi2 * eig_val[small_ind]) + angle = math.atan2(eig_vec[1, big_ind], eig_vec[0, big_ind]) + plot_ellipse(x, y, a, b, angle, color=color, ax=ax) + + +def plot_ellipse(x, y, a, b, angle, color="-r", ax=None, **kwargs): + """ + This function plots an ellipse based on the given parameters. + + Parameters + ---------- + x : (float) The x-coordinate of the center of the ellipse. + y : (float) The y-coordinate of the center of the ellipse. + a : (float) The length of the semi-major axis of the ellipse. + b : (float) The length of the semi-minor axis of the ellipse. + angle : (float) The rotation angle of the ellipse, in radians. + color : (str, optional) The color and line style of the ellipse plot, following matplotlib conventions. Defaults to "-r" (a red solid line). + ax : (matplotlib.axes.Axes, optional) The Axes object to draw the ellipse on. If None (default), a new figure and axes are created. + **kwargs: Additional keyword arguments to pass to plt.plot or ax.plot. + + Returns + --------- + None. This function plots the ellipse based on the specified parameters. + """ + + t = np.arange(0, 2 * math.pi + 0.1, 0.1) + px = [a * math.cos(it) for it in t] + py = [b * math.sin(it) for it in t] + fx = rot_mat_2d(angle) @ (np.array([px, py])) + px = np.array(fx[0, :] + x).flatten() + py = np.array(fx[1, :] + y).flatten() + if ax is None: + plt.plot(px, py, color, **kwargs) + else: + ax.plot(px, py, color, **kwargs) + def plot_arrow(x, y, yaw, arrow_length=1.0, origin_point_plot_style="xr", @@ -132,7 +195,6 @@ def plot_3d_vector_arrow(ax, p1, p2): ) - def plot_triangle(p1, p2, p3, ax): ax.add_collection3d(art3d.Poly3DCollection([[p1, p2, p3]], color='b')) @@ -163,3 +225,10 @@ def set_equal_3d_axis(ax, x_lims, y_lims, z_lims): ax.set_xlim(mid_x - max_range, mid_x + max_range) ax.set_ylim(mid_y - max_range, mid_y + max_range) ax.set_zlim(mid_z - max_range, mid_z + max_range) + + +if __name__ == '__main__': + plot_ellipse(0, 0, 1, 2, np.deg2rad(15)) + plt.axis('equal') + plt.show() +