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//! let (d8, d8_nd) = d8_pointer(&dem, -1.0, 8.0, 8.0);
25//! ```
26use rayon::prelude::*;
27use std::collections::{BinaryHeap, VecDeque};
28use std::cmp::Ordering;
29use std::cmp::Ordering::Equal;
30use ndarray::Array2;
31
32
33
34#[derive(PartialEq, Debug)]
35struct GridCell {
36 row: usize,
37 column: usize,
38 priority: f64,
39}
40
41impl Eq for GridCell {}
42
43impl PartialOrd for GridCell {
44 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
45 other.priority.partial_cmp(&self.priority)
46 }
47}
48
49impl Ord for GridCell {
50 fn cmp(&self, other: &Self) -> Ordering {
51 self.partial_cmp(other).unwrap()
52 }
53}
54
55#[derive(PartialEq, Debug)]
56struct GridCell2 {
57 row: usize,
58 column: usize,
59 z: f64,
60 priority: f64,
61}
62
63impl Eq for GridCell2 {}
64
65impl PartialOrd for GridCell2 {
66 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
67 other.priority.partial_cmp(&self.priority)
68 }
69}
70
71impl Ord for GridCell2 {
72 fn cmp(&self, other: &Self) -> Ordering {
73 self.partial_cmp(other).unwrap()
74 }
75}
76
77
78/// Fills depressions (sinks) in a digital elevation model (DEM).
79///
80/// More-or-less the contents of
81/// [whitebox fill_depressions](https://github.com/jblindsay/whitebox-tools/blob/master/whitebox-tools-app/src/tools/hydro_analysis/fill_depressions.rs)
82///
83/// This function modifies the input `dem` to ensure that all depressions (local minima that do not
84/// drain) are removed, making the surface hydrologically correct. It also considers no-data values
85/// and can optionally fix flat areas.
86///
87/// # Parameters
88///
89/// - `dem`: A mutable reference to a 2D array (`Array2<f64>`) representing the elevation data.
90/// - `nodata`: The value representing no-data cells in the DEM.
91/// - `resx`: The horizontal resolution (grid spacing in the x-direction).
92/// - `resy`: The vertical resolution (grid spacing in the y-direction).
93/// - `fix_flats`: A boolean flag to determine whether flat areas should be slightly sloped.
94///
95/// # Example
96///
97/// ```
98/// use ndarray::Array2;
99/// use hydro_analysis::fill_depressions;
100///
101/// let mut dem = Array2::from_shape_vec(
102/// (3, 3),
103/// vec![
104/// 10.0, 12.0, 10.0,
105/// 12.0, 9.0, 12.0,
106/// 10.0, 12.0, 10.0,
107/// ],
108/// ).expect("Failed to create DEM");
109///
110/// fill_depressions(&mut dem, -3.0, 8.0, 8.0, true);
111/// ```
112pub fn fill_depressions(
113 dem: &mut Array2<f64>, nodata: f64, resx: f64, resy: f64, fix_flats: bool
114)
115{
116 let (rows, columns) = (dem.nrows(), dem.ncols());
117 let small_num = {
118 let diagres = (resx * resx + resy * resy).sqrt();
119 let elev_digits = (dem.iter().cloned().fold(f64::NEG_INFINITY, f64::max) as i64).to_string().len();
120 let elev_multiplier = 10.0_f64.powi((9 - elev_digits) as i32);
121 1.0_f64 / elev_multiplier as f64 * diagres.ceil()
122 };
123
124 let dx = [1, 1, 1, 0, -1, -1, -1, 0];
125 let dy = [-1, 0, 1, 1, 1, 0, -1, -1];
126
127 // Find pit cells. This step is parallelizable.
128 let mut pits: Vec<_> = (1..rows - 1)
129 .into_par_iter()
130 .flat_map(|row| {
131 let mut local_pits = Vec::new();
132 for col in 1..columns - 1 {
133 let z = dem[[row, col]];
134 if z == nodata {
135 continue;
136 }
137 let mut apit = true;
138 // is anything lower than me?
139 for n in 0..8 {
140 let zn = dem[[(row as isize + dy[n]) as usize, (col as isize + dx[n]) as usize]];
141 if zn < z || zn == nodata {
142 apit = false;
143 break;
144 }
145 }
146 // no, so I am a pit
147 if apit {
148 local_pits.push((row, col, z));
149 }
150 }
151 local_pits
152 }).collect();
153
154 // Now we need to perform an in-place depression filling
155 let mut minheap = BinaryHeap::new();
156 let mut minheap2 = BinaryHeap::new();
157 let mut visited = Array2::<u8>::zeros((rows, columns));
158 let mut flats = Array2::<u8>::zeros((rows, columns));
159 let mut possible_outlets = vec![];
160 let mut queue = VecDeque::new();
161
162 // go through pits from highest to lowest
163 pits.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(Equal));
164 while let Some(cell) = pits.pop() {
165 let row: usize = cell.0;
166 let col: usize = cell.1;
167
168 // if it's already in a solved site, don't do it a second time.
169 if flats[[row, col]] != 1 {
170 // First there is a priority region-growing operation to find the outlets.
171 minheap.clear();
172 minheap.push(GridCell {
173 row: row,
174 column: col,
175 priority: dem[[row, col]],
176 });
177 visited[[row, col]] = 1;
178 let mut outlet_found = false;
179 let mut outlet_z = f64::INFINITY;
180 if !queue.is_empty() {
181 queue.clear();
182 }
183 while let Some(cell2) = minheap.pop() {
184 let z = cell2.priority;
185 if outlet_found && z > outlet_z {
186 break;
187 }
188 if !outlet_found {
189 for n in 0..8 {
190 let cn: usize = (cell2.column as isize + dx[n]) as usize;
191 let rn: usize = (cell2.row as isize + dy[n]) as usize;
192 if rn < rows && cn < columns && visited[[rn, cn]] == 0 {
193 let zn = dem[[rn, cn]];
194 if !outlet_found {
195 if zn >= z && zn != nodata {
196 minheap.push(GridCell {
197 row: rn,
198 column: cn,
199 priority: zn,
200 });
201 visited[[rn, cn]] = 1;
202 } else if zn != nodata {
203 // zn < z
204 // 'cell' has a lower neighbour that hasn't already passed through minheap.
205 // Therefore, 'cell' is a pour point cell.
206 outlet_found = true;
207 outlet_z = z;
208 queue.push_back((cell2.row, cell2.column));
209 possible_outlets.push((cell2.row, cell2.column));
210 }
211 } else if zn == outlet_z {
212 // We've found the outlet but are still looking for additional depression cells.
213 minheap.push(GridCell {
214 row: rn,
215 column: cn,
216 priority: zn,
217 });
218 visited[[rn, cn]] = 1;
219 }
220 }
221 }
222 } else {
223 // We've found the outlet but are still looking for additional depression cells and potential outlets.
224 if z == outlet_z {
225 let mut anoutlet = false;
226 for n in 0..8 {
227 let cn: usize = (cell2.column as isize + dx[n]) as usize;
228 let rn: usize = (cell2.row as isize + dy[n]) as usize;
229 if rn < rows && cn < columns && visited[[rn, cn]] == 0 {
230 let zn = dem[[rn, cn]];
231 if zn < z {
232 anoutlet = true;
233 } else if zn == outlet_z {
234 minheap.push(GridCell {
235 row: rn,
236 column: cn,
237 priority: zn,
238 });
239 visited[[rn, cn]] = 1;
240 }
241 }
242 }
243 if anoutlet {
244 queue.push_back((cell2.row, cell2.column));
245 possible_outlets.push((cell2.row, cell2.column));
246 } else {
247 visited[[cell2.row, cell2.column]] = 1;
248 }
249 }
250 }
251 }
252
253 if outlet_found {
254 // Now that we have the outlets, raise the interior of the depression.
255 // Start from the outlets.
256 while let Some(cell2) = queue.pop_front() {
257 for n in 0..8 {
258 let rn: usize = (cell2.0 as isize + dy[n]) as usize;
259 let cn: usize = (cell2.1 as isize + dx[n]) as usize;
260 if rn < rows && cn < columns && visited[[rn, cn]] == 1 {
261 visited[[rn, cn]] = 0;
262 queue.push_back((rn, cn));
263 let z = dem[[rn, cn]];
264 if z < outlet_z {
265 dem[[rn, cn]] = outlet_z;
266 flats[[rn, cn]] = 1;
267 } else if z == outlet_z {
268 flats[[rn, cn]] = 1;
269 }
270 }
271 }
272 }
273 } else {
274 queue.push_back((row, col)); // start at the pit cell and clean up visited
275 while let Some(cell2) = queue.pop_front() {
276 for n in 0..8 {
277 let rn: usize = (cell2.0 as isize + dy[n]) as usize;
278 let cn: usize = (cell2.1 as isize + dx[n]) as usize;
279 if visited[[rn, cn]] == 1 {
280 visited[[rn, cn]] = 0;
281 queue.push_back((rn, cn));
282 }
283 }
284 }
285 }
286 }
287
288 }
289
290 drop(visited);
291
292 //let (mut col, mut row): (isize, isize);
293 //let (mut rn, mut cn): (isize, isize);
294 //let (mut z, mut zn): (f64, f64);
295 //let mut flag: bool;
296
297 if small_num > 0.0 && fix_flats {
298 // Some of the potential outlets really will have lower cells.
299 minheap.clear();
300 while let Some(cell) = possible_outlets.pop() {
301 let z = dem[[cell.0, cell.1]];
302 let mut anoutlet = false;
303 for n in 0..8 {
304 let rn: usize = (cell.0 as isize + dy[n]) as usize;
305 let cn: usize = (cell.1 as isize + dx[n]) as usize;
306 if rn >= rows || cn >= columns {
307 break;
308 }
309 let zn = dem[[rn, cn]];
310 if zn < z && zn != nodata {
311 anoutlet = true;
312 break;
313 }
314 }
315 if anoutlet {
316 minheap.push(GridCell {
317 row: cell.0,
318 column: cell.1,
319 priority: z,
320 });
321 }
322 }
323
324 let mut outlets = vec![];
325 while let Some(cell) = minheap.pop() {
326 if flats[[cell.row, cell.column]] != 3 {
327 let z = dem[[cell.row, cell.column]];
328 flats[[cell.row, cell.column]] = 3;
329 if !outlets.is_empty() {
330 outlets.clear();
331 }
332 outlets.push(cell);
333 // Are there any other outlet cells at the same elevation (likely for the same feature)
334 let mut flag = true;
335 while flag {
336 match minheap.peek() {
337 Some(cell2) => {
338 if cell2.priority == z {
339 flats[[cell2.row, cell2.column]] = 3;
340 outlets
341 .push(minheap.pop().expect("Error during pop operation."));
342 } else {
343 flag = false;
344 }
345 }
346 None => {
347 flag = false;
348 }
349 }
350 }
351 if !minheap2.is_empty() {
352 minheap2.clear();
353 }
354 for cell2 in &outlets {
355 let z = dem[[cell2.row, cell2.column]];
356 for n in 0..8 {
357 let cn: usize = (cell2.column as isize + dx[n]) as usize;
358 let rn: usize = (cell2.row as isize + dy[n]) as usize;
359 if rn < rows && cn < columns && flats[[rn, cn]] != 3 {
360 let zn = dem[[rn, cn]];
361 if zn == z && zn != nodata {
362 minheap2.push(GridCell2 {
363 row: rn,
364 column: cn,
365 z: z,
366 priority: dem[[rn, cn]],
367 });
368 dem[[rn, cn]] = z + small_num;
369 flats[[rn, cn]] = 3;
370 }
371 }
372 }
373 }
374 // Now fix the flats
375 while let Some(cell2) = minheap2.pop() {
376 let z = dem[[cell2.row, cell2.column]];
377 for n in 0..8 {
378 let cn: usize = (cell2.column as isize + dx[n]) as usize;
379 let rn: usize = (cell2.row as isize + dy[n]) as usize;
380 if rn < rows && cn < columns && flats[[rn, cn]] != 3 {
381 let zn = dem[[rn, cn]];
382 if zn < z + small_num && zn >= cell2.z && zn != nodata {
383 minheap2.push(GridCell2 {
384 row: rn,
385 column: cn,
386 z: cell2.z,
387 priority: dem[[rn, cn]],
388 });
389 dem[[rn, cn]] = z + small_num;
390 flats[[rn, cn]] = 3;
391 }
392 }
393 }
394 }
395 }
396
397 }
398 }
399}
400
401/// Calculates the D8 flow direction from a digital elevation model (DEM).
402///
403/// More-or-less the contents of
404/// [whitebox d8_pointer](https://github.com/jblindsay/whitebox-tools/blob/master/whitebox-tools-app/src/tools/hydro_analysis/d8_pointer.rs)
405///
406/// This function computes the D8 flow direction for each cell in the provided DEM:
407///
408/// | . | . | . |
409/// |:--:|:---:|:--:|
410/// | 64 | 128 | 1 |
411/// | 32 | 0 | 2 |
412/// | 16 | 8 | 4 |
413///
414/// Grid cells that have no lower neighbours are assigned a flow direction of zero. In a DEM that
415/// has been pre-processed to remove all depressions and flat areas, this condition will only occur
416/// along the edges of the grid.
417///
418/// Grid cells possessing the NoData value in the input DEM are assigned the NoData value in the
419/// output image.
420///
421/// # Parameters
422/// - `dem`: A 2D array representing the digital elevation model (DEM)
423/// - `nodata`: The nodata in the DEM
424/// - `resx`: The resolution of the DEM in the x-direction in meters
425/// - `resy`: The resolution of the DEM in the y-direction in meters
426///
427/// # Returns
428/// - A tuple containing:
429/// - An `Array2<u8>` representing the D8 flow directions for each cell.
430/// - A `u8` nodata value (255)
431///
432/// # Example
433/// ```rust
434/// let dem = Array2::from_shape_vec(
435/// (3, 3),
436/// vec![
437/// 10.0, 12.0, 10.0,
438/// 12.0, 13.0, 12.0,
439/// 10.0, 12.0, 10.0,
440/// ],
441/// ).expect("Failed to create DEM");
442/// let nodata = -9999.0;
443/// let resx = 8.0;
444/// let resy = 8.0;
445/// let (d8, nd) = d8_pointer(&dem, nodata, resx, resy);
446/// ```
447pub fn d8_pointer(dem: &Array2<f64>, nodata: f64, resx: f64, resy: f64) -> (Array2<u8>, u8)
448{
449 let (nrows, ncols) = (dem.nrows(), dem.ncols());
450 let out_nodata: u8 = 255;
451 let mut d8: Array2<u8> = Array2::from_elem((nrows, ncols), out_nodata);
452
453 let diag = (resx * resx + resy * resy).sqrt();
454 let grid_lengths = [diag, resx, diag, resy, diag, resx, diag, resy];
455
456 let dx = [1, 1, 1, 0, -1, -1, -1, 0];
457 let dy = [-1, 0, 1, 1, 1, 0, -1, -1];
458
459 d8.axis_iter_mut(ndarray::Axis(0))
460 .into_par_iter()
461 .enumerate()
462 .for_each(|(row, mut d8_row)| {
463 for col in 0..ncols {
464 let z = dem[[row, col]];
465 if z == nodata {
466 continue;
467 }
468
469 let mut dir = 0;
470 let mut max_slope = f64::MIN;
471 for i in 0..8 {
472 let rn: isize = row as isize + dy[i];
473 let cn: isize = col as isize + dx[i];
474 if rn < 0 || rn >= nrows as isize || cn < 0 || cn >= ncols as isize {
475 continue;
476 }
477 let z_n = dem[[rn as usize, cn as usize]];
478 if z_n != nodata {
479 let slope = (z - z_n) / grid_lengths[i];
480 if slope > max_slope && slope > 0.0 {
481 max_slope = slope;
482 dir = i;
483 }
484 }
485 }
486
487 if max_slope >= 0.0 {
488 d8_row[col] = 1 << dir;
489 } else {
490 d8_row[col] = 0u8;
491 }
492 }
493 });
494
495 return (d8, out_nodata);
496}
497
498