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