From c78eda078266463bbf56d503f790296757dc78de Mon Sep 17 00:00:00 2001
From: Adam Sitabkhan <adam.sitabkhan@gmail.com>
Date: Sat, 11 Jan 2025 12:59:29 -0600
Subject: [PATCH] Added plotting function for rotated grids and command line
args
---
guided_mrmp/controllers/place_grid.py | 112 ++++++++++++++++++--------
1 file changed, 78 insertions(+), 34 deletions(-)
diff --git a/guided_mrmp/controllers/place_grid.py b/guided_mrmp/controllers/place_grid.py
index 192ad8f..a14fbf5 100644
--- a/guided_mrmp/controllers/place_grid.py
+++ b/guided_mrmp/controllers/place_grid.py
@@ -77,15 +77,16 @@ def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5), return_loss=Fals
return origin.value, cell_centers.value, prob.value
return origin.value, cell_centers.value
+
def place_grid_with_rotation(robot_locations, cell_size=1, grid_shape=(5, 5), num_angles=18, return_angle_info=False):
"""
- Place a grid to cover robot locations with alignment to centers. Allows rotations of the grid, but only samples
+ Place a grid to cover robot locations with alignment to centers.
inputs:
- robot_locations (list): locations of robots involved in conflict [[x,y], [x,y], ...]
- cell_size (float): the width of each grid cell in continuous space
- grid_shape (tuple): (# of rows, # of columns) of the grid
- - num_angles (int): number of evenly spaced angles to sample between 0 and 90 degrees
+ - sample_size (int): number of evenly spaced angles to sample between 0 and pi/2 radians
outputs:
- grid_center (tuple): center of the grid in continuous space
- cell_centers (list): centers of grid cells for each robot (same order as robot_locations)
@@ -99,8 +100,9 @@ def place_grid_with_rotation(robot_locations, cell_size=1, grid_shape=(5, 5), nu
min_loss_angle = 0
min_loss = None
- for angle in np.linspace(start=0, stop=np.pi/2, num=num_angles, endpoint=False):
- print(angle)
+ from tqdm import tqdm
+
+ for angle in tqdm(np.linspace(start=0, stop=np.pi/2, num=num_angles, endpoint=False), desc="Processing..."):
rotation_matrix = np.array([[np.cos(angle), -np.sin(angle)],
[np.sin(angle), np.cos(angle)]])
@@ -123,33 +125,13 @@ def place_grid_with_rotation(robot_locations, cell_size=1, grid_shape=(5, 5), nu
angle_info[angle] = grid_center, cell_centers, loss
if return_angle_info:
- return angle_info[min_loss_angle][0], angle_info[min_loss_angle][1], angle_info
- return angle_info[min_loss_angle][0], angle_info[min_loss_angle][1]
+ return *angle_info[min_loss_angle][:2], min_loss_angle, angle_info
+ return *angle_info[min_loss_angle][:2], min_loss_angle
-def main(rotations):
- np.random.seed(52)
- robot_locations = np.random.uniform(low=0, high=5, size=(5, 2))
- cell_size = 1
- grid_shape = (5, 5)
-
- import matplotlib.pyplot as plt
-
- if rotations:
- grid_center, cell_centers, angle_info = place_grid_with_rotation(robot_locations, cell_size, grid_shape, num_angles=30, return_angle_info=True)
- # angles = angle_info.keys()
- # losses = [loss for (_, _, loss) in angle_info.values()]
-
- # plt.scatter(angles, losses, c='r')
-
- # plt.show()
- else:
- origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape)
- print("Grid Origin (Bottom-Left Corner):", origin)
- print(cell_centers)
+def plot_grid(origin, cell_size, grid_shape):
import matplotlib.pyplot as plt
- plt.figure(figsize=(4, 4))
# Draw the grid
for i in range(grid_shape[1] + 1):
# Draw vertical lines
@@ -159,7 +141,58 @@ def main(rotations):
# Draw horizontal lines
plt.plot([origin[0], origin[0] + grid_shape[1] * cell_size],
[origin[1] + i * cell_size, origin[1] + i * cell_size], 'k-')
+
+
+def plot_grid_with_rotation(grid_center, cell_size, grid_shape, rotation_angle):
+ rotation_matrix = np.array([[np.cos(rotation_angle), -np.sin(rotation_angle)],
+ [np.sin(rotation_angle), np.cos(rotation_angle)]])
+ grid_center = np.array(grid_center)
+
+ import matplotlib.pyplot as plt
+
+ # Draw the grid
+ for i in range(grid_shape[1] + 1):
+ # Construct vertical line endpoints
+ line_points = np.array([[(-grid_shape[1] / 2 + i) * cell_size, (-grid_shape[1] / 2 + i) * cell_size],
+ [-grid_shape[0] / 2, grid_shape[0] / 2]])
+ # Rotate vertical lines and draw
+ rotated_line_points = rotation_matrix @ line_points + grid_center[:, None]
+ plt.plot(rotated_line_points[0], rotated_line_points[1], 'k-')
+ for i in range(grid_shape[1] + 1):
+ # Construct horizontal line endpoints
+ line_points = np.array([[-grid_shape[1] / 2, grid_shape[1] / 2],
+ [(-grid_shape[0] / 2 + i) * cell_size, (-grid_shape[0] / 2 + i) * cell_size]])
+ # Rotate horizontal lines and draw
+ rotated_line_points = rotation_matrix @ line_points + grid_center[:, None]
+ plt.plot(rotated_line_points[0], rotated_line_points[1], 'k-')
+
+def main(allow_rotations, num_robots, num_angles, seed):
+ if seed is not None:
+ np.random.seed(seed)
+
+ robot_locations = np.random.uniform(low=0, high=5, size=(num_robots, 2))
+ cell_size = 1
+ grid_shape = (5, 5)
+
+ import matplotlib.pyplot as plt
+
+ plt.figure(figsize=(6, 6))
+
+ if allow_rotations:
+ grid_center, cell_centers, rotation_angle = place_grid_with_rotation(robot_locations, cell_size, grid_shape, num_angles=num_angles)
+ print("Grid Center:", grid_center)
+ # angles = angle_info.keys()
+ # losses = [loss for (_, _, loss) in angle_info.values()]
+ # plt.scatter(angles, losses, c='r')
+
+ plot_grid_with_rotation(grid_center, cell_size, grid_shape, rotation_angle)
+ else:
+ origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape)
+ print("Grid Origin (Bottom-Left Corner):", origin)
+
+ plot_grid(origin, cell_size, grid_shape)
+
# Plot robot locations
robot_locations = np.array(robot_locations)
plt.scatter(robot_locations[:, 0], robot_locations[:, 1], c='r', label='Robot Locations')
@@ -167,10 +200,6 @@ def main(rotations):
# Plot cell centers
cell_centers = np.array(cell_centers)
plt.scatter(cell_centers[:, 0], cell_centers[:, 1], c='b', label='Cell Centers')
- for (cx, cy) in cell_centers:
- x = [cx - cell_size/2, cx + cell_size/2, cx + cell_size/2, cx - cell_size/2, cx - cell_size/2]
- y = [cy - cell_size/2, cy - cell_size/2, cy + cell_size/2, cy + cell_size/2, cy - cell_size/2]
- plt.plot(x, y, c='r')
plt.legend(loc='upper left')
@@ -181,10 +210,25 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
- "--rotations",
+ "--allow_rotations",
type=bool,
- required=True
+ default=False
+ )
+ parser.add_argument(
+ "--num_robots",
+ type=int,
+ default=2
+ )
+ parser.add_argument(
+ "--num_angles",
+ type=int,
+ default=18
+ )
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=None
)
args = parser.parse_args()
- main(args.rotations)
\ No newline at end of file
+ main(args.allow_rotations, args.num_robots, args.num_angles, args.seed)
\ No newline at end of file
--
GitLab