oxigdal_gpu/
algebra.rs

1//! GPU-accelerated raster algebra operations.
2//!
3//! Provides element-wise band math operations (`AlgebraOp`), a structured
4//! expression tree (`BandExpression`) for composing multi-band formulas, and
5//! the top-level `GpuAlgebra` driver that dispatches to GPU (future) or CPU.
6
7use crate::error::GpuError;
8
9/// Element-wise raster algebra operation.
10#[derive(Debug, Clone, PartialEq)]
11pub enum AlgebraOp {
12    /// `A + B`
13    Add,
14    /// `A - B`
15    Subtract,
16    /// `A * B`
17    Multiply,
18    /// `A / B` (outputs nodata when `|B| < 1e-10`)
19    Divide,
20    /// `min(A, B)`
21    Min,
22    /// `max(A, B)`
23    Max,
24    /// `sqrt(max(0, A))`
25    Sqrt,
26    /// `|A|`
27    Abs,
28    /// `A ^ exp`
29    Power(f32),
30    /// `clamp(A, min, max)`
31    Clamp { min: f32, max: f32 },
32    /// Linear stretch: maps `[src_min, src_max]` → `[dst_min, dst_max]`
33    Normalize {
34        src_min: f32,
35        src_max: f32,
36        dst_min: f32,
37        dst_max: f32,
38    },
39}
40
41/// Pure-Rust raster algebra executor.
42pub struct GpuAlgebra;
43
44impl GpuAlgebra {
45    /// Execute an algebra operation pixel-by-pixel (CPU fallback).
46    ///
47    /// `band_b` is required for binary operations (`Add`, `Subtract`,
48    /// `Multiply`, `Divide`, `Min`, `Max`).  For unary operations it is
49    /// ignored.
50    ///
51    /// # Errors
52    ///
53    /// Returns [`GpuError::InvalidKernelParams`] if `band_a` is empty.
54    pub fn execute(
55        band_a: &[f32],
56        band_b: Option<&[f32]>,
57        op: AlgebraOp,
58        nodata: Option<f32>,
59    ) -> Result<Vec<f32>, GpuError> {
60        if band_a.is_empty() {
61            return Err(GpuError::invalid_kernel_params("band_a must not be empty"));
62        }
63
64        let nodata_val = nodata.unwrap_or(f32::NAN);
65        let mut output = Vec::with_capacity(band_a.len());
66
67        for (i, &a) in band_a.iter().enumerate() {
68            // Nodata check for band A.
69            if nodata.is_some() && Self::is_nodata(a, nodata_val) {
70                output.push(nodata_val);
71                continue;
72            }
73
74            let b = band_b.and_then(|bb| bb.get(i)).copied().unwrap_or(0.0_f32);
75
76            // Nodata check for band B.
77            if nodata.is_some() && band_b.is_some() && Self::is_nodata(b, nodata_val) {
78                output.push(nodata_val);
79                continue;
80            }
81
82            let result = match &op {
83                AlgebraOp::Add => a + b,
84                AlgebraOp::Subtract => a - b,
85                AlgebraOp::Multiply => a * b,
86                AlgebraOp::Divide => {
87                    if b.abs() > 1e-10 {
88                        a / b
89                    } else {
90                        nodata_val
91                    }
92                }
93                AlgebraOp::Min => a.min(b),
94                AlgebraOp::Max => a.max(b),
95                AlgebraOp::Sqrt => a.max(0.0).sqrt(),
96                AlgebraOp::Abs => a.abs(),
97                AlgebraOp::Power(exp) => a.powf(*exp),
98                AlgebraOp::Clamp { min, max } => a.clamp(*min, *max),
99                AlgebraOp::Normalize {
100                    src_min,
101                    src_max,
102                    dst_min,
103                    dst_max,
104                } => {
105                    let range = src_max - src_min;
106                    if range.abs() < 1e-10 {
107                        *dst_min
108                    } else {
109                        (a - src_min) / range * (dst_max - dst_min) + dst_min
110                    }
111                }
112            };
113
114            output.push(result);
115        }
116
117        Ok(output)
118    }
119
120    /// Evaluate a multi-band expression for every pixel.
121    ///
122    /// All bands in `bands` must have the same length.  Pixels where any
123    /// band holds the nodata value are written as nodata without evaluating
124    /// the expression.
125    ///
126    /// # Errors
127    ///
128    /// Returns [`GpuError::InvalidKernelParams`] if no bands are provided.
129    /// Propagates any error from `expression.evaluate`.
130    pub fn evaluate_expression(
131        bands: &[&[f32]],
132        expression: &BandExpression,
133        nodata: Option<f32>,
134    ) -> Result<Vec<f32>, GpuError> {
135        if bands.is_empty() {
136            return Err(GpuError::invalid_kernel_params("no bands provided"));
137        }
138
139        let len = bands[0].len();
140        let nodata_val = nodata.unwrap_or(f32::NAN);
141
142        let mut output = Vec::with_capacity(len);
143        for i in 0..len {
144            // Check nodata across all bands.
145            let has_nodata = nodata.is_some()
146                && bands.iter().any(|b| {
147                    b.get(i)
148                        .map(|v| Self::is_nodata(*v, nodata_val))
149                        .unwrap_or(false)
150                });
151
152            if has_nodata {
153                output.push(nodata_val);
154                continue;
155            }
156
157            let vals: Vec<f32> = bands
158                .iter()
159                .map(|b| b.get(i).copied().unwrap_or(0.0))
160                .collect();
161            output.push(expression.evaluate(&vals)?);
162        }
163
164        Ok(output)
165    }
166
167    #[inline]
168    fn is_nodata(value: f32, nodata: f32) -> bool {
169        (value - nodata).abs() < 1e-6
170    }
171}
172
173/// A composable expression tree for multi-band raster math.
174///
175/// Leaf nodes are either a `Band` index or a scalar `Constant`.
176/// Interior nodes are arithmetic operators.
177#[derive(Debug, Clone)]
178pub enum BandExpression {
179    /// Reference to band at the given index.
180    Band(usize),
181    /// Scalar constant.
182    Constant(f32),
183    /// Addition: `A + B`
184    Add(Box<BandExpression>, Box<BandExpression>),
185    /// Subtraction: `A - B`
186    Sub(Box<BandExpression>, Box<BandExpression>),
187    /// Multiplication: `A * B`
188    Mul(Box<BandExpression>, Box<BandExpression>),
189    /// Division: `A / B` (errors on divide-by-zero)
190    Div(Box<BandExpression>, Box<BandExpression>),
191    /// Square root: `sqrt(max(0, A))`
192    Sqrt(Box<BandExpression>),
193    /// Absolute value: `|A|`
194    Abs(Box<BandExpression>),
195    /// Negation: `-A`
196    Neg(Box<BandExpression>),
197    /// Element-wise minimum: `min(A, B)`
198    Min(Box<BandExpression>, Box<BandExpression>),
199    /// Element-wise maximum: `max(A, B)`
200    Max(Box<BandExpression>, Box<BandExpression>),
201    /// Power: `A ^ B` (uses `pow(A, B)` in WGSL).
202    Pow(Box<BandExpression>, Box<BandExpression>),
203    /// Natural logarithm: `ln(A)`.
204    Log(Box<BandExpression>),
205    /// Exponential: `exp(A)`.
206    Exp(Box<BandExpression>),
207    /// Clamp: `clamp(A, lo, hi)` — equivalent to `min(max(A, lo), hi)`.
208    Clamp {
209        value: Box<BandExpression>,
210        lo: Box<BandExpression>,
211        hi: Box<BandExpression>,
212    },
213}
214
215impl BandExpression {
216    /// Evaluate the expression for one pixel given per-band values.
217    ///
218    /// # Errors
219    ///
220    /// Returns [`GpuError::InvalidKernelParams`] when a `Band` index is out
221    /// of range or a `Div` node encounters a zero denominator.
222    pub fn evaluate(&self, bands: &[f32]) -> Result<f32, GpuError> {
223        match self {
224            BandExpression::Band(idx) => bands.get(*idx).copied().ok_or_else(|| {
225                GpuError::invalid_kernel_params(format!(
226                    "band index {} out of range (have {} bands)",
227                    idx,
228                    bands.len()
229                ))
230            }),
231            BandExpression::Constant(v) => Ok(*v),
232            BandExpression::Add(a, b) => Ok(a.evaluate(bands)? + b.evaluate(bands)?),
233            BandExpression::Sub(a, b) => Ok(a.evaluate(bands)? - b.evaluate(bands)?),
234            BandExpression::Mul(a, b) => Ok(a.evaluate(bands)? * b.evaluate(bands)?),
235            BandExpression::Div(a, b) => {
236                let denom = b.evaluate(bands)?;
237                if denom.abs() < 1e-10 {
238                    Err(GpuError::invalid_kernel_params(
239                        "division by zero in BandExpression",
240                    ))
241                } else {
242                    Ok(a.evaluate(bands)? / denom)
243                }
244            }
245            BandExpression::Sqrt(a) => Ok(a.evaluate(bands)?.max(0.0).sqrt()),
246            BandExpression::Abs(a) => Ok(a.evaluate(bands)?.abs()),
247            BandExpression::Neg(a) => Ok(-a.evaluate(bands)?),
248            BandExpression::Min(a, b) => Ok(a.evaluate(bands)?.min(b.evaluate(bands)?)),
249            BandExpression::Max(a, b) => Ok(a.evaluate(bands)?.max(b.evaluate(bands)?)),
250            BandExpression::Pow(a, b) => Ok(a.evaluate(bands)?.powf(b.evaluate(bands)?)),
251            BandExpression::Log(a) => Ok(a.evaluate(bands)?.ln()),
252            BandExpression::Exp(a) => Ok(a.evaluate(bands)?.exp()),
253            BandExpression::Clamp { value, lo, hi } => {
254                let v = value.evaluate(bands)?;
255                let l = lo.evaluate(bands)?;
256                let h = hi.evaluate(bands)?;
257                Ok(v.clamp(l, h))
258            }
259        }
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266
267    #[test]
268    fn test_execute_empty_band_a() {
269        let result = GpuAlgebra::execute(&[], None, AlgebraOp::Add, None);
270        assert!(result.is_err());
271    }
272
273    #[test]
274    fn test_execute_add() {
275        let a = vec![1.0_f32, 2.0, 3.0];
276        let b = vec![4.0_f32, 5.0, 6.0];
277        let out = GpuAlgebra::execute(&a, Some(&b), AlgebraOp::Add, None).expect("execute failed");
278        assert_eq!(out, vec![5.0, 7.0, 9.0]);
279    }
280
281    #[test]
282    fn test_expression_band_out_of_range() {
283        let expr = BandExpression::Band(5);
284        assert!(expr.evaluate(&[1.0, 2.0]).is_err());
285    }
286
287    #[test]
288    fn test_expression_div_by_zero() {
289        let expr = BandExpression::Div(
290            Box::new(BandExpression::Band(0)),
291            Box::new(BandExpression::Constant(0.0)),
292        );
293        assert!(expr.evaluate(&[1.0]).is_err());
294    }
295}
oxigdal_gpu/algebra.rs

oxigdal_gpu/
algebra.rs
