cli/
augment_scenario.rs

1use rand::prelude::SliceRandom;
2use rand::{Rng, SeedableRng};
3use rand_xorshift::XorShiftRng;
4
5use abstutil::{prettyprint_usize, Timer};
6use geom::{Distance, Duration, FindClosest};
7use map_model::{AmenityType, BuildingID, Map};
8use synthpop::{IndividTrip, Scenario, ScenarioModifier, TripEndpoint, TripMode, TripPurpose};
9
10pub fn run(
11    input_scenario: String,
12    should_add_return_trips: bool,
13    should_add_lunch_trips: bool,
14    modifiers: Vec<ScenarioModifier>,
15    should_delete_cancelled_trips: bool,
16    rng_seed: u64,
17) {
18    let mut rng = XorShiftRng::seed_from_u64(rng_seed);
19    let mut timer = Timer::new("augment scenario");
20
21    let mut scenario: Scenario = abstio::must_read_object(input_scenario, &mut timer);
22    let map = Map::load_synchronously(scenario.map_name.path(), &mut timer);
23
24    if should_add_return_trips {
25        add_return_trips(&mut scenario, &mut rng);
26    }
27    if should_add_lunch_trips {
28        add_lunch_trips(&mut scenario, &map, &mut rng, &mut timer);
29    }
30
31    for m in modifiers {
32        scenario = m.apply(&map, scenario, &mut rng);
33    }
34
35    if should_delete_cancelled_trips {
36        scenario = delete_cancelled_trips(scenario);
37    }
38
39    scenario.save();
40}
41
42fn add_return_trips(scenario: &mut Scenario, rng: &mut XorShiftRng) {
43    let mut cnt = 0;
44    for person in &mut scenario.people {
45        if person.trips.len() != 1 {
46            continue;
47        }
48
49        // Assume a uniform distribution of 4-12 hour workday
50        let depart =
51            person.trips[0].depart + rand_duration(rng, Duration::hours(4), Duration::hours(12));
52        person.trips.push(IndividTrip::new(
53            depart,
54            TripPurpose::Home,
55            person.trips[0].destination,
56            person.trips[0].origin,
57            person.trips[0].mode,
58        ));
59        cnt += 1;
60    }
61    info!("Added return trips to {} people", prettyprint_usize(cnt));
62}
63
64fn rand_duration(rng: &mut XorShiftRng, low: Duration, high: Duration) -> Duration {
65    Duration::seconds(rng.gen_range(low.inner_seconds()..high.inner_seconds()))
66}
67
68fn add_lunch_trips(scenario: &mut Scenario, map: &Map, rng: &mut XorShiftRng, timer: &mut Timer) {
69    // First let's build up a quadtree of lunch spots.
70    timer.start("index lunch spots");
71    let mut closest_spots: FindClosest<BuildingID> = FindClosest::new();
72    for b in map.all_buildings() {
73        if b.amenities
74            .iter()
75            .any(|a| AmenityType::categorize(&a.amenity_type) == Some(AmenityType::Food))
76        {
77            closest_spots.add_polygon(b.id, &b.polygon);
78        }
79    }
80    timer.stop("index lunch spots");
81
82    timer.start_iter("add lunch trips", scenario.people.len());
83    let mut cnt = 0;
84    for person in &mut scenario.people {
85        timer.next();
86        let num_trips = person.trips.len();
87        // Only handle people with their final trip going back home.
88        if num_trips <= 1 || person.trips[num_trips - 1].destination != person.trips[0].origin {
89            continue;
90        }
91
92        let work = match person.trips[num_trips - 2].destination {
93            TripEndpoint::Building(b) => b,
94            _ => continue,
95        };
96        let has_bike = person.trips[num_trips - 2].mode == TripMode::Bike;
97        let (restaurant, mode) =
98            if let Some(pair) = pick_lunch_spot(work, has_bike, &closest_spots, map, rng) {
99                pair
100            } else {
101                continue;
102            };
103
104        // Insert the break in the middle of their workday
105        let t1 = person.trips[num_trips - 2].depart;
106        let t2 = person.trips[num_trips - 1].depart;
107        let depart = t1 + (t2 - t1) / 2.0;
108        let return_home = person.trips.pop().unwrap();
109        person.trips.push(IndividTrip::new(
110            depart,
111            TripPurpose::Meal,
112            TripEndpoint::Building(work),
113            TripEndpoint::Building(restaurant),
114            mode,
115        ));
116        person.trips.push(IndividTrip::new(
117            depart + Duration::minutes(30),
118            TripPurpose::Work,
119            TripEndpoint::Building(restaurant),
120            TripEndpoint::Building(work),
121            mode,
122        ));
123        person.trips.push(return_home);
124        cnt += 1;
125    }
126    info!("Added lunch trips to {} people", prettyprint_usize(cnt));
127}
128
129fn pick_lunch_spot(
130    work: BuildingID,
131    has_bike: bool,
132    closest_spots: &FindClosest<BuildingID>,
133    map: &Map,
134    rng: &mut XorShiftRng,
135) -> Option<(BuildingID, TripMode)> {
136    // We have a list of candidate shops and the Euclidean distance there. Use that distance to
137    // make a weighted choice.
138    let choices =
139        closest_spots.all_close_pts(map.get_b(work).polygon.center(), Distance::miles(10.0));
140    let (b, _, dist) = choices
141        .choose_weighted(rng, |(_, _, dist)| dist.inner_meters())
142        .ok()?;
143    // Simple hardcoded mode thresholds for now
144    let mode = if *dist <= Distance::miles(1.0) {
145        TripMode::Walk
146    } else if has_bike {
147        TripMode::Bike
148    } else {
149        TripMode::Drive
150    };
151    Some((*b, mode))
152}
153
154fn delete_cancelled_trips(mut scenario: Scenario) -> Scenario {
155    for person in &mut scenario.people {
156        person.trips.retain(|trip| !trip.cancelled);
157    }
158    scenario.remove_weird_schedules(false)
159}