1use std::collections::BTreeSet;
5
6use rand::seq::SliceRandom;
7use rand::Rng;
8use rand_xorshift::XorShiftRng;
9use serde::{Deserialize, Serialize};
10
11use abstutil::Timer;
12use geom::{Duration, Time};
13use map_model::{IntersectionID, Map};
14
15use crate::{IndividTrip, PersonSpec, Scenario, TripEndpoint, TripMode, TripPurpose};
16
17#[derive(Clone, Serialize, Deserialize, Debug)]
20pub struct ScenarioGenerator {
21 pub scenario_name: String,
22
23 pub only_seed_buses: Option<BTreeSet<String>>,
24 pub spawn_over_time: Vec<SpawnOverTime>,
25 pub border_spawn_over_time: Vec<BorderSpawnOverTime>,
26}
27
28#[derive(Clone, Serialize, Deserialize, Debug)]
32pub struct SpawnOverTime {
33 pub num_agents: usize,
34 pub start_time: Time,
36 pub stop_time: Time,
37 pub goal: Option<TripEndpoint>,
38 pub percent_driving: f64,
39 pub percent_biking: f64,
40 pub percent_use_transit: f64,
41}
42
43#[derive(Clone, Serialize, Deserialize, Debug)]
44pub struct BorderSpawnOverTime {
45 pub num_peds: usize,
46 pub num_cars: usize,
47 pub num_bikes: usize,
48 pub percent_use_transit: f64,
49 pub start_time: Time,
51 pub stop_time: Time,
52 pub start_from_border: IntersectionID,
53 pub goal: Option<TripEndpoint>,
54}
55
56impl ScenarioGenerator {
57 pub fn generate(&self, map: &Map, rng: &mut XorShiftRng, timer: &mut Timer) -> Scenario {
59 let mut scenario = Scenario::empty(map, &self.scenario_name);
60 scenario.only_seed_buses = self.only_seed_buses.clone();
61
62 timer.start(format!("Generating scenario {}", self.scenario_name));
63
64 for s in &self.spawn_over_time {
65 timer.start_iter("SpawnOverTime each agent", s.num_agents);
66 for _ in 0..s.num_agents {
67 timer.next();
68 s.spawn_agent(rng, &mut scenario, map);
69 }
70 }
71
72 timer.start_iter("BorderSpawnOverTime", self.border_spawn_over_time.len());
73 for s in &self.border_spawn_over_time {
74 timer.next();
75 for _ in 0..s.num_peds {
76 let mode = if rng.gen_bool(s.percent_use_transit) {
77 TripMode::Transit
78 } else {
79 TripMode::Walk
80 };
81 s.spawn(rng, &mut scenario, mode, map);
82 }
83 for _ in 0..s.num_cars {
84 s.spawn(rng, &mut scenario, TripMode::Drive, map);
85 }
86 for _ in 0..s.num_bikes {
87 s.spawn(rng, &mut scenario, TripMode::Bike, map);
88 }
89 }
90
91 timer.stop(format!("Generating scenario {}", self.scenario_name));
92 scenario.remove_weird_schedules(true)
93 }
94
95 pub fn small_run(map: &Map) -> ScenarioGenerator {
96 let mut s = ScenarioGenerator {
97 scenario_name: "small_run".to_string(),
98 only_seed_buses: None,
99 spawn_over_time: vec![SpawnOverTime {
100 num_agents: 100,
101 start_time: Time::START_OF_DAY,
102 stop_time: Time::START_OF_DAY + Duration::seconds(5.0),
103 goal: None,
104 percent_driving: 0.5,
105 percent_biking: 0.5,
106 percent_use_transit: 0.5,
107 }],
108 border_spawn_over_time: map
111 .all_incoming_borders()
112 .into_iter()
113 .map(|i| BorderSpawnOverTime {
114 num_peds: 10,
115 num_cars: 10,
116 num_bikes: 10,
117 start_time: Time::START_OF_DAY,
118 stop_time: Time::START_OF_DAY + Duration::seconds(5.0),
119 start_from_border: i.id,
120 goal: None,
121 percent_use_transit: 0.5,
122 })
123 .collect(),
124 };
125 for i in map.all_outgoing_borders() {
126 s.spawn_over_time.push(SpawnOverTime {
127 num_agents: 10,
128 start_time: Time::START_OF_DAY,
129 stop_time: Time::START_OF_DAY + Duration::seconds(5.0),
130 goal: Some(TripEndpoint::Border(i.id)),
131 percent_driving: 0.5,
132 percent_biking: 0.5,
133 percent_use_transit: 0.5,
134 });
135 }
136 s
137 }
138
139 pub fn empty(name: &str) -> ScenarioGenerator {
140 ScenarioGenerator {
141 scenario_name: name.to_string(),
142 only_seed_buses: Some(BTreeSet::new()),
143 spawn_over_time: Vec::new(),
144 border_spawn_over_time: Vec::new(),
145 }
146 }
147}
148
149impl SpawnOverTime {
150 fn spawn_agent(&self, rng: &mut XorShiftRng, scenario: &mut Scenario, map: &Map) {
151 let depart = rand_time(rng, self.start_time, self.stop_time);
152 let from_bldg = map.all_buildings().choose(rng).unwrap().id;
155 let mode = if rng.gen_bool(self.percent_driving) {
156 TripMode::Drive
157 } else if rng.gen_bool(self.percent_biking) {
158 TripMode::Bike
159 } else if rng.gen_bool(self.percent_use_transit) {
160 TripMode::Transit
161 } else {
162 TripMode::Walk
163 };
164 scenario.people.push(PersonSpec {
165 orig_id: None,
166 trips: vec![IndividTrip::new(
167 depart,
168 TripPurpose::Shopping,
169 TripEndpoint::Building(from_bldg),
170 self.goal.unwrap_or_else(|| {
171 TripEndpoint::Building(map.all_buildings().choose(rng).unwrap().id)
172 }),
173 mode,
174 )],
175 });
176 }
177}
178
179impl BorderSpawnOverTime {
180 fn spawn(&self, rng: &mut XorShiftRng, scenario: &mut Scenario, mode: TripMode, map: &Map) {
181 let depart = rand_time(rng, self.start_time, self.stop_time);
182 scenario.people.push(PersonSpec {
183 orig_id: None,
184 trips: vec![IndividTrip::new(
185 depart,
186 TripPurpose::Shopping,
187 TripEndpoint::Border(self.start_from_border),
188 self.goal.unwrap_or_else(|| {
189 TripEndpoint::Building(map.all_buildings().choose(rng).unwrap().id)
190 }),
191 mode,
192 )],
193 });
194 }
195}
196
197fn rand_time(rng: &mut XorShiftRng, low: Time, high: Time) -> Time {
198 assert!(high > low);
199 Time::START_OF_DAY + Duration::seconds(rng.gen_range(low.inner_seconds()..high.inner_seconds()))
200}