Update google_solver/google_model.py

google_solver/google_model.py

@@ -1,52 +1,39 @@
 from __future__ import print_function
-from ortools.constraint_solver import routing_enums_pb2
-from ortools.constraint_solver import pywrapcp
+from ortools.constraint_solver import routing_enums_pb2, pywrapcp
 import torch
 from google_solver.convert_data import convert_data
 
 
-class GoogleActor(object):
+class GoogleActor:
+    """
+    Wrapper class to evaluate VRP solutions using Google's OR-Tools solver.
+    """
 
     def __init__(self, scale_factor=100):
-
-        if scale_factor is None:
-            self.scale_factor = 1
-        else:
-            self.scale_factor = scale_factor
-
+        self.scale_factor = scale_factor if scale_factor is not None else 1
 
     def __call__(self, input):
-
         drive_times = []
         data = convert_data(input, self.scale_factor)
-        batch_size = len(data)
         for datum in data:
             routing, assignment = self.compute_route(datum)
             total_time = self.compute_total_time(datum, routing, assignment)
             drive_times.append(total_time)
 
-        drive_times = torch.tensor(drive_times).float()
-        return drive_times
-
-
-    def compute_distance(self, routing, assignment, num_nodes):
-        """Prints solution on console."""
-        cumulative_route_distance = 0
-        for vehicle_id in range(num_nodes):
-            index = routing.Start(vehicle_id)
-            route_distance = 0
-            while not routing.IsEnd(index):
-                previous_index = index
-                index = assignment.Value(routing.NextVar(index))
-                route_distance += routing.GetArcCostForVehicle(
-                    previous_index, index, vehicle_id)
-            cumulative_route_distance += route_distance
+        return torch.tensor(drive_times).float()
 
-        cumulative_route_distance = cumulative_route_distance / self.scale_factor
-        return cumulative_route_distance
+    def compute_total_time(self, data, routing, assignment):
+        """
+        Computes the total time spent across all routes.
 
+        Args:
+            data (dict): Problem data with time matrix and vehicle count.
+            routing (RoutingModel): OR-Tools routing model.
+            assignment (Assignment): OR-Tools assignment solution.
 
-    def compute_total_time(self, data, routing, assignment):
+        Returns:
+            float: Total time (scaled back).
+        """
         time_dimension = routing.GetDimensionOrDie('Time')
         total_time = 0
         for vehicle_id in range(data['num_vehicles']):
@@ -55,74 +42,82 @@ class GoogleActor(object):
                 index = assignment.Value(routing.NextVar(index))
             time_var = time_dimension.CumulVar(index)
             total_time += assignment.Min(time_var)
-        total_time = total_time/self.scale_factor
-        return total_time
-
 
+        return total_time / self.scale_factor
 
     def compute_route(self, input):
+        """
+        Solves the routing problem using OR-Tools.
+
+        Args:
+            input (dict): Data containing distance, time matrix, time windows, and depot index.
 
+        Returns:
+            RoutingModel, Assignment: OR-Tools routing and solution.
+        """
         distance_matrix = input['distance_matrix']
         time_matrix = input['time_matrix']
         time_windows = input['time_windows']
         num_vehicles = input['num_vehicles']
         depot = input['depot']
 
-
         manager = pywrapcp.RoutingIndexManager(len(time_matrix), num_vehicles, depot)
         routing = pywrapcp.RoutingModel(manager)
 
         def time_callback(from_index, to_index):
-            """Returns the travel time between the two nodes."""
-            # Convert from routing variable Index to time matrix NodeIndex.
             from_node = manager.IndexToNode(from_index)
             to_node = manager.IndexToNode(to_index)
             return time_matrix[from_node][to_node]
 
         transit_callback_index = routing.RegisterTransitCallback(time_callback)
         routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)
-        time = 'Time'
+
         routing.AddDimension(
             transit_callback_index,
-            10000,  # allow waiting time
-            10000,  # maximum time per vehicle
-            False,  # Don't force start cumul to zero.
-            time)
-
-        time_dimension = routing.GetDimensionOrDie(time)
-        # Add time window constraints for each location except depot.
-        for location_idx, time_window in enumerate(time_windows):
-            if location_idx == 0:
+            10000,  # Allow waiting time
+            10000,  # Max time per vehicle
+            False,  # Don't force start cumul to zero
+            'Time'
+        )
+
+        time_dimension = routing.GetDimensionOrDie('Time')
+
+        # Time windows for all locations except depot
+        for location_idx, (start, end) in enumerate(time_windows):
+            if location_idx == depot:
                 continue
             index = manager.NodeToIndex(location_idx)
-            a, b = int(time_window[0]), int(time_window[1])
-            time_dimension.CumulVar(index).SetRange(a, b)
+            time_dimension.CumulVar(index).SetRange(int(start), int(end))
 
-        # Add time window constraints for each vehicle start node.
+        # Time windows for vehicle start (depot)
+        depot_start, depot_end = time_windows[depot]
         for vehicle_id in range(num_vehicles):
             index = routing.Start(vehicle_id)
-            a, b = int(time_windows[0][0]), int(time_windows[0][1])
-            time_dimension.CumulVar(index).SetRange(a, b)
+            time_dimension.CumulVar(index).SetRange(int(depot_start), int(depot_end))
 
+        # Finalizer hints for optimization
         for i in range(num_vehicles):
-            routing.AddVariableMinimizedByFinalizer(
-                time_dimension.CumulVar(routing.Start(i)))
-            routing.AddVariableMinimizedByFinalizer(
-                time_dimension.CumulVar(routing.End(i)))
+            routing.AddVariableMinimizedByFinalizer(time_dimension.CumulVar(routing.Start(i)))
+            routing.AddVariableMinimizedByFinalizer(time_dimension.CumulVar(routing.End(i)))
 
-        search_parameters = pywrapcp.DefaultRoutingSearchParameters()
-        search_parameters.first_solution_strategy = (routing_enums_pb2.FirstSolutionStrategy.AUTOMATIC)
-        assignment = routing.SolveWithParameters(search_parameters)
+        search_params = pywrapcp.DefaultRoutingSearchParameters()
+        search_params.first_solution_strategy = routing_enums_pb2.FirstSolutionStrategy.AUTOMATIC
 
+        assignment = routing.SolveWithParameters(search_params)
         return routing, assignment
 
 
+def evaluate_google_model(validation_dataset):
+    """
+    Evaluate the validation dataset using Google OR-Tools model.
 
+    Args:
+        validation_dataset (Dataset): A dataset with a get_data method.
 
-def evaluate_google_model(validation_dataset):
+    Returns:
+        torch.Tensor: Scores for each batch item.
+    """
     validation_dataset.device = 'cpu'
     data = validation_dataset.get_data()
     model = GoogleActor(scale_factor=100)
-    scores = model(data)
-    return scores
-
+    return model(data)