checkbox to toggle running the pandemic model in the population map

game/src/common/overlays.rs

@@ -28,7 +28,7 @@ pub enum Overlays {
     Elevation(Colorer, Drawable),
     Edits(Colorer),
     TripsHistogram(Time, Composite),
-    PopulationMap(Time, Option<HeatmapOptions>, Drawable, Composite),
+    PopulationMap(Time, PopulationOptions, Drawable, Composite),
 
     // These aren't selectable from the main picker; they're particular to some object.
     // TODO They should become something else, like an info panel tab.
@@ -160,7 +160,7 @@ impl Overlays {
                         _ => unreachable!(),
                     },
                     None => {
-                        let new_opts = heatmap_options(c);
+                        let new_opts = population_options(c);
                         if *opts != new_opts {
                             app.overlay = Overlays::population_map(ctx, app, new_opts);
                         }
@@ -243,7 +243,6 @@ impl Overlays {
             Btn::text_fg("bus network").build_def(ctx, hotkey(Key::U)),
             Btn::text_fg("population map").build_def(ctx, hotkey(Key::X)),
         ];
-        // TODO Grey out the inactive SVGs, and add the green checkmark
         if let Some(name) = match app.overlay {
             Overlays::Inactive => Some("None"),
             Overlays::ParkingAvailability(_, _) => Some("parking availability"),
@@ -351,7 +350,14 @@ impl Overlays {
         .maybe_cb(
             "population map",
             Box::new(|ctx, app| {
-                app.overlay = Overlays::population_map(ctx, app, Some(HeatmapOptions::new()));
+                app.overlay = Overlays::population_map(
+                    ctx,
+                    app,
+                    PopulationOptions {
+                        pandemic: false,
+                        heatmap: Some(HeatmapOptions::new()),
+                    },
+                );
                 Some(maybe_unzoom(ctx, app))
             }),
         );
@@ -813,7 +819,11 @@ impl Overlays {
 
     // TODO Disable drawing unzoomed agents... or alternatively, implement this by asking Sim to
     // return this kind of data instead!
-    fn population_map(ctx: &mut EventCtx, app: &App, opts: Option<HeatmapOptions>) -> Overlays {
+    fn population_map(ctx: &mut EventCtx, app: &App, opts: PopulationOptions) -> Overlays {
+        if opts.pandemic {
+            return Overlays::pandemic_map(ctx, app, opts);
+        }
+
         let mut pts = Vec::new();
         // Faster to grab all agent positions than individually map trips to agent positions.
         for a in app.primary.sim.get_unzoomed_agents(&app.primary.map) {
@@ -839,14 +849,68 @@ impl Overlays {
         // It's quite silly to produce triangles for the same circle over and over again. ;)
         let circle = Circle::new(Pt2D::new(0.0, 0.0), Distance::meters(10.0)).to_polygon();
         let mut batch = GeomBatch::new();
-        if let Some(ref o) = opts {
+        if let Some(ref o) = opts.heatmap {
             make_heatmap(&mut batch, app.primary.map.get_bounds(), pts, o);
         } else {
             for pt in pts {
                 batch.push(Color::RED.alpha(0.8), circle.translate(pt.x(), pt.y()));
             }
         }
-        let controls = population_controls(ctx, app, opts.as_ref());
+        let controls = population_controls(ctx, app, &opts, None);
+        Overlays::PopulationMap(app.primary.sim.time(), opts, ctx.upload(batch), controls)
+    }
+
+    // Similar to population_map, but only display infected people.
+    fn pandemic_map(ctx: &mut EventCtx, app: &App, opts: PopulationOptions) -> Overlays {
+        // TODO Why not app.primary.current_flags.sim_flags.make_rng()? Because that'll only be the
+        // same every time this code runs (frequently, as the simulation is run) if --rng_seed is
+        // specified in the flags. If you forget it, quite confusing to see the model jump around.
+        use rand::SeedableRng;
+        use rand_xorshift::XorShiftRng;
+
+        let model = PandemicModel::calculate(
+            app.primary.sim.get_analytics(),
+            app.primary.sim.time(),
+            &mut XorShiftRng::from_seed([42; 16]),
+        );
+
+        let mut infected_pts = Vec::new();
+        // Faster to grab all agent positions than individually map trips to agent positions.
+        for a in app.primary.sim.get_unzoomed_agents(&app.primary.map) {
+            if let Some(p) = a.person {
+                if model.infected.contains(&p) {
+                    infected_pts.push(a.pos);
+                }
+            }
+        }
+
+        let mut seen_bldgs = HashSet::new();
+        for person in app.primary.sim.get_all_people() {
+            match person.state {
+                // Already covered above
+                PersonState::Trip(_) => {}
+                PersonState::Inside(b) => {
+                    // Duplicate circles for the same building are expensive!
+                    if model.infected.contains(&person.id) && !seen_bldgs.contains(&b) {
+                        seen_bldgs.insert(b);
+                        infected_pts.push(app.primary.map.get_b(b).polygon.center());
+                    }
+                }
+                PersonState::OffMap | PersonState::Limbo => {}
+            }
+        }
+
+        // It's quite silly to produce triangles for the same circle over and over again. ;)
+        let circle = Circle::new(Pt2D::new(0.0, 0.0), Distance::meters(10.0)).to_polygon();
+        let mut batch = GeomBatch::new();
+        if let Some(ref o) = opts.heatmap {
+            make_heatmap(&mut batch, app.primary.map.get_bounds(), infected_pts, o);
+        } else {
+            for pt in infected_pts {
+                batch.push(Color::RED.alpha(0.8), circle.translate(pt.x(), pt.y()));
+            }
+        }
+        let controls = population_controls(ctx, app, &opts, Some(model));
         Overlays::PopulationMap(app.primary.sim.time(), opts, ctx.upload(batch), controls)
     }
 }
@@ -864,8 +928,20 @@ fn maybe_unzoom(ctx: &EventCtx, app: &mut App) -> Transition {
     ))
 }
 
+#[derive(Clone, PartialEq)]
+pub struct PopulationOptions {
+    pandemic: bool,
+    // If None, just a dot map
+    heatmap: Option<HeatmapOptions>,
+}
+
 // This function sounds more ominous than it should.
-fn population_controls(ctx: &mut EventCtx, app: &App, opts: Option<&HeatmapOptions>) -> Composite {
+fn population_controls(
+    ctx: &mut EventCtx,
+    app: &App,
+    opts: &PopulationOptions,
+    pandemic: Option<PandemicModel>,
+) -> Composite {
     let (total_ppl, ppl_in_bldg, ppl_off_map) = app.primary.sim.num_ppl();
 
     let mut col = vec![
@@ -884,22 +960,10 @@ fn population_controls(ctx: &mut EventCtx, app: &App, opts: Option<&HeatmapOptio
             format!("Off-map: {}", prettyprint_usize(ppl_off_map)).draw_text(ctx),
         ])
         .centered(),
-        Widget::checkbox(ctx, "Show heatmap", None, opts.is_some()),
+        Widget::checkbox(ctx, "Run pandemic model", None, opts.pandemic),
     ];
 
-    // TODO tmp place to put pandemic model
+    if let Some(model) = pandemic {
-    if app.opts.dev {
-        // TODO Why not app.primary.current_flags.sim_flags.make_rng()? Because that'll only be the
-        // same every time this code runs (frequently, as the simulation is run) if --rng_seed is
-        // specified in the flags. If you forget it, quite confusing to see the model jump around.
-        use rand::SeedableRng;
-        use rand_xorshift::XorShiftRng;
-
-        let model = PandemicModel::calculate(
-            app.primary.sim.get_analytics(),
-            app.primary.sim.time(),
-            &mut XorShiftRng::from_seed([42; 16]),
-        );
         col.push(
             format!(
                 "Pandemic model: {} infected ({:.1}%)",
@@ -910,7 +974,13 @@ fn population_controls(ctx: &mut EventCtx, app: &App, opts: Option<&HeatmapOptio
         );
     }
 
-    if let Some(ref o) = opts {
+    col.push(Widget::checkbox(
+        ctx,
+        "Show heatmap",
+        None,
+        opts.heatmap.is_some(),
+    ));
+    if let Some(ref o) = opts.heatmap {
         // TODO Display the value...
         col.push(Widget::row(vec![
             "Resolution (meters)".draw_text(ctx).margin(5),
@@ -942,8 +1012,8 @@ fn population_controls(ctx: &mut EventCtx, app: &App, opts: Option<&HeatmapOptio
         .build(ctx)
 }
 
-fn heatmap_options(c: &mut Composite) -> Option<HeatmapOptions> {
+fn population_options(c: &mut Composite) -> PopulationOptions {
-    if c.is_checked("Show heatmap") {
+    let heatmap = if c.is_checked("Show heatmap") {
         // Did we just change?
         if c.has_widget("resolution") {
             Some(HeatmapOptions {
@@ -956,5 +1026,9 @@ fn heatmap_options(c: &mut Composite) -> Option<HeatmapOptions> {
         }
     } else {
         None
+    };
+    PopulationOptions {
+        pandemic: c.is_checked("Run pandemic model"),
+        heatmap,
     }
 }

game/src/pregame.rs

@@ -30,9 +30,9 @@ impl TitleScreen {
                 Composite::new(
                     Widget::col(vec![
                         Widget::draw_svg(ctx, "../data/system/assets/pregame/logo.svg").margin(5),
+                        // TODO that nicer font
                         // TODO Any key
-                        // TODO The hover color is wacky
+                        Btn::text_bg2("PLAY")
-                        Btn::svg_def("../data/system/assets/pregame/start.svg")
                             .build(ctx, "start game", hotkeys(vec![Key::Space, Key::Enter]))
                             .margin(5),
                     ])

sim/src/lib.rs

@@ -487,6 +487,7 @@ pub struct CreateCar {
     pub maybe_parked_car: Option<ParkedCar>,
     // None for buses
     pub trip: Option<TripID>,
+    pub person: Option<PersonID>,
 }
 
 impl CreateCar {
@@ -496,6 +497,7 @@ impl CreateCar {
         router: Router,
         req: PathRequest,
         trip: TripID,
+        person: PersonID,
     ) -> CreateCar {
         CreateCar {
             vehicle,
@@ -504,6 +506,7 @@ impl CreateCar {
             start_dist: start_pos.dist_along(),
             maybe_parked_car: None,
             trip: Some(trip),
+            person: Some(person),
         }
     }
 
@@ -514,6 +517,7 @@ impl CreateCar {
         req: PathRequest,
         start_dist: Distance,
         trip: TripID,
+        person: PersonID,
     ) -> CreateCar {
         CreateCar {
             vehicle: parked_car.vehicle.clone(),
@@ -522,6 +526,7 @@ impl CreateCar {
             start_dist,
             maybe_parked_car: Some(parked_car),
             trip: Some(trip),
+            person: Some(person),
         }
     }
 }

sim/src/make/spawner.rs

@@ -292,7 +292,9 @@ impl TripSpawner {
                         scheduler.quick_push(
                             start_time,
                             Command::SpawnCar(
-                                CreateCar::for_appearing(vehicle, start_pos, router, req, trip),
+                                CreateCar::for_appearing(
+                                    vehicle, start_pos, router, req, trip, person,
+                                ),
                                 retry_if_no_room,
                             ),
                         );

sim/src/mechanics/car.rs

@@ -1,6 +1,6 @@
 use crate::{
-    CarStatus, DistanceInterval, DrawCarInput, ParkingSpot, Router, TimeInterval, TransitSimState,
+    CarStatus, DistanceInterval, DrawCarInput, ParkingSpot, PersonID, Router, TimeInterval,
-    TripID, Vehicle, VehicleType,
+    TransitSimState, TripID, Vehicle, VehicleType,
 };
 use geom::{Distance, Duration, PolyLine, Time};
 use map_model::{Map, Traversable};
@@ -14,6 +14,7 @@ pub struct Car {
     pub router: Router,
     // None for buses
     pub trip: Option<TripID>,
+    pub person: Option<PersonID>,
     pub started_at: Time,
     pub total_blocked_time: Duration,
 

sim/src/mechanics/driving.rs

@@ -98,6 +98,7 @@ impl DrivingSimState {
                 started_at: now,
                 total_blocked_time: Duration::ZERO,
                 trip: params.trip,
+                person: params.person,
             };
             if let Some(p) = params.maybe_parked_car {
                 car.state = CarState::Unparking(
@@ -782,6 +783,7 @@ impl DrivingSimState {
                 result.push(UnzoomedAgent {
                     vehicle_type: Some(car.vehicle.vehicle_type),
                     pos: queue.id.dist_along(dist, map).0,
+                    person: car.person,
                 });
             }
         }

sim/src/mechanics/walking.rs

@@ -1,8 +1,8 @@
 use crate::{
     AgentID, Command, CreatePedestrian, DistanceInterval, DrawPedCrowdInput, DrawPedestrianInput,
     Event, IntersectionSimState, ParkingSimState, ParkingSpot, PedCrowdLocation, PedestrianID,
-    Scheduler, SidewalkPOI, SidewalkSpot, TimeInterval, TransitSimState, TripID, TripManager,
+    PersonID, Scheduler, SidewalkPOI, SidewalkSpot, TimeInterval, TransitSimState, TripID,
-    TripPositions, UnzoomedAgent,
+    TripManager, TripPositions, UnzoomedAgent,
 };
 use abstutil::{deserialize_multimap, serialize_multimap, MultiMap};
 use geom::{Distance, Duration, Line, PolyLine, Speed, Time};
@@ -64,6 +64,7 @@ impl WalkingSimState {
             path: params.path,
             goal: params.goal,
             trip: params.trip,
+            person: params.person,
         };
         ped.state = match params.start.connection {
             SidewalkPOI::Building(b) | SidewalkPOI::ParkingSpot(ParkingSpot::Offstreet(b, _)) => {
@@ -330,6 +331,7 @@ impl WalkingSimState {
             peds.push(UnzoomedAgent {
                 vehicle_type: None,
                 pos: ped.get_draw_ped(now, map).pos,
+                person: Some(ped.person),
             });
         }
 
@@ -459,6 +461,7 @@ struct Pedestrian {
     path: Path,
     goal: SidewalkSpot,
     trip: TripID,
+    person: PersonID,
 }
 
 impl Pedestrian {

sim/src/pandemic.rs

@@ -22,14 +22,14 @@ impl PandemicModel {
     //
     // This recomputes everything every time the UI asks for it. That's fine for the scale of
     // simulations now; everything else in Analytics works the same way. The faster streaming
-    // version is very straightforward -- cache this output and only process new events to update.
+    // version is very straightforward -- cache this output and only process new events.
     pub fn calculate(analytics: &Analytics, now: Time, rng: &mut XorShiftRng) -> PandemicModel {
         let mut state = PandemicModel {
             infected: BTreeSet::new(),
             bldg_occupants: BTreeMap::new(),
         };
 
-        // Now track people's movements through buildings
+        // Track people's movements through buildings
         for (time, person, bldg, left) in &analytics.building_transitions {
             if *time > now {
                 break;
@@ -49,7 +49,8 @@ impl PandemicModel {
                             true
                         }
                     });
-                // TODO Same bug as above.
+                // TODO A person left a building, but they weren't inside of it? Bug -- few
+                // possible causes...
                 if inside_since.is_none() {
                     continue;
                 }
@@ -57,7 +58,6 @@ impl PandemicModel {
 
                 // Was this person leaving infected while they were inside?
                 if !state.infected.contains(person) {
-                    //let time_in_bldg = time - inside_since.unwrap();
                     let mut longest_overlap_with_infected = Duration::ZERO;
                     for (p, t) in &state.bldg_occupants[bldg] {
                         if !state.infected.contains(p) {

sim/src/render.rs

@@ -1,4 +1,4 @@
-use crate::{CarID, PedestrianID, VehicleType};
+use crate::{CarID, PedestrianID, PersonID, VehicleType};
 use geom::{Angle, Distance, PolyLine, Pt2D, Time};
 use map_model::{BuildingID, Map, Traversable, TurnID};
 
@@ -50,6 +50,8 @@ pub struct UnzoomedAgent {
     // None means a pedestrian.
     pub vehicle_type: Option<VehicleType>,
     pub pos: Pt2D,
+    // None means a bus.
+    pub person: Option<PersonID>,
 }
 
 // TODO Can we return borrows instead? Nice for time travel, not for main sim?

sim/src/sim.rs

@@ -257,6 +257,7 @@ impl Sim {
                         router: Router::follow_bus_route(path.clone(), end_dist),
                         maybe_parked_car: None,
                         trip: None,
+                        person: None,
                     },
                     map,
                     &self.intersections,

sim/src/trips.rs

@@ -259,6 +259,7 @@ impl TripManager {
                     req,
                     start.dist_along(),
                     trip.id,
+                    trip.person,
                 ),
                 true,
             ),
@@ -313,7 +314,7 @@ impl TripManager {
         scheduler.push(
             now,
             Command::SpawnCar(
-                CreateCar::for_appearing(vehicle, driving_pos, router, req, trip.id),
+                CreateCar::for_appearing(vehicle, driving_pos, router, req, trip.id, trip.person),
                 true,
             ),
         );