Use the weights when picking points in a zone

Urban-Analytics-Technology-Platform · Sep 12, 2024 · 986fd18 · 986fd18
1 parent 343d3a6
commit 986fd18
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 37 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/od2net/Cargo.toml b/od2net/Cargo.toml
@@ -19,8 +19,8 @@ itertools = "0.12.1"
 log = "0.4.21"
 lts = { path = "../lts" }
 osm-reader = { git = "https://github.com/a-b-street/osm-reader", features = ["serde"] }
-nanorand = { version = "0.7.0", default-features = false, features = ["wyrand"] }
 num_cpus = "1.16.0"
+rand = { version = "0.8.5", default-features = false, features = ["alloc", "std_rng"] }
 rayon = "1.9.0"
 rstar = "0.12.0"
 serde = { version = "1.0.197", features = ["derive"] }

diff --git a/od2net/src/od.rs b/od2net/src/od.rs
@@ -6,7 +6,7 @@ use fs_err::File;
 use geo::{BoundingRect, Centroid, Contains, MultiPolygon, Point};
 use geojson::{FeatureReader, Value};
 use indicatif::HumanCount;
-use nanorand::{Rng, WyRand};
+use rand::{prelude::SliceRandom, rngs::StdRng, SeedableRng};
 use rstar::{PointDistance, RTree, RTreeObject, AABB};
 use serde::Deserialize;
 
@@ -104,23 +104,23 @@ pub fn generate_requests(
             timer.stop();
 
             timer.start(format!("Generating requests from {csv_path}"));
-            let mut rng = WyRand::new_seed(rng_seed);
+            let mut rng = StdRng::seed_from_u64(rng_seed);
 
             for rec in csv::Reader::from_reader(File::open(csv_path)?).deserialize() {
                 let row: BetweenZonesRow = rec?;
+                let Some(from_points) = origins_per_zone.get(&row.from) else {
+                    bail!("Unknown zone {}", row.from);
+                };
+                let Some(to_points) = destinations_per_zone.get(&row.to) else {
+                    bail!("Unknown zone {}", row.to);
+                };
                 for _ in 0..row.count {
-                    let from = match origins_per_zone.get(&row.from) {
-                        Some(points) => points[rng.generate_range(0..points.len())],
-                        None => {
-                            bail!("Unknown zone {}", row.from);
-                        }
-                    };
-                    let to = match destinations_per_zone.get(&row.to) {
-                        Some(points) => points[rng.generate_range(0..points.len())],
-                        None => {
-                            bail!("Unknown zone {}", row.to);
-                        }
-                    };
+                    // TODO choose_weighted is O(n); there are alternatives if this ever becomes a
+                    // problem.
+                    let from = from_points
+                        .choose_weighted(&mut rng, |pt| pt.weight)
+                        .unwrap();
+                    let to = to_points.choose_weighted(&mut rng, |pt| pt.weight).unwrap();
                     requests.push(Request {
                         x1: from.lon,
                         y1: from.lat,
@@ -154,23 +154,20 @@ pub fn generate_requests(
             timer.stop();
 
             timer.start(format!("Generating requests from {csv_path}"));
-            let mut rng = WyRand::new_seed(rng_seed);
+            let mut rng = StdRng::seed_from_u64(rng_seed);
 
             for rec in csv::Reader::from_reader(File::open(csv_path)?).deserialize() {
                 let row: BetweenZonesRow = rec?;
+                let Some(from_points) = origins_per_zone.get(&row.from) else {
+                    bail!("Unknown zone {}", row.from);
+                };
+                let Some(to) = destinations.get(&row.to) else {
+                    bail!("Unknown destination {}", row.to);
+                };
                 for _ in 0..row.count {
-                    let from = match origins_per_zone.get(&row.from) {
-                        Some(points) => points[rng.generate_range(0..points.len())],
-                        None => {
-                            bail!("Unknown zone {}", row.from);
-                        }
-                    };
-                    let to = match destinations.get(&row.to) {
-                        Some(pt) => *pt,
-                        None => {
-                            bail!("Unknown destination {}", row.to);
-                        }
-                    };
+                    let from = from_points
+                        .choose_weighted(&mut rng, |pt| pt.weight)
+                        .unwrap();
                     requests.push(Request {
                         x1: from.lon,
                         y1: from.lat,