Update google_solver/convert_data.py

google_solver/convert_data.py

@@ -1,57 +1,77 @@
-
 import torch
 import numpy as np
 
 
 def convert_tensor(x):
-    x = x.long().numpy().astype('int')
+    """
+    Convert a PyTorch tensor to a nested Python list of integers.
 
-    if len(x.shape) == 1:
-        return list(x)
-    else:
-        return [list(x[i]) for i in range(x.shape[0])]
+    Args:
+        x (torch.Tensor): Tensor to convert.
 
+    Returns:
+        list: Converted nested list.
+    """
+    x = x.long().cpu().numpy().astype(int)
+    if x.ndim == 1:
+        return list(x)
+    return [list(row) for row in x]
 
 
 def make_time_windows(start_time, end_time):
+    """
+    Concatenate start and end time tensors to form time windows.
+
+    Args:
+        start_time (torch.Tensor): Start times (B, N, 1)
+        end_time (torch.Tensor): End times (B, N, 1)
+
+    Returns:
+        torch.Tensor: Time windows (B, N, 2)
+    """
     return torch.cat([start_time, end_time], dim=2)
 
 
+def convert_data(input_data, scale_factor):
+    """
+    Convert batched graph and fleet data to OR-Tools compatible format.
 
-def convert_data(input, scale_factor):
+    Args:
+        input_data (tuple): Tuple of (graph_data, fleet_data) as dictionaries.
+        scale_factor (float): Scaling factor to convert float to integer.
 
-    graph_data, fleet_data = input
+    Returns:
+        list: List of dictionaries, one per batch item, containing:
+              - distance_matrix
+              - time_matrix
+              - time_windows
+              - depot index (default 0)
+              - num_vehicles
+    """
+    graph_data, fleet_data = input_data
 
     start_times = graph_data['start_times']
     end_times = graph_data['end_times']
-
     distance_matrix = graph_data['distance_matrix']
     time_matrix = graph_data['time_matrix']
-
     time_windows = make_time_windows(start_times, end_times)
 
+    batch_size = distance_matrix.size(0)
+    converted_data = []
 
-    batch_size = distance_matrix.shape[0]
-    data = []
     for i in range(batch_size):
+        space_mat = (distance_matrix[i] * scale_factor)
+        time_mat = (time_matrix[i] * scale_factor)
+        windows = (time_windows[i] * scale_factor)
 
-        num_vehicles = distance_matrix[i].shape[1]
-
-        space_mat = distance_matrix[i] * scale_factor
-        time_mat = time_matrix[i] * scale_factor
-        windows = time_windows[i] * scale_factor
-
-        space_mat = convert_tensor(space_mat)
-        time_mat = convert_tensor(time_mat)
-        windows = convert_tensor(windows)
-
+        sample_dict = {
+            'distance_matrix': convert_tensor(space_mat),
+            'time_matrix': convert_tensor(time_mat),
+            'time_windows': convert_tensor(windows),
+            'depot': 0,
+            'num_vehicles': distance_matrix[i].shape[1]  # assuming square matrix
+        }
 
-        D = {'distance_matrix': space_mat,
-             'time_matrix': time_mat,
-             'time_windows': windows,
-             'depot': 0,
-             'num_vehicles': num_vehicles
-             }
+        converted_data.append(sample_dict)
 
-        data.append(D)
-    return data
+    return converted_data