Added plotting function for rotated grids and command line args

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