hydro_analysis/lib.rs
1//! # Hydro-analysis
2//!
3//! `hydro-analysis` provides functions for Hydrology DEM manipulation that are based on
4//! [whitebox](https://github.com/jblindsay/whitebox-tools). Whitebox is a command line tool, this
5//! crate provides some (only a couple functions at present) of that functionality via functions so
6//! can be called from your code.
7//!
8//! ## Example
9//!
10//! ```
11//! use ndarray::Array2;
12//! use hydro_analysis::fill_depressions;
13//!
14//! let mut dem = Array2::from_shape_vec(
15//! (3, 3),
16//! vec![
17//! 10.0, 12.0, 10.0,
18//! 12.0, 9.0, 12.0,
19//! 10.0, 12.0, 10.0,
20//! ],
21//! ).expect("Failed to create DEM");
22//!
23//! fill_depressions(&mut dem, -3.0, 8.0, 8.0, true);
24//! ```
25use rayon::prelude::*;
26use std::collections::{BinaryHeap, VecDeque};
27use std::cmp::Ordering;
28use std::cmp::Ordering::Equal;
29use ndarray::Array2;
30
31
32#[derive(PartialEq, Debug)]
33struct GridCell {
34 row: usize,
35 column: usize,
36 priority: f64,
37}
38
39impl Eq for GridCell {}
40
41impl PartialOrd for GridCell {
42 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
43 other.priority.partial_cmp(&self.priority)
44 }
45}
46
47impl Ord for GridCell {
48 fn cmp(&self, other: &Self) -> Ordering {
49 self.partial_cmp(other).unwrap()
50 }
51}
52
53#[derive(PartialEq, Debug)]
54struct GridCell2 {
55 row: usize,
56 column: usize,
57 z: f64,
58 priority: f64,
59}
60
61impl Eq for GridCell2 {}
62
63impl PartialOrd for GridCell2 {
64 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
65 other.priority.partial_cmp(&self.priority)
66 }
67}
68
69impl Ord for GridCell2 {
70 fn cmp(&self, other: &Self) -> Ordering {
71 self.partial_cmp(other).unwrap()
72 }
73}
74
75
76/// Fills depressions (sinks) in a digital elevation model (DEM).
77///
78/// More-or-less the contents of
79/// [whitebox fill_depressions](https://github.com/jblindsay/whitebox-tools/blob/master/whitebox-tools-app/src/tools/hydro_analysis/fill_depressions.rs)
80///
81/// This function modifies the input `dem` to ensure that all depressions (local minima that do not
82/// drain) are removed, making the surface hydrologically correct. It also considers no-data values
83/// and can optionally fix flat areas.
84///
85/// # Parameters
86///
87/// - `dem`: A mutable reference to a 2D array (`Array2<f64>`) representing the elevation data.
88/// - `nodata`: The value representing no-data cells in the DEM.
89/// - `resx`: The horizontal resolution (grid spacing in the x-direction).
90/// - `resy`: The vertical resolution (grid spacing in the y-direction).
91/// - `fix_flats`: A boolean flag to determine whether flat areas should be slightly sloped.
92///
93/// # Example
94///
95/// ```
96/// use ndarray::Array2;
97/// use hydro_analysis::fill_depressions;
98///
99/// let mut dem = Array2::from_shape_vec(
100/// (3, 3),
101/// vec![
102/// 10.0, 12.0, 10.0,
103/// 12.0, 9.0, 12.0,
104/// 10.0, 12.0, 10.0,
105/// ],
106/// ).expect("Failed to create DEM");
107///
108/// fill_depressions(&mut dem, -3.0, 8.0, 8.0, true);
109/// ```
110pub fn fill_depressions(
111 dem: &mut Array2<f64>, nodata: f64, resx: f64, resy: f64, fix_flats: bool
112)
113{
114 let (rows, columns) = (dem.nrows(), dem.ncols());
115 let small_num = {
116 let diagres = (resx * resx + resy * resy).sqrt();
117 let elev_digits = (dem.iter().cloned().fold(f64::NEG_INFINITY, f64::max) as i64).to_string().len();
118 let elev_multiplier = 10.0_f64.powi((9 - elev_digits) as i32);
119 1.0_f64 / elev_multiplier as f64 * diagres.ceil()
120 };
121
122 let dx = [1, 1, 1, 0, -1, -1, -1, 0];
123 let dy = [-1, 0, 1, 1, 1, 0, -1, -1];
124
125 // Find pit cells. This step is parallelizable.
126 let mut pits: Vec<_> = (1..rows - 1)
127 .into_par_iter()
128 .flat_map(|row| {
129 let mut local_pits = Vec::new();
130 for col in 1..columns - 1 {
131 let z = dem[[row, col]];
132 if z == nodata {
133 continue;
134 }
135 let mut apit = true;
136 // is anything lower than me?
137 for n in 0..8 {
138 let zn = dem[[(row as isize + dy[n]) as usize, (col as isize + dx[n]) as usize]];
139 if zn < z || zn == nodata {
140 apit = false;
141 break;
142 }
143 }
144 // no, so I am a pit
145 if apit {
146 local_pits.push((row, col, z));
147 }
148 }
149 local_pits
150 }).collect();
151
152 // Now we need to perform an in-place depression filling
153 let mut minheap = BinaryHeap::new();
154 let mut minheap2 = BinaryHeap::new();
155 let mut visited = Array2::<u8>::zeros((rows, columns));
156 let mut flats = Array2::<u8>::zeros((rows, columns));
157 let mut possible_outlets = vec![];
158 let mut queue = VecDeque::new();
159
160 // go through pits from highest to lowest
161 pits.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(Equal));
162 while let Some(cell) = pits.pop() {
163 let row: usize = cell.0;
164 let col: usize = cell.1;
165
166 // if it's already in a solved site, don't do it a second time.
167 if flats[[row, col]] != 1 {
168 // First there is a priority region-growing operation to find the outlets.
169 minheap.clear();
170 minheap.push(GridCell {
171 row: row,
172 column: col,
173 priority: dem[[row, col]],
174 });
175 visited[[row, col]] = 1;
176 let mut outlet_found = false;
177 let mut outlet_z = f64::INFINITY;
178 if !queue.is_empty() {
179 queue.clear();
180 }
181 while let Some(cell2) = minheap.pop() {
182 let z = cell2.priority;
183 if outlet_found && z > outlet_z {
184 break;
185 }
186 if !outlet_found {
187 for n in 0..8 {
188 let cn: usize = (cell2.column as isize + dx[n]) as usize;
189 let rn: usize = (cell2.row as isize + dy[n]) as usize;
190 if rn < rows && cn < columns && visited[[rn, cn]] == 0 {
191 let zn = dem[[rn, cn]];
192 if !outlet_found {
193 if zn >= z && zn != nodata {
194 minheap.push(GridCell {
195 row: rn,
196 column: cn,
197 priority: zn,
198 });
199 visited[[rn, cn]] = 1;
200 } else if zn != nodata {
201 // zn < z
202 // 'cell' has a lower neighbour that hasn't already passed through minheap.
203 // Therefore, 'cell' is a pour point cell.
204 outlet_found = true;
205 outlet_z = z;
206 queue.push_back((cell2.row, cell2.column));
207 possible_outlets.push((cell2.row, cell2.column));
208 }
209 } else if zn == outlet_z {
210 // We've found the outlet but are still looking for additional depression cells.
211 minheap.push(GridCell {
212 row: rn,
213 column: cn,
214 priority: zn,
215 });
216 visited[[rn, cn]] = 1;
217 }
218 }
219 }
220 } else {
221 // We've found the outlet but are still looking for additional depression cells and potential outlets.
222 if z == outlet_z {
223 let mut anoutlet = false;
224 for n in 0..8 {
225 let cn: usize = (cell2.column as isize + dx[n]) as usize;
226 let rn: usize = (cell2.row as isize + dy[n]) as usize;
227 if rn < rows && cn < columns && visited[[rn, cn]] == 0 {
228 let zn = dem[[rn, cn]];
229 if zn < z {
230 anoutlet = true;
231 } else if zn == outlet_z {
232 minheap.push(GridCell {
233 row: rn,
234 column: cn,
235 priority: zn,
236 });
237 visited[[rn, cn]] = 1;
238 }
239 }
240 }
241 if anoutlet {
242 queue.push_back((cell2.row, cell2.column));
243 possible_outlets.push((cell2.row, cell2.column));
244 } else {
245 visited[[cell2.row, cell2.column]] = 1;
246 }
247 }
248 }
249 }
250
251 if outlet_found {
252 // Now that we have the outlets, raise the interior of the depression.
253 // Start from the outlets.
254 while let Some(cell2) = queue.pop_front() {
255 for n in 0..8 {
256 let rn: usize = (cell2.0 as isize + dy[n]) as usize;
257 let cn: usize = (cell2.1 as isize + dx[n]) as usize;
258 if rn < rows && cn < columns && visited[[rn, cn]] == 1 {
259 visited[[rn, cn]] = 0;
260 queue.push_back((rn, cn));
261 let z = dem[[rn, cn]];
262 if z < outlet_z {
263 dem[[rn, cn]] = outlet_z;
264 flats[[rn, cn]] = 1;
265 } else if z == outlet_z {
266 flats[[rn, cn]] = 1;
267 }
268 }
269 }
270 }
271 } else {
272 queue.push_back((row, col)); // start at the pit cell and clean up visited
273 while let Some(cell2) = queue.pop_front() {
274 for n in 0..8 {
275 let rn: usize = (cell2.0 as isize + dy[n]) as usize;
276 let cn: usize = (cell2.1 as isize + dx[n]) as usize;
277 if visited[[rn, cn]] == 1 {
278 visited[[rn, cn]] = 0;
279 queue.push_back((rn, cn));
280 }
281 }
282 }
283 }
284 }
285
286 }
287
288 drop(visited);
289
290 //let (mut col, mut row): (isize, isize);
291 //let (mut rn, mut cn): (isize, isize);
292 //let (mut z, mut zn): (f64, f64);
293 //let mut flag: bool;
294
295 if small_num > 0.0 && fix_flats {
296 // Some of the potential outlets really will have lower cells.
297 minheap.clear();
298 while let Some(cell) = possible_outlets.pop() {
299 let z = dem[[cell.0, cell.1]];
300 let mut anoutlet = false;
301 for n in 0..8 {
302 let rn: usize = (cell.0 as isize + dy[n]) as usize;
303 let cn: usize = (cell.1 as isize + dx[n]) as usize;
304 if rn >= rows || cn >= columns {
305 break;
306 }
307 let zn = dem[[rn, cn]];
308 if zn < z && zn != nodata {
309 anoutlet = true;
310 break;
311 }
312 }
313 if anoutlet {
314 minheap.push(GridCell {
315 row: cell.0,
316 column: cell.1,
317 priority: z,
318 });
319 }
320 }
321
322 let mut outlets = vec![];
323 while let Some(cell) = minheap.pop() {
324 if flats[[cell.row, cell.column]] != 3 {
325 let z = dem[[cell.row, cell.column]];
326 flats[[cell.row, cell.column]] = 3;
327 if !outlets.is_empty() {
328 outlets.clear();
329 }
330 outlets.push(cell);
331 // Are there any other outlet cells at the same elevation (likely for the same feature)
332 let mut flag = true;
333 while flag {
334 match minheap.peek() {
335 Some(cell2) => {
336 if cell2.priority == z {
337 flats[[cell2.row, cell2.column]] = 3;
338 outlets
339 .push(minheap.pop().expect("Error during pop operation."));
340 } else {
341 flag = false;
342 }
343 }
344 None => {
345 flag = false;
346 }
347 }
348 }
349 if !minheap2.is_empty() {
350 minheap2.clear();
351 }
352 for cell2 in &outlets {
353 let z = dem[[cell2.row, cell2.column]];
354 for n in 0..8 {
355 let cn: usize = (cell2.column as isize + dx[n]) as usize;
356 let rn: usize = (cell2.row as isize + dy[n]) as usize;
357 if rn < rows && cn < columns && flats[[rn, cn]] != 3 {
358 let zn = dem[[rn, cn]];
359 if zn == z && zn != nodata {
360 minheap2.push(GridCell2 {
361 row: rn,
362 column: cn,
363 z: z,
364 priority: dem[[rn, cn]],
365 });
366 dem[[rn, cn]] = z + small_num;
367 flats[[rn, cn]] = 3;
368 }
369 }
370 }
371 }
372 // Now fix the flats
373 while let Some(cell2) = minheap2.pop() {
374 let z = dem[[cell2.row, cell2.column]];
375 for n in 0..8 {
376 let cn: usize = (cell2.column as isize + dx[n]) as usize;
377 let rn: usize = (cell2.row as isize + dy[n]) as usize;
378 if rn < rows && cn < columns && flats[[rn, cn]] != 3 {
379 let zn = dem[[rn, cn]];
380 if zn < z + small_num && zn >= cell2.z && zn != nodata {
381 minheap2.push(GridCell2 {
382 row: rn,
383 column: cn,
384 z: cell2.z,
385 priority: dem[[rn, cn]],
386 });
387 dem[[rn, cn]] = z + small_num;
388 flats[[rn, cn]] = 3;
389 }
390 }
391 }
392 }
393 }
394
395 }
396 }
397}
398
399
400