Added brute force method for testing grid rotation

guided_mrmp/controllers/place_grid.py

 import cvxpy as cp
 import numpy as np
 
-def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5)):
+def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5), return_loss=False):
     """
         Place a grid to cover robot locations with alignment to centers.
 
@@ -11,6 +11,8 @@ def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5)):
             - grid_shape (tuple): (# of rows, # of columns) of the grid
         outputs:
             - origin (tuple): bottom-left corner of the grid in continuous space
+            - cell_centers (list): centers of grid cells for each robot (same order as robot_locations)
+            - loss: when return_loss=True, sum of squared differences loss
     """
     robot_locations = np.array(robot_locations)
     N = len(robot_locations)
@@ -70,17 +72,80 @@ def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5)):
     if prob.status not in ["optimal", "optimal_inaccurate"]:
         print("Problem could not be solved to optimality.")
         return None
+    
+    if return_loss:
+        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 
+
+        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
+        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)
+            - angle_info (dict): when return_angle_info=True, dict of { grid rotation angle : ( grid_center, cell_centers, loss ) for that angle
+    """ 
+    
+    # Dictionary of { grid rotation angle : ( grid_center, cell_centers, loss ) for that angle }
+    angle_info = {}
+    
+    # The rotation angle from the angles sampled which minimizes loss
+    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)
+        rotation_matrix = np.array([[np.cos(angle), -np.sin(angle)],
+                                    [np.sin(angle), np.cos(angle)]])
         
-def main():
-    np.random.seed(54)
+        # Rotate robot locations about the origin in continuous space
+        rotated_robot_locations = robot_locations @ rotation_matrix
+        
+        # Run place_grid on transformed robot locations
+        rotated_origin, rotated_cell_centers, loss = place_grid(rotated_robot_locations, cell_size, grid_shape, return_loss=True)
+        rotated_grid_center = rotated_origin + np.array(grid_shape) * cell_size / 2
+        
+        # Undo the rotation transformation for the origin and cell centers 
+        grid_center = rotated_grid_center @ np.linalg.inv(rotation_matrix)
+        cell_centers = rotated_cell_centers @ np.linalg.inv(rotation_matrix)
+        
+        # Check if the loss is the smallest observed so far
+        if min_loss is None or loss < min_loss:
+            min_loss_angle = angle
+            min_loss = loss
+            
+        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]
+
+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)
     
-    origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape)
-    print("Grid Origin (Bottom-Left Corner):", origin)
-    print(cell_centers)
+    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)
 
     import matplotlib.pyplot as plt
     
@@ -112,4 +177,14 @@ def main():
     plt.show()
 
 if __name__ == "__main__":
-    main()
+    import argparse
\ No newline at end of file
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--rotations", 
+        type=bool, 
+        required=True
+    )
+    args = parser.parse_args()
+
+    main(args.rotations)
\ No newline at end of file