terra_texture_rs/curvature.rs
1//! Profile and planform curvature of a DEM, as a fused kernel.
2//!
3//! Mirrors `TerraTexture.derivatives`' `curvatures()` exactly: the same
4//! "gradient of gradient" approximation (`np.gradient` applied twice),
5//! the same `edge_order=1` boundary handling, and the same Zevenbergen &
6//! Thorne curvature formulas. Verified bit-parity against
7//! `derivatives.py` in `tests/test_derivatives_rust.py`.
8//!
9//! # NaN handling
10//!
11//! The DEM passed in must be NaN-free. The Python caller nan-fills it
12//! first (`derivatives.py`'s `_fill_nan_nearest`) and re-masks the
13//! NaN/void cells in the outputs afterwards. Rust owns the plain numeric
14//! math on a clean `f32` array; Python owns the NaN bookkeeping.
15
16use ndarray::{Array2, ArrayView2};
17use rayon::prelude::*;
18
19use crate::common::{gradient2d, PARALLEL_THRESHOLD};
20
21/// Profile and planform curvature for one cell.
22///
23/// Arguments are the cell's partial derivatives, all `f32`:
24/// `p` = dz/dx, `q` = dz/dy, `r` = d²z/dx², `t` = d²z/dy²,
25/// `s` = d²z/dxdy.
26///
27/// Returns `(profile, planform)` as `f32`. Flat cells
28/// (`p² + q² < 1e-9`) return `(0.0, 0.0)`, and any non-finite result
29/// becomes `0.0`, matching `np.nan_to_num(..., nan=0, posinf=0, neginf=0)`.
30#[inline]
31fn curvature_pixel(p: f32, q: f32, r: f32, t: f32, s: f32) -> (f32, f32) {
32 let p2q2 = p * p + q * q;
33 if p2q2 < 1e-9 {
34 return (0.0, 0.0); // flat cell -- matches derivatives.py's `flat` mask
35 }
36 let profile_raw = -(r * p * p + 2.0 * s * p * q + t * q * q) / (p2q2 * (1.0 + p2q2).powf(1.5));
37 let planform_raw = -(r * q * q - 2.0 * s * p * q + t * p * p) / p2q2.powf(1.5);
38
39 // matches np.nan_to_num(..., nan=0.0, posinf=0.0, neginf=0.0)
40 let clean = |v: f32| if v.is_finite() { v } else { 0.0 };
41 (clean(profile_raw), clean(planform_raw))
42}
43
44/// Compute profile and planform curvature of a DEM.
45///
46/// Profile curvature is curvature in the direction of steepest slope
47/// (affects flow acceleration); planform curvature is curvature
48/// perpendicular to it (affects flow convergence). Both use the
49/// Zevenbergen & Thorne sign convention of `derivatives.py`.
50///
51/// # Arguments
52///
53/// * `dem` - `ArrayView2<f32>`, shape (H, W): elevations. **Must be
54/// NaN-free** (see module docs). Any memory layout. H and W should
55/// each be at least 2; see **Edge cases**.
56/// * `cellsize` - `f32`: grid spacing, in the same units as the
57/// elevations, applied to both axes.
58/// * `profile_out` - `&mut Array2<f32>`, shape (H, W), C-contiguous:
59/// overwritten with profile curvature.
60/// * `planform_out` - `&mut Array2<f32>`, shape (H, W), C-contiguous:
61/// overwritten with planform curvature.
62///
63/// Flat cells and any non-finite results are written as `0.0`.
64///
65/// Runs serially below [`PARALLEL_THRESHOLD`]
66/// elements and in parallel at or above it. The gradient passes
67/// themselves always run serially.
68///
69/// # Edge cases
70///
71/// Where numpy's `np.gradient` would raise `ValueError` because an axis
72/// has length 1, this treats that axis's derivatives as zero instead.
73///
74/// # Panics
75///
76/// * If `profile_out` or `planform_out` is not C-contiguous. Arrays made
77/// with `Array2::zeros(shape)` always are.
78/// * If `profile_out` or `planform_out` has more elements than `dem`.
79/// * If exactly one of H and W is 0.
80///
81/// Output shapes are not otherwise checked: outputs with fewer elements
82/// than `dem` are partly filled, and a different shape with the same
83/// element count is filled in the wrong layout. Always allocate both
84/// outputs with `dem.raw_dim()`.
85///
86/// # Example
87///
88/// ```
89/// use ndarray::Array2;
90/// use terra_texture_rs::curvatures_core;
91///
92/// // a uniformly tilted plane has no curvature anywhere
93/// let dem = Array2::from_shape_fn((5, 5), |(i, j)| (i + 2 * j) as f32);
94/// let mut profile = Array2::<f32>::zeros(dem.raw_dim());
95/// let mut planform = Array2::<f32>::zeros(dem.raw_dim());
96/// curvatures_core(dem.view(), 1.0, &mut profile, &mut planform);
97///
98/// assert!(profile.iter().all(|&v| v.abs() < 1e-6));
99/// assert!(planform.iter().all(|&v| v.abs() < 1e-6));
100/// ```
101pub fn curvatures_core(
102 dem: ArrayView2<f32>,
103 cellsize: f32,
104 profile_out: &mut Array2<f32>,
105 planform_out: &mut Array2<f32>,
106) {
107 let (zy, zx) = gradient2d(dem, cellsize);
108 let (zxy, zxx) = gradient2d(zx.view(), cellsize);
109 let (zyy, _zyx) = gradient2d(zy.view(), cellsize); // zyx discarded, matches derivatives.py
110
111 let n = dem.len();
112 // 2 outputs + 5 inputs = 7 producers, one over ndarray::Zip's max
113 // arity of 6 -- fall back to plain contiguous slices + rayon here
114 // instead (all arrays are freshly `zeros()`-allocated, hence
115 // standard/C-contiguous, so `.as_slice()` is always `Some`).
116 let zx_s = zx.as_slice().expect("gradient2d output not contiguous");
117 let zy_s = zy.as_slice().expect("gradient2d output not contiguous");
118 let zxx_s = zxx.as_slice().expect("gradient2d output not contiguous");
119 let zyy_s = zyy.as_slice().expect("gradient2d output not contiguous");
120 let zxy_s = zxy.as_slice().expect("gradient2d output not contiguous");
121 let profile_s = profile_out.as_slice_mut().expect("profile_out not contiguous");
122 let planform_s = planform_out.as_slice_mut().expect("planform_out not contiguous");
123
124 let compute = |i: usize, po: &mut f32, plo: &mut f32| {
125 let (profile, planform) = curvature_pixel(zx_s[i], zy_s[i], zxx_s[i], zyy_s[i], zxy_s[i]);
126 *po = profile;
127 *plo = planform;
128 };
129
130 if n >= PARALLEL_THRESHOLD {
131 profile_s
132 .par_iter_mut()
133 .zip(planform_s.par_iter_mut())
134 .enumerate()
135 .for_each(|(i, (po, plo))| compute(i, po, plo));
136 } else {
137 profile_s
138 .iter_mut()
139 .zip(planform_s.iter_mut())
140 .enumerate()
141 .for_each(|(i, (po, plo))| compute(i, po, plo));
142 }
143}