Extract some GTFS data and put it in RawMap. #372

convert_osm/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 abstio = { path = "../abstio" }
 abstutil = { path = "../abstutil" }
 anyhow = "1.0.38"
+csv = "1.1.4"
 geom = { path = "../geom" }
 kml = { path = "../kml" }
 log = "0.4.14"

convert_osm/src/gtfs.rs

@@ -0,0 +1,215 @@
+use std::collections::{HashMap, HashSet};
+use std::fs::File;
+
+use anyhow::Result;
+use serde::Deserialize;
+
+use abstutil::MultiMap;
+use geom::{LonLat, PolyLine, Pt2D};
+use map_model::raw::{RawMap, RawTransitRoute, RawTransitStop};
+use map_model::PathConstraints;
+
+pub fn import(map: &mut RawMap) -> Result<()> {
+    // Collect metadata about routes
+    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/routes.txt"))?)
+        .deserialize()
+    {
+        let rec: Route = rec?;
+        // See https://developers.google.com/transit/gtfs/reference#routestxt
+        let route_type = match rec.route_type {
+            3 => PathConstraints::Bus,
+            // These aren't distinguished in the map model yet. Trams and streetcars might
+            // particularly mess up...  or just fail to snap to a road later.
+            0 | 1 | 2 => PathConstraints::Train,
+            _ => continue,
+        };
+        map.transit_routes.push(RawTransitRoute {
+            long_name: rec.route_long_name,
+            short_name: rec.route_short_name,
+            gtfs_id: rec.route_id.0,
+            shape: PolyLine::dummy(),
+            stops: Vec::new(),
+            route_type,
+        });
+    }
+
+    // Map route_id to shape_id
+    let mut route_to_shapes = MultiMap::new();
+    // Map (route_id, shape_id) to trip_id
+    let mut route_and_shape_to_trips = MultiMap::new();
+    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/trips.txt"))?)
+        .deserialize()
+    {
+        let rec: Trip = rec?;
+        route_to_shapes.insert(rec.route_id.clone(), rec.shape_id.clone());
+        route_and_shape_to_trips.insert((rec.route_id, rec.shape_id), rec.trip_id);
+    }
+
+    // Scrape all shape data. Map from shape_id to points and the sequence number
+    let mut raw_shapes: HashMap<ShapeID, Vec<(Pt2D, usize)>> = HashMap::new();
+    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/shapes.txt"))?)
+        .deserialize()
+    {
+        let rec: Shape = rec?;
+        let pt = LonLat::new(rec.shape_pt_lon, rec.shape_pt_lat).to_pt(&map.gps_bounds);
+        raw_shapes
+            .entry(rec.shape_id)
+            .or_insert_with(Vec::new)
+            .push((pt, rec.shape_pt_sequence));
+    }
+
+    // Build a PolyLine for every route
+    let mut transit_routes = Vec::new();
+    let mut route_to_shape = HashMap::new();
+    for mut route in map.transit_routes.drain(..) {
+        let shape_ids = route_to_shapes.get(RouteID(route.gtfs_id.clone()));
+        if shape_ids.is_empty() {
+            warn!("Route {} has no shape", route.gtfs_id);
+            continue;
+        }
+        if shape_ids.len() > 1 {
+            warn!(
+                "Route {} has several shapes, choosing one arbitrarily: {:?}",
+                route.gtfs_id, shape_ids
+            );
+        }
+        let shape_id = shape_ids.into_iter().next().unwrap();
+        route_to_shape.insert(RouteID(route.gtfs_id.clone()), shape_id.clone());
+        let mut pts = if let Some(pts) = raw_shapes.remove(shape_id) {
+            pts
+        } else {
+            warn!("Route {} is missing its shape", route.gtfs_id);
+            continue;
+        };
+        // Points are usually sorted, but just in case...
+        pts.sort_by_key(|(_, seq)| *seq);
+        let pts: Vec<Pt2D> = pts.into_iter().map(|(pt, _)| pt).collect();
+        match PolyLine::new(pts) {
+            Ok(pl) => {
+                route.shape = pl;
+                transit_routes.push(route);
+            }
+            Err(err) => {
+                warn!("Route {} has a weird shape: {}", route.gtfs_id, err);
+                continue;
+            }
+        }
+    }
+    map.transit_routes = transit_routes;
+
+    // For now, every route uses exactly one trip ID, and there's no schedule. Just pick an
+    // arbitrary trip per route.
+    let mut route_to_trip = HashMap::new();
+    for (route_id, shape_id) in &route_to_shape {
+        let trips = route_and_shape_to_trips.get((route_id.clone(), shape_id.clone()));
+        if let Some(trip_id) = trips.into_iter().next() {
+            route_to_trip.insert(route_id.clone(), trip_id);
+        }
+    }
+
+    // Scrape the trip ID -> (stop ID, sequence number)
+    let mut trip_to_stops: HashMap<TripID, Vec<(StopID, usize)>> = HashMap::new();
+    for rec in
+        csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/stop_times.txt"))?)
+            .deserialize()
+    {
+        let rec: StopTime = rec?;
+        trip_to_stops
+            .entry(rec.trip_id)
+            .or_insert_with(Vec::new)
+            .push((rec.stop_id, rec.stop_sequence));
+    }
+
+    // Assign the stops for every route
+    let mut stop_ids = HashSet::new();
+    for route in &mut map.transit_routes {
+        let trip_id = route_to_trip[&RouteID(route.gtfs_id.clone())];
+        let mut stops = trip_to_stops.remove(&trip_id).unwrap_or_else(Vec::new);
+        stops.sort_by_key(|(_, seq)| *seq);
+        for (stop_id, _) in stops {
+            route.stops.push(stop_id.0.clone());
+            stop_ids.insert(stop_id);
+        }
+    }
+
+    // Scrape stop metadata
+    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/stops.txt"))?)
+        .deserialize()
+    {
+        let rec: Stop = rec?;
+        if stop_ids.contains(&rec.stop_id) {
+            let position = LonLat::new(rec.stop_lon, rec.stop_lat).to_pt(&map.gps_bounds);
+            if map.boundary_polygon.contains_pt(position) {
+                map.transit_stops.insert(
+                    rec.stop_id.0.clone(),
+                    RawTransitStop {
+                        gtfs_id: rec.stop_id.0,
+                        position,
+                        name: rec.stop_name,
+                    },
+                );
+            }
+        }
+    }
+
+    // Make sure all of the stops are valid and used by some route
+    let mut used_stops = HashSet::new();
+    for route in &mut map.transit_routes {
+        route.stops.retain(|stop_id| {
+            used_stops.insert(stop_id.clone());
+            map.transit_stops.contains_key(stop_id)
+        });
+    }
+    map.transit_routes.retain(|route| !route.stops.is_empty());
+    map.transit_stops
+        .retain(|stop_id, _| used_stops.contains(stop_id));
+
+    Ok(())
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
+struct ShapeID(String);
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
+struct TripID(String);
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
+struct StopID(String);
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
+struct RouteID(String);
+
+#[derive(Deserialize)]
+struct Route {
+    route_id: RouteID,
+    route_short_name: String,
+    route_long_name: String,
+    route_type: usize,
+}
+
+#[derive(Deserialize)]
+struct Trip {
+    route_id: RouteID,
+    shape_id: ShapeID,
+    trip_id: TripID,
+}
+
+#[derive(Deserialize)]
+struct Shape {
+    shape_id: ShapeID,
+    shape_pt_lat: f64,
+    shape_pt_lon: f64,
+    shape_pt_sequence: usize,
+}
+
+#[derive(Deserialize)]
+struct Stop {
+    stop_id: StopID,
+    stop_lon: f64,
+    stop_lat: f64,
+    stop_name: String,
+}
+
+#[derive(Deserialize)]
+struct StopTime {
+    trip_id: TripID,
+    stop_id: StopID,
+    stop_sequence: usize,
+}

convert_osm/src/lib.rs

@@ -17,6 +17,7 @@ use serde::{Deserialize, Serialize};
 mod clip;
 mod elevation;
 mod extract;
+mod gtfs;
 pub mod osm_geom;
 mod parking;
 pub mod reader;
@@ -119,6 +120,10 @@ pub fn convert(
         filter_crosswalks(&mut map, crosswalks, pt_to_road, timer);
     }
 
+    if opts.gtfs_url.is_some() {
+        gtfs::import(&mut map).unwrap();
+    }
+
     map.config = opts.map_config;
     map
 }

geom/src/polyline.rs

@@ -990,6 +990,11 @@ impl PolyLine {
 
         PolyLine::new(output).unwrap()
     }
+
+    /// An arbitrary placeholder value, when Option types aren't worthwhile
+    pub fn dummy() -> PolyLine {
+        PolyLine::must_new(vec![Pt2D::new(0.0, 0.0), Pt2D::new(0.1, 0.1)])
+    }
 }
 
 impl fmt::Display for PolyLine {

map_model/src/raw.rs

@@ -16,6 +16,7 @@ use geom::{Distance, GPSBounds, PolyLine, Polygon, Pt2D};
 use crate::make::initial::lane_specs::get_lane_specs_ltr;
 use crate::{
     osm, Amenity, AreaType, Direction, DrivingSide, IntersectionType, LaneType, MapConfig,
+    PathConstraints,
 };
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -39,6 +40,12 @@ pub struct RawMap {
     pub areas: Vec<RawArea>,
     pub parking_lots: Vec<RawParkingLot>,
     pub parking_aisles: Vec<(osm::WayID, Vec<Pt2D>)>,
+    pub transit_routes: Vec<RawTransitRoute>,
+    #[serde(
+        serialize_with = "serialize_btreemap",
+        deserialize_with = "deserialize_btreemap"
+    )]
+    pub transit_stops: BTreeMap<String, RawTransitStop>,
 
     pub boundary_polygon: Polygon,
     pub gps_bounds: GPSBounds,
@@ -114,6 +121,8 @@ impl RawMap {
             areas: Vec::new(),
             parking_lots: Vec::new(),
             parking_aisles: Vec::new(),
+            transit_routes: Vec::new(),
+            transit_stops: BTreeMap::new(),
             // Some nonsense thing
             boundary_polygon: Polygon::rectangle(1.0, 1.0),
             gps_bounds: GPSBounds::new(),
@@ -720,3 +729,24 @@ impl RestrictionType {
         }
     }
 }
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct RawTransitRoute {
+    pub long_name: String,
+    pub short_name: String,
+    pub gtfs_id: String,
+    /// This may begin and/or end inside or outside the map boundary.
+    pub shape: PolyLine,
+    /// Entries into transit_stops
+    pub stops: Vec<String>,
+    pub route_type: PathConstraints,
+    // TODO Schedule
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct RawTransitStop {
+    pub gtfs_id: String,
+    /// Only stops within a map's boundary are kept
+    pub position: Pt2D,
+    pub name: String,
+}