[two_step_routing] Add initial support for pickup and delivery.

The input scenario can now contain both delivery-only and pickup-only
shipments, both handled directoy from the vehicle or from a parking
location. The direct shipments translate naturally in the solution.

For shipments delivered from a parking location:
- in the local model they are represented as pickup and delivery
  shipments with a pickup at the parking location and delivery at the
  customer address, or pickup from the customer address and delivery at
  the parking location.
- in the global model, the visits to the parking are represented as
  delivery-only shipments where the load requirements are the maximum
  over the requirements of the shipments delivered from the parking and
  picked up from the parking.
- in the merged model, the shipments are pickup-only or delivery-only
  again.

To use the pickup & delivery requests, just add `pickups` or
`deliveries` as in a vanilla CFR request.

While the internal representations and models have changed
significantly, the public API of two_step_routing.py did not change.

python/cfr/analysis/analysis.py

@@ -377,7 +377,13 @@ def group_global_visits(
     vehicle_index: int,
     split_by_breaks: bool = False,
 ) -> Iterable[
-    tuple[two_step_routing.ParkingTag, int, Sequence[cfr_json.Shipment]]
+    tuple[
+        two_step_routing.ParkingTag,
+        int,
+        Sequence[cfr_json.Shipment],
+        int,
+        int,
+    ]
 ]:
   """Iterates over groups of "global" visits and their shipments on a route.
 
@@ -398,8 +404,9 @@ def group_global_visits(
       three visits to the parking.
 
   Yields:
-    A sequence of triples `(parking_tag, num_rounds, shipments)` where each
-    triple represents one group of "global" visits.
+    A sequence of five tuple `(parking_tag, num_rounds, shipments,
+    arrival_visit_index, departure_visit_index)` where each five tuple
+    represents one group of "global" visits.
 
     When the global visit is a shipment that is delivered directly,
     `parking_tag` is `None`, `num_rounds` is 1, and `shipments` is a list that
@@ -410,6 +417,11 @@ def group_global_visits(
     of delivery rounds in the sequence (the number of visits in the global model
     used by the two_step_routing library), and `shipments` is the list of
     shipments delivered in all delivery rounds in the group.
+
+    `arrival_visit_index` and `departure_visit_index` are both the index of
+    the visit in the route.  `arrival_visit_index` is the index of the visit
+    to the "arrival to parking" virtual shipment, and `departure_visit_index`
+    is the index of the visit to the "departure from parking" virtual shipment.
   """
   parking_data = scenario.parking_location_data
   route = scenario.routes[vehicle_index]
@@ -425,10 +437,20 @@ def group_global_visits(
     parking_tag, arrival_visit_index, departure_visit_index = global_visits[
         global_visit_index
     ]
+
+    first_arrival_visit_index = arrival_visit_index
+    last_departure_visit_index = departure_visit_index
+
     if parking_tag is None:
       # Shipment delivered directly.
       shipment_index = visits[arrival_visit_index].get("shipmentIndex", 0)
-      yield (None, 1, (shipments[shipment_index],))
+      yield (
+          None,
+          1,
+          (shipments[shipment_index],),
+          first_arrival_visit_index,
+          last_departure_visit_index,
+      )
       global_visit_index += 1
       continue
 
@@ -459,12 +481,20 @@ def group_global_visits(
               departure_visit_index - 1,
           )
       )
+      last_departure_visit_index = departure_visit_index
+
       num_rounds += 1
       # After the last iteration of this loop, global_visit_index will be the
       # index of the last global visit in this (potential) ping-pong.
       global_visit_index += 1
 
-    yield (parking_tag, num_rounds, group_shipments)
+    yield (
+        parking_tag,
+        num_rounds,
+        group_shipments,
+        first_arrival_visit_index,
+        last_departure_visit_index
+    )
 
     global_visit_index += 1
 
@@ -507,7 +537,7 @@ def get_num_ping_pongs(
   """
   num_ping_pongs = 0
   bad_ping_pong_parking_tags = []
-  for parking_tag, num_rounds, group_shipments in group_global_visits(
+  for parking_tag, num_rounds, group_shipments, _, _ in group_global_visits(
       scenario, vehicle_index, split_by_breaks=split_by_breaks
   ):
     if num_rounds == 1:
@@ -672,7 +702,7 @@ def get_num_sandwiches(
   bad_sandwich_tags = []
   last_visit_to_parking = {}
 
-  for parking_tag, _, group_shipments in group_global_visits(
+  for parking_tag, _, group_shipments, _, _ in group_global_visits(
       scenario, vehicle_index
   ):
     last_visit_shipments = last_visit_to_parking.get(parking_tag)

python/cfr/analysis/analysis_test.py

@@ -57,7 +57,16 @@ def test_grouped_shipments_no_breaks(self):
         )
     )
     for group, expected_group in zip(groups, expected_groups, strict=True):
-      tag, num_rounds, shipments = group
+      (tag, num_rounds, shipments, arrival_visit_index,
+       departure_visit_index) = group
+
+      visits = cfr_json.get_visits(self._scenario.routes[0])
+      if tag is not None:
+        self.assertEqual(
+            visits[arrival_visit_index]["shipmentLabel"], f"{tag} arrival")
+        self.assertEqual(
+            visits[departure_visit_index]["shipmentLabel"], f"{tag} departure")
+
       expected_tag, expected_num_rounds, expected_num_shipments = expected_group
       with self.subTest(f"group {expected_group!r}"):
         self.assertEqual(tag, expected_tag)
@@ -92,7 +101,19 @@ def test_grouped_shipments_with_breaks(self):
         )
     )
     for group, expected_group in zip(groups, expected_groups, strict=True):
-      tag, num_rounds, shipments = group
+      tag, num_rounds, shipments, arrival_visit_index, departure_visit_index = (
+          group
+      )
+
+      visits = cfr_json.get_visits(self._scenario.routes[0])
+      if tag is not None:
+        self.assertEqual(
+            visits[arrival_visit_index]["shipmentLabel"], f"{tag} arrival"
+        )
+        self.assertEqual(
+            visits[departure_visit_index]["shipmentLabel"], f"{tag} departure"
+        )
+
       expected_tag, expected_num_rounds, expected_num_shipments = expected_group
       with self.subTest(f"group {expected_group!r}"):
         self.assertEqual(tag, expected_tag)

python/cfr/analysis/cfr-json-analysis.ipynb

@@ -40,6 +40,7 @@
       "source": [
         "### Prerequisities\n",
         "\n",
+        "\u003c!-- _and_replace\n",
         "If you're not familiar with Colab notebooks, check out the\n",
         "[Welcome to Colaboratory](https://colab.research.google.com/notebooks/intro.ipynb)\n",
         "notebook first for a tutorial.\n",
@@ -59,6 +60,7 @@
         "    [local runtime guide](https://research.google.com/colaboratory/local-runtimes.html).\n",
         "    As of 2023-10-23, using the `COLAB_KERNEL_MANAGER_PROXY_PORT` option is\n",
         "    needed to make file upload work correctly.\n",
+        "-->\n",
         "\n",
         "### How to use the colab\n",
         "\n",
@@ -78,8 +80,10 @@
         "\n",
         "### Don't panic!\n",
         "\n",
+        "\u003c!-- _and_replace\n",
         "If you have any questions or run into issues with using the colab, contact\n",
-        "ondrasej at google dot com.\n"
+        "ondrasej at google dot com.\n",
+        "-->"
       ]
     },
     {
@@ -122,13 +126,15 @@
         "import ipywidgets\n",
         "import pandas as pd\n",
         "\n",
-        "# Clone the CFR library from GitHub, and import from there.\n",
-        "!git clone https://github.com/google/cfr\n",
-        "from cfr.python.cfr import utils\n",
-        "from cfr.python.cfr.analysis import analysis\n",
-        "from cfr.python.cfr.json import cfr_json\n",
-        "from cfr.python.cfr.json import human_readable\n",
-        "from cfr.python.cfr.two_step_routing import two_step_routing\n"
+        "# _and_replace_begin\n",
+        "# # Clone the CFR library from GitHub, and import from there.\n",
+        "# !git clone https://github.com/google/cfr\n",
+        "# from cfr.python.cfr import utils\n",
+        "# from cfr.python.cfr.analysis import analysis\n",
+        "# from cfr.python.cfr.json import cfr_json\n",
+        "# from cfr.python.cfr.json import human_readable\n",
+        "# from cfr.python.cfr.two_step_routing import two_step_routing\n",
+        "# "
       ]
     },
     {

python/cfr/json/cfr_json.py

@@ -9,6 +9,7 @@
 from collections.abc import Collection, Iterable, Mapping, Sequence, Set
 import datetime
 import itertools
+import logging
 from typing import TypeAlias, TypedDict
 
 # A duration in a string format following the protocol buffers specification in
@@ -405,9 +406,11 @@ def combined_allowed_vehicle_indices(
   """Returns the list of allowed vehicle indices that can serve all shipments."""
   allowed_vehicles = None
   for shipment in shipments:
+    logging.debug("shipment: %r", shipment)
     shipment_allowed_vehicles = shipment.get("allowedVehicleIndices")
     if shipment_allowed_vehicles is None:
       continue
+    logging.debug("allowed_vehicles = %r", allowed_vehicles)
     if allowed_vehicles is None:
       allowed_vehicles = set(shipment_allowed_vehicles)
     else:
@@ -420,15 +423,46 @@ def combined_allowed_vehicle_indices(
 
 
 def combined_load_demands(shipments: Collection[Shipment]) -> dict[str, Load]:
-  """Computes the combined load demands of all shipments in `shipments`."""
-  demands = collections.defaultdict(int)
+  """Computes the combined load demands of all shipments in `shipments`.
+
+  Assumes that shipments are shipments that are picked up and delivered at the
+  same location (the same parking location), and computes load demands that are
+  sufficient both for the pickups and for the deliveries (assuming that pickups
+  happen before deliveries).
+
+  For each unit, computes the amount that is delivered and the amount that is
+  picked up, and takes the max of the two. The output load demands are the
+  smallest load demands that need to be reserved on the vehicle for the whole
+  route to guarantee that all shipments can be performed.
+
+  Args:
+    shipments: The list of shipments to be processed.
+
+  Returns:
+    The combined load demands as described above. Never contains negative
+    numbers.
+  """
+  delivery_demands = collections.defaultdict(int)
+  pickup_demands = collections.defaultdict(int)
   for shipment in shipments:
     shipment_demands = shipment.get("loadDemands")
     if shipment_demands is None:
       continue
+
+    is_pickup = get_pickup_or_none(shipment) is not None
+    combined_demands = pickup_demands if is_pickup else delivery_demands
     for unit, amount in shipment_demands.items():
-      demands[unit] += int(amount.get("amount", 0))
-  return {unit: {"amount": str(amount)} for unit, amount in demands.items()}
+      combined_demands[unit] += int(amount.get("amount", 0))
+
+  demands = {}
+  for unit in set(itertools.chain(delivery_demands, pickup_demands)):
+    demands[unit] = {
+        "amount": str(
+            max(delivery_demands.get(unit, 0), pickup_demands.get(unit, 0))
+        )
+    }
+
+  return demands
 
 
 _DEFAULT_GLOBAL_START_TIME = datetime.datetime.fromtimestamp(
@@ -1020,6 +1054,7 @@ def recompute_route_metrics(
     )
 
   if check_consistency:
+    assert len(visits) + 1 == len(get_transitions(route))
     if (
         route_total_duration
         != route_travel_duration

python/cfr/json/cfr_json_test.py

@@ -369,7 +369,7 @@ class CombinedLoadDemandsTest(unittest.TestCase):
   def test_no_shipments(self):
     self.assertEqual(cfr_json.combined_load_demands(()), {})
 
-  def test_some_shipments(self):
+  def test_some_deliveries(self):
     shipments = [
         cfr_json.make_shipment(
             "S001",
@@ -398,6 +398,64 @@ def test_some_shipments(self):
         },
     )
 
+  def test_some_pickups(self):
+    shipments: Sequence[cfr_json.Shipment] = (
+        cfr_json.make_shipment(
+            "S001",
+            pickup_latlng=(48.86471, 2.34901),
+            pickup_duration="120s",
+            load_demands={"wheat": 3, "wood": 1},
+        ),
+        cfr_json.make_shipment(
+            "S002",
+            pickup_latlng=(48.86471, 2.34901),
+            pickup_duration="120s",
+            load_demands={"wood": 5, "ore": 2},
+        ),
+        cfr_json.make_shipment(
+            "S002",
+            pickup_latlng=(48.86471, 2.34901),
+            pickup_duration="120s",
+        ),
+    )
+    self.assertEqual(
+        cfr_json.combined_load_demands(shipments),
+        {
+            "wheat": {"amount": "3"},
+            "wood": {"amount": "6"},
+            "ore": {"amount": "2"},
+        },
+    )
+
+  def test_pickups_and_deliveries(self):
+    shipments: Sequence[cfr_json.Shipment] = (
+        cfr_json.make_shipment(
+            "S001",
+            pickup_latlng=(48.86471, 2.34901),
+            pickup_duration="120s",
+            load_demands={"wheat": 3, "wood": 1},
+        ),
+        cfr_json.make_shipment(
+            "S002",
+            delivery_latlng=(48.86471, 2.34901),
+            delivery_duration="120s",
+            load_demands={"wood": 5, "ore": 2},
+        ),
+        cfr_json.make_shipment(
+            "S002",
+            pickup_latlng=(48.86471, 2.34901),
+            pickup_duration="120s",
+        ),
+    )
+    self.assertEqual(
+        cfr_json.combined_load_demands(shipments),
+        {
+            "wheat": {"amount": "3"},
+            "wood": {"amount": "5"},
+            "ore": {"amount": "2"},
+        },
+    )
+
 
 class GetShipmentsTest(unittest.TestCase):
   """Tests for get_shipments."""