support subway drivers generation

mosstool/map/public_transport/public_transport_post.py

@@ -188,7 +188,7 @@ def route_length(
                             "departure_times": departure_times,
                             "offset_times": offset_times,
                         },
-                        "capacity":STATION_CAPACITY[pub["type"]],
+                        "capacity": STATION_CAPACITY[pub["type"]],
                         "taz_costs": [],
                     }
                 )
@@ -308,6 +308,7 @@ def get_taz_id(x, y):
 
 
 def _post_compute(m: dict, workers: int):
+    m = deepcopy(m)
     header = m["header"]
     aois = {a["id"]: a for a in m["aois"]}
     lanes = {l["id"]: l for l in m["lanes"]}
@@ -338,7 +339,24 @@ def _station_distance(road_ids, s_start: float, s_end: float) -> float:
         res += s_end
         return res
 
+    # add subline_driving_lane_pairs
+    for _, aoi in aois.items():
+        aoi["subline_driving_lane_pairs"] = []
+
+    def aoi2driving_lane_id(aoi: dict, road_id: int):
+        road = roads[road_id]
+        rightest_lane_id = [
+            lanes[lid]
+            for lid in road["lane_ids"]
+            if lanes[lid]["type"] == mapv2.LANE_TYPE_DRIVING
+        ][-1]["id"]
+        for gate, pos in zip(aoi["driving_gates"], aoi["driving_positions"]):
+            if pos["lane_id"] == rightest_lane_id:
+                return pos["lane_id"]
+        raise ValueError(f"{rightest_lane_id} not in AOI {aoi['id']} gates")
+
     for subline in sublines_data:
+        subline_id = subline["id"]
         station_aois = {aoi_id: aois[aoi_id] for aoi_id in subline["aoi_ids"]}
         sta_aoi_ids = list(station_aois.keys())
         route_lengths = []
@@ -352,13 +370,26 @@ def _station_distance(road_ids, s_start: float, s_end: float) -> float:
                 if lane["parent_id"] == road_ids[0]:
                     s_start = d["s"]
                     break
+            aoi_start["subline_driving_lane_pairs"].append(
+                {
+                    "subline_id": subline_id,
+                    "driving_lane_id": aoi2driving_lane_id(aoi_start, road_ids[0]),
+                }
+            )
             aoi_end = station_aois[sta_aoi_ids[i + 1]]
             s_end = 0
             for d in aoi_end["driving_positions"]:
                 lane = lanes[d["lane_id"]]
                 if lane["parent_id"] == road_ids[-1]:
                     s_end = d["s"]
                     break
+            if i + 1 == len(sta_aoi_ids) - 1:
+                aoi_end["subline_driving_lane_pairs"].append(
+                    {
+                        "subline_id": subline_id,
+                        "driving_lane_id": aoi2driving_lane_id(aoi_end, road_ids[-1]),
+                    }
+                )
             route_lengths.append(_station_distance(road_ids, s_start, s_end))
         arg = (subline, station_aois, (x_min, x_step, y_min, y_step), route_lengths)
         taz_cost_args.append(arg)
@@ -373,14 +404,6 @@ def _station_distance(road_ids, s_start: float, s_end: float) -> float:
     for subline in sublines_data:
         subline_id = subline["id"]
         subline["taz_costs"] = subline_id2taz_costs[subline_id]
-    # for sl in sublines_data:
-    #     if not sl["type"] == mapv2.SUBLINE_TYPE_SUBWAY:
-    #         continue
-    #     transfer_stations = []
-    #     for aid in sl["aoi_ids"]:
-    #         station = aois[aid]
-    #         if len(station["subline_ids"]) > 1:
-    #             transfer_stations.append(station)
 
     return m
 

mosstool/trip/generator/_util/utils.py

@@ -1,16 +1,25 @@
 from multiprocessing import Pool, cpu_count
-from typing import Dict, List
+from typing import Dict, List, Literal, Optional, Set, Tuple, Union, cast
 
 import numpy as np
+from pycityproto.city.person.v1.person_pb2 import (BusAttribute, BusType,
+                                                   PersonAttribute,
+                                                   PersonProfile)
+from pycityproto.city.routing.v2.routing_pb2 import (DrivingJourneyBody,
+                                                     Journey, JourneyType)
+from pycityproto.city.trip.v2.trip_pb2 import TripStop
 
 from ....map._map_util.const import *
-from ....type import Consumption, Education, TripMode
+from ....type import (AoiPosition, Consumption, Education, LanePosition, Map,
+                      Person, Position, Schedule, Trip, TripMode)
+from ....util.format_converter import dict2pb, pb2dict
 from .const import *
 
 __all__ = [
     "is_walking",
     "gen_profiles",
     "recalculate_trip_mode_prob",
+    "gen_bus_drivers",
 ]
 
 
@@ -230,3 +239,134 @@ def extract_HWEO_from_od_matrix(
     other_od = np.nan_to_num(other_od)
     other_od[np.where(sum_other_od_j == 0)] = 1
     return home_dist, work_od, educate_od, other_od
+
+
+def gen_bus_drivers(
+    person_id: int,
+    person_template: Person,
+    depart_times: List[float],
+    stop_duration_time: float,
+    road_aoi_id2d_pos: Dict[Tuple[int, int], geov2.LanePosition],
+    subline,
+) -> Tuple[int, List[Person]]:
+    def _transfer_conn_road_ids(
+        station_connection_road_ids: List[List[int]],
+    ) -> List[int]:
+        assert (
+            len(station_connection_road_ids) > 0
+            and len(station_connection_road_ids[0]) > 0
+        ), f"Bad conn_road_ids {station_connection_road_ids}"
+        route_road_ids = []
+        for next_road_ids in station_connection_road_ids:
+            if len(route_road_ids) > 0 and route_road_ids[-1] == next_road_ids[0]:
+                route_road_ids += next_road_ids[1:]
+            else:
+                route_road_ids += next_road_ids
+        return route_road_ids
+
+    def _aoi_road_ids(station_connection_road_ids) -> List[int]:
+        cur_road_ids = [rids[0] for rids in station_connection_road_ids]
+        # road that connect to the end station
+        cur_road_ids.append(station_connection_road_ids[-1][-1])
+        return cur_road_ids
+
+    sl_id = subline.id
+    # basic attributes
+    sl_type = subline.type
+    bus_type = BusType.BUS_TYPE_UNSPECIFIED
+    if sl_type == mapv2.SUBLINE_TYPE_BUS:
+        sl_capacity = STATION_CAPACITY["BUS"]
+        sl_attributes = PT_DRIVER_ATTRIBUTES["BUS"]
+        bus_type = BusType.BUS_TYPE_BUS
+    elif sl_type == mapv2.SUBLINE_TYPE_SUBWAY:
+        sl_capacity = STATION_CAPACITY["SUBWAY"]
+        sl_attributes = PT_DRIVER_ATTRIBUTES["SUBWAY"]
+        bus_type = BusType.BUS_TYPE_SUBWAY
+    elif sl_type == mapv2.SUBLINE_TYPE_UNSPECIFIED:
+        sl_capacity = STATION_CAPACITY["UNSPECIFIED"]
+        sl_attributes = PT_DRIVER_ATTRIBUTES["UNSPECIFIED"]
+        bus_type = BusType.BUS_TYPE_UNSPECIFIED
+    else:
+        raise ValueError(f"Bad Subline Type {sl_type}")
+    # road ids
+    station_connection_road_ids = [
+        [rid for rid in rids.road_ids] for rids in subline.station_connection_road_ids
+    ]
+    route_road_ids = _transfer_conn_road_ids(station_connection_road_ids)
+    trip_stop_road_ids = _aoi_road_ids(
+        station_connection_road_ids
+    )  # stop road ids during the trip
+    # trip stops
+    aoi_ids = list(subline.aoi_ids)
+    home_aoi_id, end_aoi_id = aoi_ids[0], aoi_ids[-1]
+    trip_stop_aoi_ids = aoi_ids  # stop aoi ids during the trip
+    trip_stop_lane_id_s = []
+    for cur_road_id, cur_aoi_id in zip(trip_stop_road_ids, trip_stop_aoi_ids):
+        road_aoi_key = (cur_road_id, cur_aoi_id)
+        if road_aoi_key not in road_aoi_id2d_pos:
+            raise ValueError(f"bad road and AOI pair {road_aoi_key}")
+        d_pos = road_aoi_id2d_pos[road_aoi_key]
+        d_lane_id, d_s = d_pos.lane_id, d_pos.s
+        trip_stop_lane_id_s.append((d_lane_id, d_s))
+    assert len(trip_stop_lane_id_s) == len(
+        trip_stop_aoi_ids
+    ), f"Bad PublicTransport Route at {aoi_ids}"
+    p_trip_stops = []
+    # bus attribute
+    p_bus_attr = BusAttribute(
+        subline_id=sl_id, capacity=sl_capacity, type=bus_type, model=""
+    )
+    for (d_lane_id, d_s), aoi_id in zip(trip_stop_lane_id_s, trip_stop_aoi_ids):
+        trip_stop = TripStop()
+        trip_stop.lane_position.CopyFrom(LanePosition(lane_id=d_lane_id, s=d_s))
+        trip_stop.aoi_position.CopyFrom(AoiPosition(aoi_id=aoi_id))
+        trip_stop.duration = stop_duration_time
+        p_trip_stops.append(trip_stop)
+    # eta for bus journey
+    bus_eta = sum(subline.schedules.offset_times)
+    sl_drivers = []
+    if bus_type == BusType.BUS_TYPE_BUS:
+        for tm in depart_times:
+            p = Person()
+            p.CopyFrom(person_template)
+            p.id = person_id
+            if sl_attributes:
+                p.attribute.CopyFrom(dict2pb(sl_attributes, PersonAttribute()))
+            p.bus_attribute.CopyFrom(p_bus_attr)
+            p.home.CopyFrom(Position(aoi_position=AoiPosition(aoi_id=home_aoi_id)))
+            schedule = cast(Schedule, p.schedules.add())
+            schedule.departure_time = tm
+            schedule.loop_count = 1
+            trip = Trip(
+                mode=cast(
+                    TripMode,
+                    CAR,
+                ),
+                end=Position(aoi_position=AoiPosition(aoi_id=end_aoi_id)),
+                activity="",
+                model="",
+                trip_stops=p_trip_stops,
+                routes=[
+                    Journey(
+                        driving=DrivingJourneyBody(
+                            road_ids=route_road_ids, eta=bus_eta
+                        ),
+                        type=JourneyType.JOURNEY_TYPE_DRIVING,
+                    )
+                ],
+            )
+            schedule.trips.append(trip)
+            person_id += 1
+            sl_drivers.append(p)
+    elif bus_type == BusType.BUS_TYPE_SUBWAY:
+        # empty schedule
+        p = Person()
+        p.CopyFrom(person_template)
+        p.id = person_id
+        if sl_attributes:
+            p.attribute.CopyFrom(dict2pb(sl_attributes, PersonAttribute()))
+        p.bus_attribute.CopyFrom(p_bus_attr)
+        p.home.CopyFrom(Position(aoi_position=AoiPosition(aoi_id=home_aoi_id)))
+        person_id += 1
+        sl_drivers.append(p)
+    return (person_id, sl_drivers)