convert_osm/
gtfs.rs

1use std::collections::{BTreeMap, HashMap, HashSet};
2
3use anyhow::Result;
4use fs_err::File;
5use serde::Deserialize;
6
7use abstutil::MultiMap;
8use geom::{LonLat, PolyLine, Pt2D};
9use kml::{ExtraShape, ExtraShapes};
10use raw_map::{RawMap, RawTransitRoute, RawTransitStop, RawTransitType};
11
12pub fn import(map: &mut RawMap) -> Result<()> {
13    // Collect metadata about routes
14    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/routes.txt"))?)
15        .deserialize()
16    {
17        let rec: Route = rec?;
18        // See https://developers.google.com/transit/gtfs/reference#routestxt
19        let route_type = match rec.route_type {
20            3 => RawTransitType::Bus,
21            // These aren't distinguished in the map model yet. Trams and streetcars might
22            // particularly mess up...  or just fail to snap to a road later.
23            0 | 1 | 2 => RawTransitType::Train,
24            _ => continue,
25        };
26        map.transit_routes.push(RawTransitRoute {
27            long_name: if rec.route_long_name.is_empty() {
28                rec.route_desc
29            } else {
30                rec.route_long_name
31            },
32            short_name: rec.route_short_name,
33            gtfs_id: rec.route_id.0,
34            shape: PolyLine::dummy(),
35            stops: Vec::new(),
36            route_type,
37        });
38    }
39
40    // Map route_id to shape_id
41    let mut route_to_shapes = MultiMap::new();
42    // Map (route_id, shape_id) to trip_id
43    let mut route_and_shape_to_trips = MultiMap::new();
44    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/trips.txt"))?)
45        .deserialize()
46    {
47        let rec: Trip = rec?;
48        route_to_shapes.insert(rec.route_id.clone(), rec.shape_id.clone());
49        route_and_shape_to_trips.insert((rec.route_id, rec.shape_id), rec.trip_id);
50    }
51
52    // Scrape all shape data. Map from shape_id to points and the sequence number
53    //
54    // If this file is missing, one idea is to just draw straight lines between stops. We only use
55    // the shape currently to pick an entry/exit border, so this could be a half-reasonable
56    // workaround.
57    let mut raw_shapes: HashMap<ShapeID, Vec<(Pt2D, usize)>> = HashMap::new();
58    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/shapes.txt"))?)
59        .deserialize()
60    {
61        let rec: Shape = rec?;
62        let pt = LonLat::new(rec.shape_pt_lon, rec.shape_pt_lat).to_pt(&map.streets.gps_bounds);
63        raw_shapes
64            .entry(rec.shape_id)
65            .or_insert_with(Vec::new)
66            .push((pt, rec.shape_pt_sequence));
67    }
68
69    // Build a PolyLine for every route
70    let mut transit_routes = Vec::new();
71    let mut route_to_shape = HashMap::new();
72    for mut route in map.transit_routes.drain(..) {
73        let shape_ids = route_to_shapes.get(RouteID(route.gtfs_id.clone()));
74        if shape_ids.is_empty() {
75            warn!("Route {} has no shape", route.gtfs_id);
76            continue;
77        }
78        if shape_ids.len() > 1 {
79            warn!(
80                "Route {} has several shapes, choosing one arbitrarily: {:?}",
81                route.gtfs_id, shape_ids
82            );
83        }
84        let shape_id = shape_ids.into_iter().next().unwrap();
85        route_to_shape.insert(RouteID(route.gtfs_id.clone()), shape_id.clone());
86        let mut pts = if let Some(pts) = raw_shapes.remove(shape_id) {
87            pts
88        } else {
89            warn!("Route {} is missing its shape", route.gtfs_id);
90            continue;
91        };
92        // Points are usually sorted, but just in case...
93        pts.sort_by_key(|(_, seq)| *seq);
94        let pts: Vec<Pt2D> = pts.into_iter().map(|(pt, _)| pt).collect();
95        match PolyLine::new(pts) {
96            Ok(pl) => {
97                route.shape = pl;
98                transit_routes.push(route);
99            }
100            Err(err) => {
101                warn!("Route {} has a weird shape: {}", route.gtfs_id, err);
102                continue;
103            }
104        }
105    }
106    map.transit_routes = transit_routes;
107
108    // For now, every route uses exactly one trip ID, and there's no schedule. Just pick an
109    // arbitrary trip per route.
110    let mut route_to_trip = HashMap::new();
111    for (route_id, shape_id) in &route_to_shape {
112        let trips = route_and_shape_to_trips.get((route_id.clone(), shape_id.clone()));
113        if let Some(trip_id) = trips.into_iter().next() {
114            route_to_trip.insert(route_id.clone(), trip_id);
115        }
116    }
117
118    // Scrape the trip ID -> (stop ID, sequence number)
119    let mut trip_to_stops: HashMap<TripID, Vec<(StopID, usize)>> = HashMap::new();
120    for rec in
121        csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/stop_times.txt"))?)
122            .deserialize()
123    {
124        let rec: StopTime = rec?;
125        trip_to_stops
126            .entry(rec.trip_id)
127            .or_insert_with(Vec::new)
128            .push((rec.stop_id, rec.stop_sequence));
129    }
130
131    // Assign the stops for every route
132    let mut stop_ids = HashSet::new();
133    for route in &mut map.transit_routes {
134        let trip_id = route_to_trip[&RouteID(route.gtfs_id.clone())];
135        let mut stops = trip_to_stops.remove(&trip_id).unwrap_or_else(Vec::new);
136        stops.sort_by_key(|(_, seq)| *seq);
137        for (stop_id, _) in stops {
138            route.stops.push(stop_id.0.clone());
139            stop_ids.insert(stop_id);
140        }
141    }
142
143    // Scrape stop metadata
144    for rec in csv::Reader::from_reader(File::open(map.name.city.input_path("gtfs/stops.txt"))?)
145        .deserialize()
146    {
147        let rec: Stop = rec?;
148        if stop_ids.contains(&rec.stop_id) {
149            let position = LonLat::new(rec.stop_lon, rec.stop_lat).to_pt(&map.streets.gps_bounds);
150            if map.streets.boundary_polygon.contains_pt(position) {
151                map.transit_stops.insert(
152                    rec.stop_id.0.clone(),
153                    RawTransitStop {
154                        gtfs_id: rec.stop_id.0,
155                        position,
156                        name: rec.stop_name,
157                    },
158                );
159            }
160        }
161    }
162
163    // Make sure all of the stops are valid and used by some route
164    let mut used_stops = HashSet::new();
165    for route in &mut map.transit_routes {
166        route.stops.retain(|stop_id| {
167            used_stops.insert(stop_id.clone());
168            map.transit_stops.contains_key(stop_id)
169        });
170    }
171    map.transit_routes.retain(|route| !route.stops.is_empty());
172    map.transit_stops
173        .retain(|stop_id, _| used_stops.contains(stop_id));
174
175    if false {
176        dump_kml(map);
177    }
178
179    Ok(())
180}
181
182#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
183struct ShapeID(String);
184#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
185struct TripID(String);
186#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
187struct StopID(String);
188#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize)]
189struct RouteID(String);
190
191#[derive(Deserialize)]
192struct Route {
193    route_id: RouteID,
194    route_short_name: String,
195    route_long_name: String,
196    // Missing from São Paulo
197    #[serde(default)]
198    route_desc: String,
199    route_type: usize,
200}
201
202#[derive(Deserialize)]
203struct Trip {
204    route_id: RouteID,
205    shape_id: ShapeID,
206    trip_id: TripID,
207}
208
209#[derive(Deserialize)]
210struct Shape {
211    shape_id: ShapeID,
212    shape_pt_lat: f64,
213    shape_pt_lon: f64,
214    shape_pt_sequence: usize,
215}
216
217#[derive(Deserialize)]
218struct Stop {
219    stop_id: StopID,
220    stop_lon: f64,
221    stop_lat: f64,
222    stop_name: String,
223}
224
225#[derive(Deserialize)]
226struct StopTime {
227    trip_id: TripID,
228    stop_id: StopID,
229    stop_sequence: usize,
230}
231
232fn dump_kml(map: &RawMap) {
233    let mut shapes = Vec::new();
234
235    // One polyline per route
236    for route in &map.transit_routes {
237        let points = map.streets.gps_bounds.convert_back(route.shape.points());
238        let mut attributes = BTreeMap::new();
239        attributes.insert("long_name".to_string(), route.long_name.clone());
240        attributes.insert("short_name".to_string(), route.short_name.clone());
241        attributes.insert("gtfs_id".to_string(), route.gtfs_id.clone());
242        attributes.insert("num_stops".to_string(), route.stops.len().to_string());
243        attributes.insert("route_type".to_string(), format!("{:?}", route.route_type));
244        shapes.push(ExtraShape { points, attributes });
245    }
246
247    // One point per stop
248    for stop in map.transit_stops.values() {
249        let mut attributes = BTreeMap::new();
250        attributes.insert("gtfs_id".to_string(), stop.gtfs_id.clone());
251        attributes.insert("name".to_string(), stop.name.clone());
252        let points = vec![stop.position.to_gps(&map.streets.gps_bounds)];
253        shapes.push(ExtraShape { points, attributes });
254    }
255
256    abstio::write_binary(
257        map.name
258            .city
259            .input_path(format!("gtfs_{}.bin", map.name.map)),
260        &ExtraShapes { shapes },
261    );
262}