diff --git a/server/src/atlas/drainage.rs b/server/src/atlas/drainage.rs index 1e198efcd..8a73ee06b 100644 --- a/server/src/atlas/drainage.rs +++ b/server/src/atlas/drainage.rs @@ -45,6 +45,13 @@ const D8: [(i32, i32); 8] = [ pub struct DrainageResult { pub river_network: RiverNetwork, pub drainage_basins: Vec, + /// Per-cell flow accumulation (row-major, `w × h`): the upstream cell count + /// draining through each cell. Exposed for D-209 attractor-strength + /// normalization (`flow_accumulation[cell] / max_accumulation`). + pub flow_accumulation: Vec, + /// Maximum flow accumulation across the grid — the denominator for + /// normalized attractor strength (D-209). Always ≥ 1. + pub max_accumulation: i32, } // --------------------------------------------------------------------------- @@ -87,9 +94,15 @@ pub fn analyze(elevation: &[f32], width: u32, height: u32, sea_level: f32) -> Dr // 8. Build DrainageBasin structs. let drainage_basins = build_basins(&labels, w, h); + // Max accumulation for D-209 strength normalization (clamped ≥ 1 so the + // division is always well-defined, even on a flat/empty world). + let max_accumulation = accum.iter().copied().max().unwrap_or(1).max(1); + DrainageResult { river_network, drainage_basins, + flow_accumulation: accum, + max_accumulation, } } @@ -378,6 +391,30 @@ fn label_basins(fdir: &[i8], accum: &[i32], w: usize, h: usize) -> Vec { // Step 7: Merge small basins // --------------------------------------------------------------------------- +/// Union-find root with path compression. +fn uf_find(parent: &mut [i32], x: i32) -> i32 { + let mut root = x; + while parent[root as usize] != root { + root = parent[root as usize]; + } + let mut cur = x; + while parent[cur as usize] != root { + let next = parent[cur as usize]; + parent[cur as usize] = root; + cur = next; + } + root +} + +/// Merge small basins into their largest neighbor until the count is in +/// `[min_count, max_count]` and every basin holds ≥ 2% of the surface. +/// +/// Builds a basin adjacency graph + sizes in a single grid pass, then performs +/// all merges as union-find operations on that graph — the grid is rewritten +/// exactly once at the end. This replaces the former O(merges × n) loop (which +/// rescanned the whole grid per merge: ~250ms at 512×256) with O(n + merges). +/// Determinism: smallest basin chosen by `(size, id)`, largest neighbor by +/// `(size, then lowest id)` — both fixed orders. fn merge_small_basins( mut labels: Vec, w: usize, @@ -385,52 +422,91 @@ fn merge_small_basins( min_count: usize, max_count: usize, ) -> Vec { + use std::collections::BTreeSet; let n = w * h; let min_frac = 0.02f64; // 2% minimum basin area - for _ in 0..200 { - // Count basin sizes. - let mut sizes: std::collections::BTreeMap = std::collections::BTreeMap::new(); - for &l in &labels { - *sizes.entry(l).or_insert(0) += 1; - } - let n_basins = sizes.len(); + let max_label = labels.iter().copied().max().unwrap_or(0); + let nb = (max_label + 1) as usize; + if nb <= 1 { + return labels; // single basin — nothing to merge + } - // Stop if within target range and all basins are large enough. - if n_basins <= max_count && sizes.values().all(|&s| s as f64 / n as f64 >= min_frac) { - break; - } - if n_basins <= min_count { - break; - } - - // Find the smallest basin. - let (&smallest_id, &smallest_size) = sizes.iter().min_by_key(|(_, &s)| s).unwrap(); - - if n_basins <= max_count && smallest_size as f64 / n as f64 >= min_frac { - break; - } - - // Find its largest adjacent basin. - let nbr_id = find_largest_neighbor(&labels, smallest_id, &sizes, w, h); - let merge_into = nbr_id.unwrap_or(0); - - // Merge. - for l in labels.iter_mut() { - if *l == smallest_id { - *l = merge_into; + // One pass: basin sizes + adjacency (neighbor labels per basin). + let mut size = vec![0usize; nb]; + let mut adj: Vec> = vec![BTreeSet::new(); nb]; + for r in 0..h { + for c in 0..w { + let l = labels[r * w + c]; + size[l as usize] += 1; + for &(dr, dc) in &D8 { + let nr = r as i32 + dr; + let nc = (c as i32 + dc).rem_euclid(w as i32) as usize; + if nr >= 0 && nr < h as i32 { + let nl = labels[nr as usize * w + nc]; + if nl != l { + adj[l as usize].insert(nl); + } + } } } } - // Renumber contiguously from 0. - let unique: Vec = { - let mut set: std::collections::BTreeSet = std::collections::BTreeSet::new(); - for &l in &labels { - set.insert(l); + let mut parent: Vec = (0..nb as i32).collect(); + let mut active: BTreeSet = (0..nb as i32).collect(); + + while active.len() > min_count { + // Smallest active basin (tie → lowest id; BTreeSet iterates ascending). + let smallest = *active.iter().min_by_key(|&&b| (size[b as usize], b)).unwrap(); + let smallest_size = size[smallest as usize]; + if active.len() <= max_count && smallest_size as f64 / n as f64 >= min_frac { + break; } - set.into_iter().collect() - }; + + // Largest active neighbor (tie → lowest id). + let mut best: i32 = -1; + let mut best_size = 0usize; + for &nb_lbl in &adj[smallest as usize] { + let rep = uf_find(&mut parent, nb_lbl); + if rep == smallest { + continue; + } + let s = size[rep as usize]; + if s > best_size || (s == best_size && (best < 0 || rep < best)) { + best_size = s; + best = rep; + } + } + // No neighbor (isolated basin) → merge into the next smallest active. + let merge_into = if best >= 0 { + best + } else { + match active.iter().find(|&&b| b != smallest) { + Some(&other) => other, + None => break, + } + }; + + // Union smallest → merge_into; fold size and adjacency. + parent[smallest as usize] = merge_into; + size[merge_into as usize] += smallest_size; + let small_adj = std::mem::take(&mut adj[smallest as usize]); + for nb_lbl in small_adj { + let rep = uf_find(&mut parent, nb_lbl); + if rep != merge_into { + adj[merge_into as usize].insert(rep); + } + } + active.remove(&smallest); + } + + // Resolve every cell to its basin representative (single pass). + for l in labels.iter_mut() { + *l = uf_find(&mut parent, *l); + } + + // Renumber contiguously from 0. + let unique: BTreeSet = labels.iter().copied().collect(); let remap: std::collections::BTreeMap = unique .iter() .enumerate() @@ -443,45 +519,6 @@ fn merge_small_basins( labels } -fn find_largest_neighbor( - labels: &[i32], - target_id: i32, - sizes: &std::collections::BTreeMap, - w: usize, - h: usize, -) -> Option { - let n = w * h; - let mut neighbor_sizes: std::collections::BTreeMap = - std::collections::BTreeMap::new(); - - for i in 0..n { - if labels[i] != target_id { - continue; - } - let r = i / w; - let c = i % w; - for &(dr, dc) in &D8 { - let nr = r as i32 + dr; - let nc = (c as i32 + dc).rem_euclid(w as i32) as usize; - if nr >= 0 && nr < h as i32 { - let nbr_id = labels[nr as usize * w + nc]; - if nbr_id != target_id { - let size = sizes.get(&nbr_id).copied().unwrap_or(0); - let e = neighbor_sizes.entry(nbr_id).or_insert(0); - if size > *e { - *e = size; - } - } - } - } - } - - neighbor_sizes - .into_iter() - .max_by_key(|(_, s)| *s) - .map(|(id, _)| id) -} - // --------------------------------------------------------------------------- // Step 8: Build DrainageBasin structs // ---------------------------------------------------------------------------