Files
settled-reach/server/src/atlas/drainage.rs
T
2026-05-03 09:46:16 +02:00

650 lines
20 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! D8 drainage routing — flow direction, flow accumulation, river network
//! extraction, and drainage basin delineation (D-208).
//!
//! **Determinism (D-010, D-208):** All flow-direction comparisons use integer
//! arithmetic on scaled elevation values (`(elev * 1_000_000.0) as i64`) to
//! avoid f32 comparison non-determinism. Tie-breaking uses a fixed D8 neighbor
//! priority order. The result is bit-identical across runs on the same inputs.
//!
//! **Algorithm:**
//! 1. Scale f32 elevation to i64 integers.
//! 2. Priority-flood depression fill (iterative, convergence in ≤10 passes).
//! 3. D8 flow direction: steepest descent, 8-neighbor, wraps horizontally.
//! 4. Flow accumulation via topological sort of the D8 DAG.
//! 5. River network extraction: cells with accumulation > RIVER_THRESHOLD.
//! 6. Basin labeling: flood-fill seeded at pour points.
//!
//! The grid is row-major. Row 0 is the north pole; row H-1 is the south pole.
//! Columns wrap horizontally (the globe is equirectangular).
use std::collections::VecDeque;
use crate::atlas::body_world_state::{DrainageBasin, RiverNetwork};
/// A cell is a river cell when its flow accumulation exceeds this threshold (D-208).
pub const RIVER_THRESHOLD: i32 = 200;
/// Scale factor for converting f32 elevation to integer for deterministic comparison.
const ELEV_SCALE: f64 = 1_000_000.0;
// D8 neighbor offsets (dr, dc) in fixed priority order for deterministic tie-breaking.
// Priority: cardinal directions first (N, S, E, W), then diagonals (NE, NW, SE, SW).
const D8: [(i32, i32); 8] = [
(-1, 0), // N
(1, 0), // S
(0, 1), // E
(0, -1), // W
(-1, 1), // NE
(-1, -1), // NW
(1, 1), // SE
(1, -1), // SW
];
/// Result of the full D8 drainage analysis for one body.
#[derive(Debug, Clone)]
pub struct DrainageResult {
pub river_network: RiverNetwork,
pub drainage_basins: Vec<DrainageBasin>,
}
// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------
/// Run the full D8 drainage analysis on an elevation grid.
///
/// `elevation` is a row-major float32 grid of shape `height × width`, values
/// in [0.0, 1.0]. `sea_level` is the fraction below which terrain is ocean.
///
/// Returns `DrainageResult` with the river network and drainage basins.
pub fn analyze(elevation: &[f32], width: u32, height: u32, sea_level: f32) -> DrainageResult {
let w = width as usize;
let h = height as usize;
// 1. Scale to integers.
let scaled: Vec<i64> = elevation
.iter()
.map(|&e| (e as f64 * ELEV_SCALE) as i64)
.collect();
// 2. Depression fill.
let filled = depression_fill(&scaled, w, h);
// 3. D8 flow direction. -1 = no outflow (edge or flat peak).
let fdir = flow_direction(&filled, w, h);
// 4. Flow accumulation.
let accum = flow_accumulation(&fdir, w, h);
// 5. River network.
let river_network = extract_river_network(&accum, &fdir, w, h, sea_level, elevation);
// 6. Basin labeling.
let labels = label_basins(&fdir, &accum, w, h);
// 7. Merge small basins + clamp count to [4, 12].
let labels = merge_small_basins(labels, w, h, 4, 12);
// 8. Build DrainageBasin structs.
let drainage_basins = build_basins(&labels, w, h);
DrainageResult {
river_network,
drainage_basins,
}
}
// ---------------------------------------------------------------------------
// Step 2: Depression fill
// ---------------------------------------------------------------------------
fn depression_fill(scaled: &[i64], w: usize, h: usize) -> Vec<i64> {
let mut filled = scaled.to_vec();
for _ in 0..10 {
let mut changed = false;
for r in 1..h.saturating_sub(1) {
for c in 0..w {
let mut nbr_min = i64::MAX;
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 val = filled[nr as usize * w + nc];
if val < nbr_min {
nbr_min = val;
}
}
}
if filled[r * w + c] < nbr_min {
filled[r * w + c] = nbr_min + 1;
changed = true;
}
}
}
if !changed {
break;
}
}
filled
}
// ---------------------------------------------------------------------------
// Step 3: D8 flow direction
// ---------------------------------------------------------------------------
/// Returns per-cell flow direction index into D8 (07), or -1 for no outflow.
fn flow_direction(filled: &[i64], w: usize, h: usize) -> Vec<i8> {
let mut fdir = vec![-1i8; w * h];
for r in 0..h {
for c in 0..w {
let elev = filled[r * w + c];
let mut best_drop = 0i64;
let mut best_k: i8 = -1;
for (k, &(dr, dc)) in D8.iter().enumerate() {
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 {
continue;
}
let drop = elev - filled[nr as usize * w + nc];
if drop > best_drop {
best_drop = drop;
best_k = k as i8;
}
}
fdir[r * w + c] = best_k;
}
}
fdir
}
// ---------------------------------------------------------------------------
// Step 4: Flow accumulation
// ---------------------------------------------------------------------------
fn flow_accumulation(fdir: &[i8], w: usize, h: usize) -> Vec<i32> {
let n = w * h;
let mut in_degree = vec![0i32; n];
for r in 0..h {
for c in 0..w {
let k = fdir[r * w + c];
if k < 0 {
continue;
}
let (dr, dc) = D8[k as usize];
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 {
in_degree[nr as usize * w + nc] += 1;
}
}
}
let mut queue = VecDeque::new();
for (i, &deg) in in_degree.iter().enumerate().take(n) {
if deg == 0 {
queue.push_back(i);
}
}
let mut accum = vec![1i32; n];
while let Some(idx) = queue.pop_front() {
let r = idx / w;
let c = idx % w;
let k = fdir[idx];
if k < 0 {
continue;
}
let (dr, dc) = D8[k as usize];
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 ni = nr as usize * w + nc;
accum[ni] += accum[idx];
in_degree[ni] -= 1;
if in_degree[ni] == 0 {
queue.push_back(ni);
}
}
}
accum
}
// ---------------------------------------------------------------------------
// Step 5: River network extraction
// ---------------------------------------------------------------------------
fn extract_river_network(
accum: &[i32],
fdir: &[i8],
w: usize,
h: usize,
sea_level: f32,
elevation: &[f32],
) -> RiverNetwork {
let n = w * h;
// River cells: above threshold AND above sea level.
let is_river: Vec<bool> = (0..n)
.map(|i| accum[i] > RIVER_THRESHOLD && elevation[i] >= sea_level)
.collect();
let river_cells: Vec<(u16, u16)> = (0..n)
.filter(|&i| is_river[i])
.map(|i| ((i / w) as u16, (i % w) as u16))
.collect();
// Confluences: river cells with 2+ river neighbors flowing into them.
let mut inflow_count = vec![0u8; n];
for r in 0..h {
for c in 0..w {
let i = r * w + c;
if !is_river[i] {
continue;
}
let k = fdir[i];
if k < 0 {
continue;
}
let (dr, dc) = D8[k as usize];
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 ni = nr as usize * w + nc;
if is_river[ni] {
inflow_count[ni] = inflow_count[ni].saturating_add(1);
}
}
}
}
let confluences: Vec<(u16, u16)> = (0..n)
.filter(|&i| is_river[i] && inflow_count[i] >= 2)
.map(|i| ((i / w) as u16, (i % w) as u16))
.collect();
// Mouths: river cells that flow to a sea cell or to the polar edge.
let mouths: Vec<(u16, u16)> = (0..n)
.filter(|&i| {
if !is_river[i] {
return false;
}
let r = i / w;
let c = i % w;
let k = fdir[i];
if k < 0 {
return true; // no outflow — edge
}
let (dr, dc) = D8[k as usize];
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 {
return true; // polar edge
}
// Flows into a sub-sea-level cell = mouth
elevation[nr as usize * w + nc] < sea_level
})
.map(|i| ((i / w) as u16, (i % w) as u16))
.collect();
RiverNetwork {
river_cells,
confluences,
mouths,
}
}
// ---------------------------------------------------------------------------
// Step 6: Basin labeling
// ---------------------------------------------------------------------------
fn label_basins(fdir: &[i8], accum: &[i32], w: usize, h: usize) -> Vec<i32> {
let n = w * h;
let mut labels = vec![-1i32; n];
// Pour points: local accumulation maxima above river threshold.
let mut pour_pts: Vec<usize> = Vec::new();
for i in 0..n {
if accum[i] <= RIVER_THRESHOLD {
continue;
}
let r = i / w;
let c = i % w;
let mut is_max = true;
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 && accum[nr as usize * w + nc] > accum[i] {
is_max = false;
break;
}
}
if is_max {
pour_pts.push(i);
}
}
if pour_pts.is_empty() {
// Flat/ocean world — single basin.
labels.iter_mut().for_each(|l| *l = 0);
return labels;
}
for (basin_id, &idx) in pour_pts.iter().enumerate() {
labels[idx] = basin_id as i32;
}
// Trace remaining cells: follow fdir until a labeled cell is reached.
for start in 0..n {
if labels[start] >= 0 {
continue;
}
// Walk forward, accumulate path.
let mut path: Vec<usize> = Vec::new();
let mut cur = start;
let label = loop {
if labels[cur] >= 0 {
break labels[cur];
}
path.push(cur);
let k = fdir[cur];
if k < 0 {
break 0; // no outflow — assign to basin 0
}
let r = cur / w;
let c = cur % w;
let (dr, dc) = D8[k as usize];
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 {
break 0; // polar edge
}
let next = nr as usize * w + nc;
// Cycle guard: if we're visiting a cell already in path, stop.
if path.contains(&next) {
break 0;
}
cur = next;
};
for idx in path {
labels[idx] = label;
}
}
labels
}
// ---------------------------------------------------------------------------
// Step 7: Merge small basins
// ---------------------------------------------------------------------------
fn merge_small_basins(
mut labels: Vec<i32>,
w: usize,
h: usize,
min_count: usize,
max_count: usize,
) -> Vec<i32> {
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<i32, usize> = std::collections::BTreeMap::new();
for &l in &labels {
*sizes.entry(l).or_insert(0) += 1;
}
let n_basins = sizes.len();
// 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;
}
}
}
// Renumber contiguously from 0.
let unique: Vec<i32> = {
let mut set: std::collections::BTreeSet<i32> = std::collections::BTreeSet::new();
for &l in &labels {
set.insert(l);
}
set.into_iter().collect()
};
let remap: std::collections::BTreeMap<i32, i32> = unique
.iter()
.enumerate()
.map(|(new, &old)| (old, new as i32))
.collect();
for l in labels.iter_mut() {
*l = remap[l];
}
labels
}
fn find_largest_neighbor(
labels: &[i32],
target_id: i32,
sizes: &std::collections::BTreeMap<i32, usize>,
w: usize,
h: usize,
) -> Option<i32> {
let n = w * h;
let mut neighbor_sizes: std::collections::BTreeMap<i32, usize> =
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
// ---------------------------------------------------------------------------
fn build_basins(labels: &[i32], w: usize, h: usize) -> Vec<DrainageBasin> {
let n = w * h;
let mut basin_map: std::collections::BTreeMap<i32, Vec<usize>> =
std::collections::BTreeMap::new();
for (i, &l) in labels.iter().enumerate() {
basin_map.entry(l).or_default().push(i);
}
let mut basins: Vec<DrainageBasin> = Vec::with_capacity(basin_map.len());
let mut ids: Vec<i32> = basin_map.keys().copied().collect();
ids.sort();
for basin_id in ids {
let cells = &basin_map[&basin_id];
let area_pct = cells.len() as f32 / n as f32;
// Boundary cells: in this basin, adjacent to a different basin or edge.
let mut boundary: Vec<(u16, u16)> = Vec::new();
for &idx in cells {
let r = idx / w;
let c = idx % w;
let mut on_boundary = false;
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 {
on_boundary = true;
break;
}
if labels[nr as usize * w + nc] != basin_id {
on_boundary = true;
break;
}
}
if on_boundary {
boundary.push((r as u16, c as u16));
}
}
// Sort boundary by angle from centroid for a coherent polygon.
if !boundary.is_empty() {
let cr = boundary.iter().map(|&(r, _)| r as f32).sum::<f32>() / boundary.len() as f32;
let cc = boundary.iter().map(|&(_, c)| c as f32).sum::<f32>() / boundary.len() as f32;
boundary.sort_by(|&(r1, c1), &(r2, c2)| {
let a1 = (r1 as f32 - cr).atan2(c1 as f32 - cc);
let a2 = (r2 as f32 - cr).atan2(c2 as f32 - cc);
a1.partial_cmp(&a2).unwrap_or(std::cmp::Ordering::Equal)
});
// Subsample to ≤500 points.
if boundary.len() > 500 {
let step = boundary.len() / 500;
boundary = boundary.into_iter().step_by(step).collect();
}
}
basins.push(DrainageBasin {
basin_id: basin_id as u32,
boundary,
area_pct,
});
}
basins
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
fn flat_grid(w: u32, h: u32, val: f32) -> Vec<f32> {
vec![val; (w * h) as usize]
}
fn slope_grid(w: u32, h: u32) -> Vec<f32> {
let n = (w * h) as usize;
(0..n)
.map(|i| {
let r = i / w as usize;
let c = i % w as usize;
// Slope: higher in top-left, drains toward bottom-right.
1.0 - (r as f32 / h as f32 * 0.5 + c as f32 / w as f32 * 0.5)
})
.collect()
}
#[test]
fn flat_grid_produces_single_basin() {
let elev = flat_grid(16, 8, 0.5);
let result = analyze(&elev, 16, 8, 0.3);
// Flat world → no pour points → single basin
assert_eq!(result.drainage_basins.len(), 1);
assert!((result.drainage_basins[0].area_pct - 1.0).abs() < 0.01);
}
#[test]
fn slope_grid_has_no_river_cells_below_threshold_by_default() {
// Small 8×4 grid: max flow_accum ≤ 32, below RIVER_THRESHOLD (200).
let elev = slope_grid(8, 4);
let result = analyze(&elev, 8, 4, 0.3);
// River cells may be empty on this tiny grid — that is acceptable.
// What matters: no panic and basin count ≥ 1.
assert!(!result.drainage_basins.is_empty());
}
#[test]
fn large_grid_river_cells_nonempty() {
// 512×256: max flow accumulation ~131K >> RIVER_THRESHOLD.
let elev = slope_grid(512, 256);
let result = analyze(&elev, 512, 256, 0.3);
assert!(
!result.river_network.river_cells.is_empty(),
"Expected river cells on a large sloped grid"
);
}
#[test]
fn basin_area_pcts_sum_to_one() {
let elev = slope_grid(64, 32);
let result = analyze(&elev, 64, 32, 0.3);
let total: f32 = result.drainage_basins.iter().map(|b| b.area_pct).sum();
assert!(
(total - 1.0).abs() < 0.01,
"Basin area fractions must sum to 1, got {}",
total
);
}
#[test]
fn basin_count_within_target_range() {
let elev = slope_grid(128, 64);
let result = analyze(&elev, 128, 64, 0.3);
let n = result.drainage_basins.len();
assert!(
n >= 1 && n <= 12,
"Basin count {} out of expected range [1, 12]",
n
);
}
#[test]
fn determinism() {
// Running analyze twice on the same input must produce identical results.
let elev = slope_grid(64, 32);
let r1 = analyze(&elev, 64, 32, 0.3);
let r2 = analyze(&elev, 64, 32, 0.3);
assert_eq!(
r1.river_network.river_cells, r2.river_network.river_cells,
"River cells must be deterministic"
);
assert_eq!(
r1.drainage_basins.len(),
r2.drainage_basins.len(),
"Basin count must be deterministic"
);
}
}