[transform_request] Added an option to reduce the request to a given …

…set of vehicles.

Bonus change:
- Also update shipment indices in the injected first solution when removing
  shipments from the request.

python/cfr/json/transform_request.py

@@ -17,7 +17,7 @@
 """
 
 import argparse
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence, Set
 import dataclasses
 import enum
 import logging
@@ -110,6 +110,7 @@ class Flags:
 
   duplicate_vehicles_by_label: Sequence[str] | None
   remove_vehicles_by_label: Sequence[str] | None
+  reduce_to_vehicles_by_label: Sequence[str] | None
   infeasible_shipment_after_removing_vehicle: transforms.OnInfeasibleShipment
 
   transform_breaks: str | None
@@ -228,6 +229,17 @@ def from_command_line(cls, args: Sequence[str] | None) -> Self:
             " same label, and it appears in this list, all of them are removed."
         ),
     )
+    parser.add_argument(
+        "--reduce_to_vehicles_by_label",
+        type=_parse_comma_separated_list,
+        help=(
+            "A comma-separated list of vehicle labels. Removes all vehicles"
+            " whose labels do not appear in this list. When multiple vehicles"
+            " have the same label, and it appears in the list, all of them are"
+            " preserved. Removes all shipments that become trivially infeasible"
+            " when the vehicles are removed."
+        ),
+    )
     parser.add_argument(
         "--transform_breaks",
         required=False,
@@ -260,6 +272,36 @@ def from_command_line(cls, args: Sequence[str] | None) -> Self:
     return cls(**vars(parsed_args))
 
 
+def _get_indices_of_vehicles_with_labels(
+    model: cfr_json.ShipmentModel,
+    vehicle_labels: Iterable[str],
+) -> tuple[Sequence[int], Set[str]]:
+  """Translates vehicle labels into vehicle indices.
+
+  When multiple vehicles use the same label, the indices of all of them will be
+  in the returned sequence.
+
+  Args:
+    model: The model in which the vehicle labels are translated.
+    vehicle_labels: The vehicle labels to translate.
+
+  Returns:
+    A tuple (vehicle_indices, unseen_labels) where vehicle_indices is a sequence
+    of vehicle indices selected by the labels, and unseen_labels is a set of
+    labels from `vehicle_labels` that did not match any vehicle in the model.
+  """
+  selected_labels = set(vehicle_labels)
+  unseen_labels = set(selected_labels)
+  selected_indices = []
+  for vehicle_index, vehicle in enumerate(cfr_json.get_vehicles(model)):
+    vehicle_label = vehicle.get("label", "")
+    if vehicle_label and vehicle_label in selected_labels:
+      selected_indices.append(vehicle_index)
+      unseen_labels.remove(vehicle_label)
+
+  return selected_indices, unseen_labels
+
+
 def _remove_vehicles_by_label(
     request: cfr_json.OptimizeToursRequest,
     vehicle_labels: Iterable[str],
@@ -274,28 +316,60 @@ def _remove_vehicles_by_label(
       trivially infeasible after removing a vehicle.
   """
   model = request["model"]
-  indices_to_remove = set()
-  labels_to_remove = set(vehicle_labels)
-  removed_labels = set()
-  for vehicle_index, vehicle in enumerate(cfr_json.get_vehicles(model)):
-    vehicle_label = vehicle.get("label", "")
-    if not vehicle_label or vehicle_label not in labels_to_remove:
-      continue
-    indices_to_remove.add(vehicle_index)
-    removed_labels.add(vehicle_label)
+  indices_to_remove, unseen_labels = _get_indices_of_vehicles_with_labels(
+      model, vehicle_labels
+  )
 
-  if len(removed_labels) != len(labels_to_remove):
-    missing_labels = sorted(labels_to_remove - removed_labels)
+  if unseen_labels:
     raise ValueError(
         "Vehicle labels from --remove_vehicles_by_label do not appear in the"
-        f" model: {', '.join(repr(label) for label in missing_labels)}"
+        f" model: {', '.join(repr(label) for label in sorted(unseen_labels))}"
     )
 
-  new_vehicle_for_old_vehicle = transforms.remove_vehicles(
-      request["model"], indices_to_remove, on_infeasible_shipment
+  new_vehicle_for_old_vehicle, new_shipment_for_old_shipment = (
+      transforms.remove_vehicles(
+          model, set(indices_to_remove), on_infeasible_shipment
+      )
+  )
+  transforms.remove_vehicles_from_injected_first_solution_routes(
+      request, new_vehicle_for_old_vehicle, new_shipment_for_old_shipment
+  )
+
+
+def _reduce_to_vehicles_by_label(
+    request: cfr_json.OptimizeToursRequest,
+    vehicle_labels: Iterable[str],
+) -> None:
+  """Removes all vehicles in the request whose label is not in `vehicle_labels`.
+
+  Removes also all shipments that become trivially infeasible after removing the
+  vehicles.
+
+  Args:
+    request: The request in which the vehicles are removed.
+    vehicle_labels: The labels of vehicles to be kept in the model.
+  """
+  model = request["model"]
+  indices_to_keep, unseen_labels = _get_indices_of_vehicles_with_labels(
+      model, vehicle_labels
+  )
+
+  if unseen_labels:
+    raise ValueError(
+        "Vehicle labels from --reduce_to_vehicles_by_label do not appear in the"
+        f" model: {', '.join(repr(label) for label in sorted(unseen_labels))}"
+    )
+  num_vehicles = len(cfr_json.get_vehicles(model))
+  indices_to_remove = set(range(num_vehicles))
+  indices_to_remove.difference_update(indices_to_keep)
+
+  new_vehicle_for_old_vehicle, new_shipment_for_old_shipment = (
+      transforms.remove_vehicles(
+          model, indices_to_remove, transforms.OnInfeasibleShipment.REMOVE
+      )
   )
   transforms.remove_vehicles_from_injected_first_solution_routes(
-      request, new_vehicle_for_old_vehicle
+      request, new_vehicle_for_old_vehicle, new_shipment_for_old_shipment
   )
 
 
@@ -363,6 +437,8 @@ def main(args: Sequence[str] | None = None) -> None:
         removed_labels,
         flags.infeasible_shipment_after_removing_vehicle,
     )
+  if preserved_labels := flags.reduce_to_vehicles_by_label:
+    _reduce_to_vehicles_by_label(request, preserved_labels)
   if duplicated_labels := flags.duplicate_vehicles_by_label:
     _duplicate_vehicles_by_label(model, duplicated_labels)
   if flags.transform_breaks is not None:

python/cfr/json/transform_request_test.py

@@ -579,6 +579,61 @@ def test_transform_breaks(self):
         expected_output_request,
     )
 
+  def test_reduce_to_vehicles_by_label(self):
+    request: cfr_json.OptimizeToursRequest = {
+        "model": {
+            "shipments": [
+                {"label": "S001", "allowedVehicleIndices": [0, 2]},
+                {
+                    "label": "S002",
+                    "costsPerVehicle": [100],
+                    "costsPerVehicleIndices": [1],
+                },
+            ],
+            "vehicles": [
+                {"label": "V001", "costPerHour": 30},
+                {"label": "V002", "costPerHour": 60},
+                {"label": "V003", "costPerHour": 90},
+            ],
+        },
+        "injectedFirstSolutionRoutes": [
+            {"vehicleLabel": "V001", "visits": []},
+            {
+                "vehicleIndex": 1,
+                "vehicleLabel": "V002",
+                "visits": [{"shipmentIndex": 1, "visitRequestIndex": 0}],
+            },
+        ],
+    }
+    expected_output_request: cfr_json.OptimizeToursRequest = {
+        "model": {
+            "shipments": [
+                {
+                    "label": "S002",
+                    "costsPerVehicle": [100],
+                    "costsPerVehicleIndices": [0],
+                },
+            ],
+            "vehicles": [
+                {"label": "V002", "costPerHour": 60},
+            ],
+        },
+        "injectedFirstSolutionRoutes": [
+            {
+                "vehicleIndex": 0,
+                "vehicleLabel": "V002",
+                "visits": [{"shipmentIndex": 0, "visitRequestIndex": 0}],
+            },
+        ],
+    }
+    self.assertEqual(
+        self.run_transform_request_main(
+            request, ("--reduce_to_vehicles_by_label=V002",)
+        ),
+        expected_output_request,
+        f"{request=}\n{expected_output_request=}",
+    )
+
   def test_override_consider_road_traffic_true_from_false(self):
     request: cfr_json.OptimizeToursRequest = {
         "model": {},

python/cfr/json/transforms.py

@@ -6,7 +6,7 @@
 """A library of transformations to CFR JSON requests."""
 
 import collections
-from collections.abc import Callable, Collection, Iterable, Mapping
+from collections.abc import Callable, Collection, Iterable, Mapping, MutableMapping
 import copy
 import itertools
 import logging
@@ -142,7 +142,7 @@ def remove_vehicles(
     model: cfr_json.ShipmentModel,
     vehicle_indices: Collection[int],
     on_infeasible_shipment: OnInfeasibleShipment = OnInfeasibleShipment.FAIL,
-) -> Mapping[int, int]:
+) -> tuple[Mapping[int, int], Mapping[int, int]]:
   """Removes vehicles with the given indices from the model.
 
   Removes the vehicles from the list and updates vehicle indices in the other
@@ -225,6 +225,7 @@ def remove_vehicles(
           if vehicle_index not in removed_vehicle_indices
       ]
 
+  new_shipment_for_old_shipment: MutableMapping[int, int] = {}
   if removed_shipments:
     # Remove all trivially infeasible shipments from the model.
     new_shipments = []
@@ -237,15 +238,17 @@ def remove_vehicles(
             repr(label) if label is not None else "",
         )
       else:
+        new_shipment_for_old_shipment[shipment_index] = len(new_shipments)
         new_shipments.append(shipment)
     model["shipments"] = new_shipments
 
-  return new_vehicle_for_old_vehicle
+  return new_vehicle_for_old_vehicle, new_shipment_for_old_shipment
 
 
 def remove_vehicles_from_injected_first_solution_routes(
     request: cfr_json.OptimizeToursRequest,
     new_vehicle_for_old_vehicle: Mapping[int, int],
+    new_shipment_for_old_shipment: Mapping[int, int],
 ) -> None:
   """Removes given vehicles from the first solution hint in `request`."""
   injected_first_solution_routes = request.get("injectedFirstSolutionRoutes")
@@ -259,6 +262,13 @@ def remove_vehicles_from_injected_first_solution_routes(
     if new_vehicle_index is None:
       continue
     route["vehicleIndex"] = new_vehicle_index
+    if new_shipment_for_old_shipment:
+      for visit in cfr_json.get_visits(route):
+        old_shipment_index = visit.get("shipmentIndex", 0)
+        visit["shipmentIndex"] = new_shipment_for_old_shipment[
+            old_shipment_index
+        ]
+
     new_injected_first_solution_routes.append(route)
 
   if new_injected_first_solution_routes:

python/cfr/json/transforms_test.py

@@ -385,7 +385,7 @@ def test_no_injected_first_solution(self):
     }
     expected_request = copy.deepcopy(request)
     transforms.remove_vehicles_from_injected_first_solution_routes(
-        request, {0: 0, 2: 1, 3: 2}
+        request, {0: 0, 2: 1, 3: 2}, {}
     )
     self.assertEqual(request, expected_request)
 
@@ -400,8 +400,16 @@ def test_remove_some_vehicles(self):
         },
         "injectedFirstSolutionRoutes": [
             {"vehicleLabel": "V001"},
-            {"vehicleIndex": 1, "vehicleLabel": "V004"},
-            {"vehicleIndex": 5, "vehicleLabel": "V007"},
+            {
+                "vehicleIndex": 1,
+                "vehicleLabel": "V004",
+                "visits": [{"shipmentIndex": 3, "visitRequestIndex": 0}],
+            },
+            {
+                "vehicleIndex": 5,
+                "vehicleLabel": "V007",
+                "visits": [{"visitRequestIndex": 0}],
+            },
         ],
     }
     expected_request: cfr_json.OptimizeToursRequest = {
@@ -413,12 +421,20 @@ def test_remove_some_vehicles(self):
             ]
         },
         "injectedFirstSolutionRoutes": [
-            {"vehicleIndex": 1, "vehicleLabel": "V004"},
-            {"vehicleIndex": 2, "vehicleLabel": "V007"},
+            {
+                "vehicleIndex": 1,
+                "vehicleLabel": "V004",
+                "visits": [{"shipmentIndex": 2, "visitRequestIndex": 0}],
+            },
+            {
+                "vehicleIndex": 2,
+                "vehicleLabel": "V007",
+                "visits": [{"shipmentIndex": 0, "visitRequestIndex": 0}],
+            },
         ],
     }
     transforms.remove_vehicles_from_injected_first_solution_routes(
-        request, {1: 1, 5: 2}
+        request, {1: 1, 5: 2}, {0: 0, 1: 1, 3: 2}
     )
     self.assertEqual(request, expected_request)