Added an option to scale pickup and delivery request durations to the…

… transform script.

Bonus change:
- modified multiple flag parsers to check that the numbers they get are non-negative.

python/cfr/json/transform_request.py

@@ -68,6 +68,15 @@ def _parse_comma_separated_list(value: str) -> Sequence[str]:
   return value.split(",")
 
 
+def _non_negative_float(value: str) -> float:
+  value_as_float = float(value)
+  if value_as_float < 0:
+    raise argparse.ArgumentTypeError(
+        f"expected a non-negative value, got {value}"
+    )
+  return value_as_float
+
+
 @dataclasses.dataclass(slots=True)
 class Flags:
   """Holds values of command-line flags of this script.
@@ -84,6 +93,7 @@ class Flags:
   shipment_penalty_cost_per_item: float | None
   remove_pickups: bool
   soften_allowed_vehicle_indices_cost: float | None
+  visit_duration_scaling_factor: float | None
 
   duplicate_vehicles_by_label: Sequence[str] | None
   remove_vehicles_by_label: Sequence[str] | None
@@ -135,7 +145,7 @@ def from_command_line(cls, args: Sequence[str] | None) -> Self:
     )
     parser.add_argument(
         "--shipment_penalty_cost_per_item",
-        type=float,
+        type=_non_negative_float,
         required=False,
         help=(
             "When provided, turns all mandatory shipments in the request into"
@@ -161,7 +171,7 @@ def from_command_line(cls, args: Sequence[str] | None) -> Self:
     )
     parser.add_argument(
         "--soften_allowed_vehicle_indices_cost",
-        type=float,
+        type=_non_negative_float,
         required=False,
         help=(
             "When set, replaces all allowedVehicleIndices in the request with a"
@@ -171,6 +181,16 @@ def from_command_line(cls, args: Sequence[str] | None) -> Self:
             " this flag."
         ),
     )
+    parser.add_argument(
+        "--visit_duration_scaling_factor",
+        type=_non_negative_float,
+        required=False,
+        help=(
+            "When set, all pickup and delivery request durations in the request"
+            " are scaled by the value of this flag. Must be a non-negative"
+            " floating point number."
+        ),
+    )
     parser.add_argument(
         "--duplicate_vehicles_by_label",
         type=_parse_comma_separated_list,
@@ -276,6 +296,10 @@ def main(args: Sequence[str] | None = None) -> None:
     transforms.soften_allowed_vehicle_indices(
         model, cost=flags.soften_allowed_vehicle_indices_cost
     )
+  if flags.visit_duration_scaling_factor is not None:
+    transforms.scale_visit_request_durations(
+        model, factor=flags.visit_duration_scaling_factor
+    )
   if removed_labels := flags.remove_vehicles_by_label:
     _remove_vehicles_by_label(model, removed_labels)
   if duplicated_labels := flags.duplicate_vehicles_by_label:

python/cfr/json/transform_request_test.py

@@ -244,6 +244,54 @@ def test_soften_allowed_vehicle_indices(self):
         expected_output_request,
     )
 
+  def test_visit_duration_scaling_factor_nonzero(self):
+    request: cfr_json.OptimizeToursRequest = {
+        "model": {
+            "shipments": [
+                {"label": "S001", "pickups": [{"duration": "10s"}]},
+                {"label": "S002", "deliveries": [{"duration": "60s"}]},
+            ]
+        }
+    }
+    expected_output_request: cfr_json.OptimizeToursRequest = {
+        "model": {
+            "shipments": [
+                {"label": "S001", "pickups": [{"duration": "8s"}]},
+                {"label": "S002", "deliveries": [{"duration": "48s"}]},
+            ]
+        }
+    }
+    self.assertEqual(
+        self.run_transform_request_main(
+            request, ("--visit_duration_scaling_factor=0.8",)
+        ),
+        expected_output_request,
+    )
+
+  def test_visit_duration_scaling_factor_zero(self):
+    request: cfr_json.OptimizeToursRequest = {
+        "model": {
+            "shipments": [
+                {"label": "S001", "pickups": [{"duration": "10s"}]},
+                {"label": "S002", "deliveries": [{"duration": "60s"}]},
+            ]
+        }
+    }
+    expected_output_request: cfr_json.OptimizeToursRequest = {
+        "model": {
+            "shipments": [
+                {"label": "S001", "pickups": [{"duration": "0s"}]},
+                {"label": "S002", "deliveries": [{"duration": "0s"}]},
+            ]
+        }
+    }
+    self.assertEqual(
+        self.run_transform_request_main(
+            request, ("--visit_duration_scaling_factor=0",)
+        ),
+        expected_output_request,
+    )
+
   def test_duplicate_vehicles_by_label(self):
     request: cfr_json.OptimizeToursRequest = {
         "model": {

python/cfr/json/transforms.py

@@ -8,6 +8,7 @@
 import collections
 from collections.abc import Callable, Collection, Iterable
 import copy
+import itertools
 import math
 import re
 
@@ -269,6 +270,26 @@ def remove_load_limits(model: cfr_json.ShipmentModel) -> None:
     vehicle.pop("loadLimits", None)
 
 
+def scale_visit_request_durations(
+    model: cfr_json.ShipmentModel, factor: float
+) -> None:
+  """Scales visit durations in the model by the given factor.
+
+  Args:
+    model: The model in which the visit durations are scaled.
+    factor: The scaling factor. Must be non-negative.
+  """
+  if factor < 0:
+    raise ValueError("factor must be a non-negative number")
+  for shipment in cfr_json.get_shipments(model):
+    pickups = shipment.get("pickups", ())
+    deliveries = shipment.get("deliveries", ())
+    for visit_request in itertools.chain(pickups, deliveries):
+      duration = cfr_json.get_visit_request_duration(visit_request)
+      duration *= factor
+      visit_request["duration"] = cfr_json.as_duration_string(duration)
+
+
 def remove_pickups(model: cfr_json.ShipmentModel) -> None:
   """Removes pickups from shipments that have both pickups and deliveries.
 

python/cfr/json/transforms_test.py

@@ -531,6 +531,58 @@ def test_remove_load_limits(self):
     )
 
 
+class ScaleVisitRequestDurations(unittest.TestCase):
+  """Tests for scale_visit_request_duration."""
+
+  maxDiff = None
+
+  _MODEL: cfr_json.ShipmentModel = {
+      "shipments": [
+          {
+              "pickups": [{"duration": "100s"}, {"duration": "0s"}],
+              "deliveries": [{"duration": "60s"}],
+          },
+          {"pickups": [{"duration": "30s"}]},
+          {"deliveries": [{"duration": "120s"}]},
+      ]
+  }
+
+  def test_negative_factor(self):
+    model: cfr_json.ShipmentModel = {}
+    with self.assertRaisesRegex(ValueError, "non-negative"):
+      transforms.scale_visit_request_durations(model, -0.5)
+
+  def test_zero_factor(self):
+    model: cfr_json.ShipmentModel = copy.deepcopy(self._MODEL)
+    expected_model: cfr_json.ShipmentModel = {
+        "shipments": [
+            {
+                "pickups": [{"duration": "0s"}, {"duration": "0s"}],
+                "deliveries": [{"duration": "0s"}],
+            },
+            {"pickups": [{"duration": "0s"}]},
+            {"deliveries": [{"duration": "0s"}]},
+        ]
+    }
+    transforms.scale_visit_request_durations(model, 0)
+    self.assertEqual(model, expected_model)
+
+  def test_non_zero_factor(self):
+    model: cfr_json.ShipmentModel = copy.deepcopy(self._MODEL)
+    expected_model: cfr_json.ShipmentModel = {
+        "shipments": [
+            {
+                "pickups": [{"duration": "110s"}, {"duration": "0s"}],
+                "deliveries": [{"duration": "66s"}],
+            },
+            {"pickups": [{"duration": "33s"}]},
+            {"deliveries": [{"duration": "132s"}]},
+        ]
+    }
+    transforms.scale_visit_request_durations(model, 1.1)
+    self.assertEqual(model, expected_model)
+
+
 class RemovePickupsTest(unittest.TestCase):
   """Tests for remove_pickups."""