hydro_analysis/lib.rs
1//! # Hydro-analysis
2//!
3//! `hydro-analysis` provides functions for Hydrology DEM manipulation. There are
4//! a couple generic functions for reading/writing raster files of any common
5//! primative type (which surprizingly I couldn't find anywhere else, unless you
6//! use GDAL which I am trying to avoid). Also there are a couple functions based
7//! on [whitebox](https://github.com/jblindsay/whitebox-tools). Whitebox is a
8//! command line tool, this provides functionality via functions so can be called
9//! from your code.
10//!
11//! ## Example of reading and writing rasters
12//!
13//! ```
14//! let ifn = PathBuf::from("input.tif");
15//! let (d8, nd, crs, geo, gdir, proj) = rasterfile_to_array::<u8>(&ifn)?;
16//! /* do something with d8, or make a new array2 */
17//! let ofn = PathBuf::from("output.tif");
18//! if let Err(e) = array_to_rasterfile::<u8>(&d8, nd, &geo, &gdir, &proj, &ofn) {
19//! eprintln!("Error occured while writing {}: {:?}", ofn.display(), e);
20//! }
21//! ```
22//!
23//! ## Example of filling and d8
24//!
25//! ```
26//! use ndarray::Array2;
27//! use hydro_analysis::{fill_depressions, d8_pointer};
28//!
29//! let mut dem = Array2::from_shape_vec(
30//! (3, 3),
31//! vec![
32//! 10.0, 12.0, 10.0,
33//! 12.0, 9.0, 12.0,
34//! 10.0, 12.0, 10.0,
35//! ],
36//! ).expect("Failed to create DEM");
37//!
38//! fill_depressions(&mut dem, -3.0, 8.0, 8.0, true);
39//! let (d8, d8_nd) = d8_pointer(&dem, -1.0, 8.0, 8.0);
40//! ```
41
42use rayon::prelude::*;
43use std::collections::{BinaryHeap, VecDeque};
44use std::cmp::Ordering;
45use std::cmp::Ordering::Equal;
46use ndarray::Array2;
47
48use std::{fs::File, f64, path::PathBuf};
49use thiserror::Error;
50use bytemuck::cast_slice;
51
52use tiff::decoder::DecodingResult;
53use tiff::encoder::compression::Deflate;
54use tiff::encoder::colortype::{Gray8,Gray16,Gray32,Gray64,Gray32Float,Gray64Float,GrayI8,GrayI16,GrayI32,GrayI64};
55use tiff::tags::Tag;
56use tiff::TiffFormatError;
57
58
59#[derive(Debug, Error)]
60pub enum RasterError {
61 #[error("TIFF error: {0}")]
62 Tiff(#[from] tiff::TiffError),
63
64 #[error("I/O error: {0}")]
65 Io(#[from] std::io::Error),
66
67 #[error("NDarray: {0}")]
68 Shape(#[from] ndarray::ShapeError),
69
70 #[error("Failed to parse nodata value")]
71 ParseIntError(#[from] std::num::ParseIntError),
72
73 #[error("Failed to parse nodata value")]
74 ParseFloatError(#[from] std::num::ParseFloatError),
75
76 #[error("Unsupported type: {0}")]
77 UnsupportedType(String)
78}
79
80/// Reads a single-band grayscale GeoTIFF raster file returning ndarry2 and metadata
81///
82/// # Type Parameters
83/// - `T`: The pixel value type. u8, u16, i16, f64 etc
84///
85/// # Parameters
86/// - `fname`: Path to the input `.tif` GeoTIFF file.
87///
88/// # Returns
89/// A `Result` with a tuple containing:
90/// - `Array2<T>`: The raster data in a 2D array.
91/// - `T`: nodata
92/// - `u16`: CRS (e.g. 2193)
93/// - `[f64; 6]`: The affine GeoTransform in the format:
94/// `[origin_x, pixel_size_x, rotation_x, origin_y, rotation_y, pixel_size_y]`.
95/// - `Vec<u64>`: raw GeoKeyDirectoryTag values (needed for writing to file)
96/// - `String`: PROJ string (needed for writing to file)
97///
98/// # Errors
99/// - Returns `RasterError` variants if reading fails, the type conversion for data or metadata
100/// fails, or required tags are missing from the TIFF file.
101///
102/// # Example
103/// ```
104/// let path = PathBuf::from("input.tif");
105/// let (d8, nd, crs, geo, gdir, proj) = rasterfile_to_array::<u8>(&path)?;
106/// ```
107pub fn rasterfile_to_array<T>(fname: &PathBuf) -> Result<
108 (
109 Array2<T>,
110 T, // nodata
111 u16, // crs
112 [f64; 6], // geo transform [start_x, psize_x, rotation, starty, rotation, psize_y]
113 Vec<u64>, // geo dir, it has the crs in it
114 String // the projection string
115 ),
116 RasterError
117>
118 where T: std::str::FromStr + num::FromPrimitive,
119 <T as std::str::FromStr>::Err: std::fmt::Debug,
120 RasterError: std::convert::From<<T as std::str::FromStr>::Err>
121{
122 // Open the file
123 let file = File::open(fname)?;
124
125 // Create a TIFF decoder
126 let mut decoder = tiff::decoder::Decoder::new(file)?;
127 decoder = decoder.with_limits(tiff::decoder::Limits::unlimited());
128
129 // Read the image dimensions
130 let (width, height) = decoder.dimensions()?;
131
132 fn estr<T>(etype: &'static str) -> RasterError {
133 RasterError::Tiff(TiffFormatError::Format(format!("Raster is {}, I was expecting {}", etype, std::any::type_name::<T>()).into()).into())
134 }
135 let data: Vec<T> = match decoder.read_image()? {
136 DecodingResult::I8(buf) => buf.into_iter().map(|v| <T>::from_i8(v).ok_or(estr::<T>("I8"))).collect::<Result<_, _>>(),
137 DecodingResult::I16(buf) => buf.into_iter().map(|v| <T>::from_i16(v).ok_or(estr::<T>("I16"))).collect::<Result<_, _>>(),
138 DecodingResult::I32(buf) => buf.into_iter().map(|v| <T>::from_i32(v).ok_or(estr::<T>("I32"))).collect::<Result<_, _>>(),
139 DecodingResult::I64(buf) => buf.into_iter().map(|v| <T>::from_i64(v).ok_or(estr::<T>("I64"))).collect::<Result<_, _>>(),
140 DecodingResult::U8(buf) => buf.into_iter().map(|v| <T>::from_u8(v).ok_or(estr::<T>("U8"))).collect::<Result<_, _>>(),
141 DecodingResult::U16(buf) => buf.into_iter().map(|v| <T>::from_u16(v).ok_or(estr::<T>("U16"))).collect::<Result<_, _>>(),
142 DecodingResult::U32(buf) => buf.into_iter().map(|v| <T>::from_u32(v).ok_or(estr::<T>("U32"))).collect::<Result<_, _>>(),
143 DecodingResult::U64(buf) => buf.into_iter().map(|v| <T>::from_u64(v).ok_or(estr::<T>("U64"))).collect::<Result<_, _>>(),
144 DecodingResult::F32(buf) => buf.into_iter().map(|v| <T>::from_f32(v).ok_or(estr::<T>("F32"))).collect::<Result<_, _>>(),
145 DecodingResult::F64(buf) => buf.into_iter().map(|v| <T>::from_f64(v).ok_or(estr::<T>("F64"))).collect::<Result<_, _>>(),
146 }?;
147
148 // Convert the flat vector into an ndarray::Array2
149 let array: Array2<T> = Array2::from_shape_vec((height as usize, width as usize), data)?;
150
151 // nodata value
152 let nodata: T = decoder.get_tag_ascii_string(Tag::GdalNodata)?.trim().parse::<T>()?;
153
154 // pixel scale [pixel scale x, pixel scale y, ...]
155 let pscale: Vec<f64> = decoder.get_tag_f64_vec(Tag::ModelPixelScaleTag)?.into_iter().collect();
156
157 // tie point [0 0 0 startx starty 0]
158 let tie: Vec<f64> = decoder.get_tag_f64_vec(Tag::ModelTiepointTag)?.into_iter().collect();
159
160 // transform, the zeros are the rotations [start x, x pixel size, 0, start y, 0, y pixel size]
161 let geotrans: [f64; 6] = [tie[3], pscale[0], 0.0, tie[4], 0.0, -pscale[1]];
162
163 let projection: String = decoder.get_tag_ascii_string(Tag::GeoAsciiParamsTag)?;
164 let geokeydir: Vec<u64> = decoder .get_tag_u64_vec(Tag::GeoKeyDirectoryTag)?;
165
166 // try and get the CRS out of the geokeydir, it is the bit after 3072
167 let crs = geokeydir.windows(4).find(|w| w[0] == 3072).map(|w| w[3])
168 .ok_or(RasterError::Tiff(tiff::TiffFormatError::InvalidTagValueType(Tag::GeoKeyDirectoryTag).into()))? as u16;
169
170 Ok((array, nodata, crs, geotrans, geokeydir, projection))
171}
172
173/// Writes a 2D array of values to a GeoTIFF raster with geo metadata.
174///
175/// # Type Parameters
176/// - `T`: The element type of the array, which must implement `bytemuck::Pod`
177/// (for safe byte casting) and `ToString` (for writing NoData values to
178/// metadata).
179///
180/// # Parameters
181/// - `data`: A 2D array (`ndarray::Array2<T>`) containing raster pixel values.
182/// - `nd`: NoData value
183/// - `geotrans`: A 6-element array defining the affine geotransform:
184/// `[origin_x, pixel_size_x, rotation_x, origin_y, rotation_y, pixel_size_y]`.
185/// - `geokeydir`: &[u64] the GeoKeyDirectoryTag (best got from reading a raster)
186/// - `proj`: PROJ string (best got from reading a raster)
187/// - `outfile`: The path to the output `.tif` file.
188///
189/// # Returns
190/// Ok() or a `RasterError`
191///
192/// # Errors
193/// - Returns `RasterError::UnsupportedType` if `T` can't be mapped to a TIFF format.
194/// - Propagates I/O and TIFF writing errors
195///
196/// # Example
197/// ```
198/// let path = PathBuf::from("input.tif");
199/// let (d8, nd, crs, geo, gdir, proj) = rasterfile_to_array::<u8>(&path)?;
200/// /* do something with d8, or make a new array2 */
201/// let out = PathBuf::from("output.tif");
202/// if let Err(e) = array_to_rasterfile::<u8>(&d8, nd, &geo, &gdir, &proj, &out) {
203/// eprintln!("Error occured while writing {}: {:?}", out.display(), e);
204/// }
205/// ```
206pub fn array_to_rasterfile<T>(
207 data: &Array2<T>,
208 nd: T, // nodata
209 geotrans: &[f64; 6], // geo transform [start_x, psize_x, rotation, starty, rotation, psize_y]
210 geokeydir: &[u64], // geo dir, it has the crs in it
211 proj: &str, // the projection string
212 outfile: &PathBuf
213) -> Result<(), RasterError>
214 where T: bytemuck::Pod + ToString
215{
216 let (nrows, ncols) = (data.nrows(), data.ncols());
217
218 let fh = File::create(outfile)?;
219 let mut encoder = tiff::encoder::TiffEncoder::new(fh)?;
220
221 // Because image doesn't have traits I couldn't figure out how to do this with generics
222 // This macro takes the tiff colortype
223 macro_rules! writit {
224 ($pix:ty) => {{
225 let mut image = encoder.new_image_with_compression::<$pix, Deflate>(ncols as u32, nrows as u32, Deflate::default())?;
226 image.encoder().write_tag(Tag::GdalNodata, &nd.to_string()[..])?;
227 image.encoder().write_tag(Tag::ModelPixelScaleTag, &[geotrans[1], geotrans[5], 0.0][..])?;
228 image.encoder().write_tag(Tag::ModelTiepointTag, &[0.0, 0.0, 0.0, geotrans[0], geotrans[3], 0.0][..])?;
229 image.encoder().write_tag(Tag::GeoKeyDirectoryTag, geokeydir)?;
230 image.encoder().write_tag(Tag::GeoAsciiParamsTag, &proj)?;
231 image.write_data(cast_slice(data.as_slice().unwrap()))?;
232 }};
233 }
234
235 match std::any::TypeId::of::<T>() {
236 id if id == std::any::TypeId::of::<u8>() => writit!(Gray8),
237 id if id == std::any::TypeId::of::<u16>() => writit!(Gray16),
238 id if id == std::any::TypeId::of::<u32>() => writit!(Gray32),
239 id if id == std::any::TypeId::of::<u64>() => writit!(Gray64),
240 id if id == std::any::TypeId::of::<f32>() => writit!(Gray32Float),
241 id if id == std::any::TypeId::of::<f64>() => writit!(Gray64Float),
242 id if id == std::any::TypeId::of::<i8>() => writit!(GrayI8),
243 id if id == std::any::TypeId::of::<i16>() => writit!(GrayI16),
244 id if id == std::any::TypeId::of::<i32>() => writit!(GrayI32),
245 id if id == std::any::TypeId::of::<i64>() => writit!(GrayI64),
246 _ => return Err(RasterError::UnsupportedType(format!("Cannot handle type {}", std::any::type_name::<T>())))
247 };
248
249 Ok(())
250}
251
252
253#[derive(PartialEq, Debug)]
254struct GridCell {
255 row: usize,
256 column: usize,
257 priority: f64,
258}
259
260impl Eq for GridCell {}
261
262impl PartialOrd for GridCell {
263 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
264 other.priority.partial_cmp(&self.priority)
265 }
266}
267
268impl Ord for GridCell {
269 fn cmp(&self, other: &Self) -> Ordering {
270 self.partial_cmp(other).unwrap()
271 }
272}
273
274#[derive(PartialEq, Debug)]
275struct GridCell2 {
276 row: usize,
277 column: usize,
278 z: f64,
279 priority: f64,
280}
281
282impl Eq for GridCell2 {}
283
284impl PartialOrd for GridCell2 {
285 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
286 other.priority.partial_cmp(&self.priority)
287 }
288}
289
290impl Ord for GridCell2 {
291 fn cmp(&self, other: &Self) -> Ordering {
292 self.partial_cmp(other).unwrap()
293 }
294}
295
296
297/// Fills depressions (sinks) in a digital elevation model (DEM).
298///
299/// More-or-less the contents of
300/// [whitebox fill_depressions](https://github.com/jblindsay/whitebox-tools/blob/master/whitebox-tools-app/src/tools/hydro_analysis/fill_depressions.rs)
301///
302/// This function modifies the input `dem` to ensure that all depressions (local minima that do not
303/// drain) are removed, making the surface hydrologically correct. It also considers no-data values
304/// and can optionally fix flat areas.
305///
306/// # Parameters
307///
308/// - `dem`: A mutable reference to a 2D array (`Array2<f64>`) representing the elevation data.
309/// - `nodata`: The value representing no-data cells in the DEM.
310/// - `resx`: The horizontal resolution (grid spacing in the x-direction).
311/// - `resy`: The vertical resolution (grid spacing in the y-direction).
312/// - `fix_flats`: A boolean flag to determine whether flat areas should be slightly sloped.
313///
314/// # Example
315///
316/// ```
317/// use ndarray::Array2;
318/// use hydro_analysis::fill_depressions;
319///
320/// let mut dem = Array2::from_shape_vec(
321/// (3, 3),
322/// vec![
323/// 10.0, 12.0, 10.0,
324/// 12.0, 9.0, 12.0,
325/// 10.0, 12.0, 10.0,
326/// ],
327/// ).expect("Failed to create DEM");
328///
329/// fill_depressions(&mut dem, -3.0, 8.0, 8.0, true);
330/// ```
331pub fn fill_depressions(
332 dem: &mut Array2<f64>, nodata: f64, resx: f64, resy: f64, fix_flats: bool
333)
334{
335 let (rows, columns) = (dem.nrows(), dem.ncols());
336 let small_num = {
337 let diagres = (resx * resx + resy * resy).sqrt();
338 let elev_digits = (dem.iter().cloned().fold(f64::NEG_INFINITY, f64::max) as i64).to_string().len();
339 let elev_multiplier = 10.0_f64.powi((9 - elev_digits) as i32);
340 1.0_f64 / elev_multiplier as f64 * diagres.ceil()
341 };
342
343 //let input = dem.clone(); // original whitebox used an input and output while doing fixing flats, don't think you really need it
344
345 let dx = [1, 1, 1, 0, -1, -1, -1, 0];
346 let dy = [-1, 0, 1, 1, 1, 0, -1, -1];
347
348 // Find pit cells (we don't care about pits around edge, they won't be considered a pit). This step is parallelizable.
349 let mut pits: Vec<_> = (1..rows - 1)
350 .into_par_iter()
351 .flat_map(|row| {
352 let mut local_pits = Vec::new();
353 for col in 1..columns - 1 {
354 let z = dem[[row, col]];
355 if z == nodata {
356 continue;
357 }
358 let mut apit = true;
359 // is anything lower than me?
360 for n in 0..8 {
361 let zn = dem[[(row as isize + dy[n]) as usize, (col as isize + dx[n]) as usize]];
362 if zn < z || zn == nodata {
363 apit = false;
364 break;
365 }
366 }
367 // no, so I am a pit
368 if apit {
369 local_pits.push((row, col, z));
370 }
371 }
372 local_pits
373 }).collect();
374
375 // Now we need to perform an in-place depression filling
376 let mut minheap = BinaryHeap::new();
377 let mut minheap2 = BinaryHeap::new();
378 let mut visited = Array2::<u8>::zeros((rows, columns));
379 let mut flats = Array2::<u8>::zeros((rows, columns));
380 let mut possible_outlets = vec![];
381 let mut queue = VecDeque::new();
382
383 // go through pits from highest to lowest
384 pits.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(Equal));
385 while let Some(cell) = pits.pop() {
386 let row: usize = cell.0;
387 let col: usize = cell.1;
388
389 // if it's already in a solved site, don't do it a second time.
390 if flats[[row, col]] != 1 {
391 // First there is a priority region-growing operation to find the outlets.
392 minheap.clear();
393 minheap.push(GridCell {
394 row: row,
395 column: col,
396 priority: dem[[row, col]],
397 });
398 visited[[row, col]] = 1;
399 let mut outlet_found = false;
400 let mut outlet_z = f64::INFINITY;
401 if !queue.is_empty() {
402 queue.clear();
403 }
404 while let Some(cell2) = minheap.pop() {
405 let z = cell2.priority;
406 if outlet_found && z > outlet_z {
407 break;
408 }
409 if !outlet_found {
410 for n in 0..8 {
411 let cn = cell2.column as isize + dx[n];
412 let rn = cell2.row as isize + dy[n];
413 if rn < 0 || rn as usize >= rows || cn < 0 || cn as usize >= columns {
414 continue;
415 }
416 let cn = cn as usize;
417 let rn = rn as usize;
418 if visited[[rn, cn]] == 0 {
419 let zn = dem[[rn, cn]];
420 if !outlet_found {
421 if zn >= z && zn != nodata {
422 minheap.push(GridCell {
423 row: rn,
424 column: cn,
425 priority: zn,
426 });
427 visited[[rn, cn]] = 1;
428 } else if zn != nodata {
429 // zn < z
430 // 'cell' has a lower neighbour that hasn't already passed through minheap.
431 // Therefore, 'cell' is a pour point cell.
432 outlet_found = true;
433 outlet_z = z;
434 queue.push_back((cell2.row, cell2.column));
435 possible_outlets.push((cell2.row, cell2.column));
436 }
437 } else if zn == outlet_z {
438 // We've found the outlet but are still looking for additional depression cells.
439 minheap.push(GridCell {
440 row: rn,
441 column: cn,
442 priority: zn,
443 });
444 visited[[rn, cn]] = 1;
445 }
446 }
447 }
448 } else {
449 // We've found the outlet but are still looking for additional depression cells and potential outlets.
450 if z == outlet_z {
451 let mut anoutlet = false;
452 for n in 0..8 {
453 let cn = cell2.column as isize + dx[n];
454 let rn = cell2.row as isize + dy[n];
455 if rn < 0 || rn as usize >= rows || cn < 0 || cn as usize >= columns {
456 continue;
457 }
458 let cn = cn as usize;
459 let rn = rn as usize;
460 if visited[[rn, cn]] == 0 {
461 let zn = dem[[rn, cn]];
462 if zn < z {
463 anoutlet = true;
464 } else if zn == outlet_z {
465 minheap.push(GridCell {
466 row: rn,
467 column: cn,
468 priority: zn,
469 });
470 visited[[rn, cn]] = 1;
471 }
472 }
473 }
474 if anoutlet {
475 queue.push_back((cell2.row, cell2.column));
476 possible_outlets.push((cell2.row, cell2.column));
477 } else {
478 visited[[cell2.row, cell2.column]] = 1;
479 }
480 }
481 }
482 }
483
484 if outlet_found {
485 // Now that we have the outlets, raise the interior of the depression.
486 // Start from the outlets.
487 while let Some(cell2) = queue.pop_front() {
488 for n in 0..8 {
489 let cn = cell2.1 as isize + dx[n];
490 let rn = cell2.0 as isize + dy[n];
491 if rn < 0 || rn as usize >= rows || cn < 0 || cn as usize >= columns {
492 continue;
493 }
494 let cn = cn as usize;
495 let rn = rn as usize;
496 if visited[[rn, cn]] == 1 {
497 visited[[rn, cn]] = 0;
498 queue.push_back((rn, cn));
499 let z = dem[[rn, cn]];
500 if z < outlet_z {
501 dem[[rn, cn]] = outlet_z;
502 flats[[rn, cn]] = 1;
503 } else if z == outlet_z {
504 flats[[rn, cn]] = 1;
505 }
506 }
507 }
508 }
509 } else {
510 queue.push_back((row, col)); // start at the pit cell and clean up visited
511 while let Some(cell2) = queue.pop_front() {
512 for n in 0..8 {
513 let cn = cell2.1 as isize + dx[n];
514 let rn = cell2.0 as isize + dy[n];
515 if rn < 0 || rn as usize >= rows || cn < 0 || cn as usize >= columns {
516 continue;
517 }
518 let cn = cn as usize;
519 let rn = rn as usize;
520 if visited[[rn, cn]] == 1 {
521 visited[[rn, cn]] = 0;
522 queue.push_back((rn, cn));
523 }
524 }
525 }
526 }
527 }
528
529 }
530
531 drop(visited);
532
533 if small_num > 0.0 && fix_flats {
534 // Some of the potential outlets really will have lower cells.
535 minheap.clear();
536 while let Some(cell) = possible_outlets.pop() {
537 let z = dem[[cell.0, cell.1]];
538 let mut anoutlet = false;
539 for n in 0..8 {
540 let rn: isize = cell.0 as isize + dy[n];
541 let cn: isize = cell.1 as isize + dx[n];
542 if rn < 0 || rn as usize >= rows || cn < 0 || cn as usize >= columns {
543 continue;
544 }
545 let zn = dem[[rn as usize, cn as usize]];
546 if zn < z && zn != nodata {
547 anoutlet = true;
548 break;
549 }
550 }
551 if anoutlet {
552 minheap.push(GridCell {
553 row: cell.0,
554 column: cell.1,
555 priority: z,
556 });
557 }
558 }
559
560 let mut outlets = vec![];
561 while let Some(cell) = minheap.pop() {
562 if flats[[cell.row, cell.column]] != 3 {
563 let z = dem[[cell.row, cell.column]];
564 flats[[cell.row, cell.column]] = 3;
565 if !outlets.is_empty() {
566 outlets.clear();
567 }
568 outlets.push(cell);
569 // Are there any other outlet cells at the same elevation (likely for the same feature)
570 let mut flag = true;
571 while flag {
572 match minheap.peek() {
573 Some(cell2) => {
574 if cell2.priority == z {
575 flats[[cell2.row, cell2.column]] = 3;
576 outlets
577 .push(minheap.pop().expect("Error during pop operation."));
578 } else {
579 flag = false;
580 }
581 }
582 None => {
583 flag = false;
584 }
585 }
586 }
587 if !minheap2.is_empty() {
588 minheap2.clear();
589 }
590 for cell2 in &outlets {
591 let z = dem[[cell2.row, cell2.column]];
592 for n in 0..8 {
593 let cn = cell2.column as isize + dx[n];
594 let rn = cell2.row as isize + dy[n];
595 if rn < 0 || rn as usize >= rows || cn < 0 || cn as usize >= columns {
596 continue;
597 }
598 let cn = cn as usize;
599 let rn = rn as usize;
600 if flats[[rn, cn]] != 3 {
601 let zn = dem[[rn, cn]];
602 if zn == z && zn != nodata {
603 minheap2.push(GridCell2 {
604 row: rn,
605 column: cn,
606 z: z,
607 priority: dem[[rn, cn]], // FIXME
608 });
609 dem[[rn, cn]] = z + small_num;
610 flats[[rn, cn]] = 3;
611 }
612 }
613 }
614 }
615
616 // Now fix the flats
617 while let Some(cell2) = minheap2.pop() {
618 let z = dem[[cell2.row, cell2.column]];
619 for n in 0..8 {
620 let cn = cell2.column as isize + dx[n];
621 let rn = cell2.row as isize + dy[n];
622 if rn < 0 || rn as usize >= rows || cn < 0 || cn as usize >= columns {
623 continue;
624 }
625 let cn = cn as usize;
626 let rn = rn as usize;
627 if flats[[rn, cn]] != 3 {
628 let zn = dem[[rn, cn]];
629 if zn < z + small_num && zn >= cell2.z && zn != nodata {
630 minheap2.push(GridCell2 {
631 row: rn,
632 column: cn,
633 z: cell2.z,
634 priority: dem[[rn, cn]], // FIXME
635 });
636 dem[[rn, cn]] = z + small_num;
637 flats[[rn, cn]] = 3;
638 }
639 }
640 }
641 }
642 }
643
644 }
645 }
646}
647
648/// Calculates the D8 flow direction from a digital elevation model (DEM).
649///
650/// More-or-less the contents of
651/// [whitebox d8_pointer](https://github.com/jblindsay/whitebox-tools/blob/master/whitebox-tools-app/src/tools/hydro_analysis/d8_pointer.rs)
652///
653/// This function computes the D8 flow direction for each cell in the provided DEM:
654///
655/// | . | . | . |
656/// |:--:|:---:|:--:|
657/// | 64 | 128 | 1 |
658/// | 32 | 0 | 2 |
659/// | 16 | 8 | 4 |
660///
661/// Grid cells that have no lower neighbours are assigned a flow direction of zero. In a DEM that
662/// has been pre-processed to remove all depressions and flat areas, this condition will only occur
663/// along the edges of the grid.
664///
665/// Grid cells possessing the NoData value in the input DEM are assigned the NoData value in the
666/// output image.
667///
668/// # Parameters
669/// - `dem`: A 2D array representing the digital elevation model (DEM)
670/// - `nodata`: The nodata in the DEM
671/// - `resx`: The resolution of the DEM in the x-direction in meters
672/// - `resy`: The resolution of the DEM in the y-direction in meters
673///
674/// # Returns
675/// - A tuple containing:
676/// - An `Array2<u8>` representing the D8 flow directions for each cell.
677/// - A `u8` nodata value (255)
678///
679/// # Example
680/// ```rust
681/// let dem = Array2::from_shape_vec(
682/// (3, 3),
683/// vec![
684/// 10.0, 12.0, 10.0,
685/// 12.0, 13.0, 12.0,
686/// 10.0, 12.0, 10.0,
687/// ],
688/// ).expect("Failed to create DEM");
689/// let nodata = -9999.0;
690/// let resx = 8.0;
691/// let resy = 8.0;
692/// let (d8, nd) = d8_pointer(&dem, nodata, resx, resy);
693/// ```
694pub fn d8_pointer(dem: &Array2<f64>, nodata: f64, resx: f64, resy: f64) -> (Array2<u8>, u8)
695{
696 let (nrows, ncols) = (dem.nrows(), dem.ncols());
697 let out_nodata: u8 = 255;
698 let mut d8: Array2<u8> = Array2::from_elem((nrows, ncols), out_nodata);
699
700 let diag = (resx * resx + resy * resy).sqrt();
701 let grid_lengths = [diag, resx, diag, resy, diag, resx, diag, resy];
702
703 let dx = [1, 1, 1, 0, -1, -1, -1, 0];
704 let dy = [-1, 0, 1, 1, 1, 0, -1, -1];
705
706 d8.axis_iter_mut(ndarray::Axis(0))
707 .into_par_iter()
708 .enumerate()
709 .for_each(|(row, mut d8_row)| {
710 for col in 0..ncols {
711 let z = dem[[row, col]];
712 if z == nodata {
713 continue;
714 }
715
716 let mut dir = 0;
717 let mut max_slope = f64::MIN;
718 for i in 0..8 {
719 let rn: isize = row as isize + dy[i];
720 let cn: isize = col as isize + dx[i];
721 if rn < 0 || rn >= nrows as isize || cn < 0 || cn >= ncols as isize {
722 continue;
723 }
724 let z_n = dem[[rn as usize, cn as usize]];
725 if z_n != nodata {
726 let slope = (z - z_n) / grid_lengths[i];
727 if slope > max_slope && slope > 0.0 {
728 max_slope = slope;
729 dir = i;
730 }
731 }
732 }
733
734 if max_slope >= 0.0 {
735 d8_row[col] = 1 << dir;
736 } else {
737 d8_row[col] = 0u8;
738 }
739 }
740 });
741
742 return (d8, out_nodata);
743}
744
745