Skip to main content

geop_core_math/
dual.rs

1//! Forward-mode automatic differentiation, for the residuals of the
2//! constraint solvers (sketches, assemblies).
3//!
4//! Every residual is written once, generically over [`Number`], and
5//! evaluated either with a plain [`Scalar`] value (the value only) or with
6//! [`Dual`] (the value plus its exact partial derivatives with respect to
7//! the constraint's own variables). A constraint touches only a handful of variables, so a dual
8//! number carries a small fixed-size gradient rather than one entry per
9//! solver variable.
10//!
11//! Building `Dual` over the interval [`Scalar`] rather than a bare `f64`
12//! means a residual's value is a rigorous enclosure of the true real number
13//! it stands for, not a float that silently absorbs whatever rounding `sin`,
14//! `sqrt` or `PI` introduced along the way — the sketch solver's own
15//! uncertainty about an irrational quantity shows up as interval width
16//! instead of vanishing.
17
18use crate::{
19    geop_error::GeopResult,
20    scalars::{Number, Scalar},
21};
22
23/// A value together with its gradient with respect to up to `N` seeded
24/// variables — as many as the residual it is computed for depends on.
25#[derive(Clone, Copy, Debug)]
26pub struct Dual<const N: usize> {
27    pub v: Scalar,
28    pub d: [Scalar; N],
29}
30
31impl<const N: usize> core::fmt::Display for Dual<N> {
32    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
33        core::fmt::Display::fmt(&self.v, f)
34    }
35}
36
37impl<const N: usize> Default for Dual<N> {
38    fn default() -> Self {
39        Dual::cst(Scalar::ZERO)
40    }
41}
42
43impl<const N: usize> Dual<N> {
44    /// A constant, with no dependence on any variable.
45    pub const fn cst(v: Scalar) -> Self {
46        Dual {
47            v,
48            d: [Scalar::ZERO; N],
49        }
50    }
51
52    /// The `slot`-th independent variable, with value `v`.
53    pub fn var(v: Scalar, slot: usize) -> Self {
54        let mut d = [Scalar::ZERO; N];
55        d[slot] = Scalar::ONE;
56        Dual { v, d }
57    }
58
59    /// The same value in a gradient of `M` variables: its own first `vars`
60    /// the ones from `slot` on — and no dependence on any other (its own
61    /// beyond `vars` it has none of).
62    pub fn embed<const M: usize>(self, slot: usize, vars: usize) -> Dual<M> {
63        let mut d = [Scalar::ZERO; M];
64        d[slot..slot + vars].copy_from_slice(&self.d[..vars]);
65        Dual { v: self.v, d }
66    }
67
68    /// `f(self)` for a scalar function with value `v` and derivative `df =
69    /// f'(self.v)`: `d <- d * df` by the chain rule.
70    pub fn chain(self, v: Scalar, df: Scalar) -> Self {
71        let mut d = self.d;
72        for x in &mut d {
73            *x = x.mul(df);
74        }
75        Dual { v, d }
76    }
77}
78
79impl<const N: usize> Number for Dual<N> {
80    const ZERO: Self = Self::cst(Scalar::ZERO);
81    const ONE: Self = Self::cst(Scalar::ONE);
82    const TWO: Self = Self::cst(Scalar::TWO);
83    const PI: Self = Self::cst(Scalar::PI);
84
85    fn from_i64(v: i64) -> Self {
86        Dual::cst(Scalar::from_i64(v))
87    }
88
89    fn from_f64(v: f64) -> Self {
90        Dual::cst(Scalar::from_f64(v))
91    }
92
93    fn from_scalar(s: Scalar) -> Self {
94        Dual::cst(s)
95    }
96
97    fn from_ratio(num: i64, den: i64) -> GeopResult<Self> {
98        Ok(Dual::cst(Scalar::from_ratio(num, den)?))
99    }
100
101    fn to_f64(self) -> f64 {
102        self.v.to_f64()
103    }
104
105    fn add(self, o: Self) -> Self {
106        let mut d = self.d;
107        for (x, y) in d.iter_mut().zip(o.d) {
108            *x = x.add(y);
109        }
110        Dual {
111            v: self.v.add(o.v),
112            d,
113        }
114    }
115
116    fn sub(self, o: Self) -> Self {
117        let mut d = self.d;
118        for (x, y) in d.iter_mut().zip(o.d) {
119            *x = x.sub(y);
120        }
121        Dual {
122            v: self.v.sub(o.v),
123            d,
124        }
125    }
126
127    fn mul(self, o: Self) -> Self {
128        let mut d = [Scalar::ZERO; N];
129        for (i, x) in d.iter_mut().enumerate() {
130            *x = self.d[i].mul(o.v).add(o.d[i].mul(self.v));
131        }
132        Dual {
133            v: self.v.mul(o.v),
134            d,
135        }
136    }
137
138    fn neg(self) -> Self {
139        self.chain(self.v.neg(), Scalar::ZERO.sub(Scalar::ONE))
140    }
141
142    fn div(self, o: Self) -> GeopResult<Self> {
143        let v = self.v.div(o.v)?;
144        let mut d = [Scalar::ZERO; N];
145        for (i, x) in d.iter_mut().enumerate() {
146            // Quotient rule: (self' - v * o') / o.v.
147            *x = self.d[i].sub(v.mul(o.d[i])).div(o.v)?;
148        }
149        Ok(Dual { v, d })
150    }
151
152    // `could_be_equal`/`is_infinite`/etc. below act on the value alone,
153    // ignoring the gradient — a residual's dual number is compared and
154    // reported on the same terms a plain [`Scalar`] would be, and the
155    // interval-lattice operations (`intersect`, `midpoint`, `width`, ...)
156    // exist on [`Dual`] only because [`Number`] requires them, not because
157    // a constraint ever calls them: the sketch solver only ever adds,
158    // multiplies, divides, and takes `sqrt`/`sin`/`cos`/`abs` of a residual.
159
160    fn abs(self) -> Self {
161        // Not differentiable at a sign change; 0 is the subgradient (the
162        // straddling case falls through to the `+1` branch, matching
163        // `Scalar::abs`'s own choice to keep the value non-negative there).
164        if self.v.definitely_less(Scalar::ZERO) {
165            self.chain(self.v.abs(), Scalar::ZERO.sub(Scalar::ONE))
166        } else {
167            self.chain(self.v.abs(), Scalar::ONE)
168        }
169    }
170
171    fn sqrt(self) -> GeopResult<Self> {
172        let s = self.v.sqrt()?;
173        // `sqrt` is not differentiable at 0. Residuals only take roots of
174        // squared lengths, whose gradient vanishes there too, so 0 is the
175        // right subgradient: it leaves the other terms in charge.
176        let df = if s.definitely_greater(Scalar::ZERO) {
177            Scalar::ONE.div(Scalar::TWO.mul(s))?
178        } else {
179            Scalar::ZERO
180        };
181        Ok(self.chain(s, df))
182    }
183
184    fn sin(self) -> Self {
185        self.chain(self.v.sin(), self.v.cos())
186    }
187
188    fn cos(self) -> Self {
189        let df = self.v.sin().neg();
190        self.chain(self.v.cos(), df)
191    }
192
193    /// `d acos(x) = -1 / sqrt(1 - x^2)`, which fails where it is infinite,
194    /// at `x = ±1`.
195    fn acos(self) -> GeopResult<Self> {
196        let root = Scalar::ONE.sub(self.v.mul(self.v)).sqrt()?;
197        let df = Scalar::ONE.div(root)?.neg();
198        Ok(self.chain(self.v.acos()?, df))
199    }
200
201    fn could_be_equal(self, other: Self) -> bool {
202        self.v.could_be_equal(other.v)
203    }
204
205    fn definitely_not_equal(self, other: Self) -> bool {
206        self.v.definitely_not_equal(other.v)
207    }
208
209    fn could_be_greater(self, other: Self) -> bool {
210        self.v.could_be_greater(other.v)
211    }
212
213    fn definitely_greater(self, other: Self) -> bool {
214        self.v.definitely_greater(other.v)
215    }
216
217    fn could_be_less(self, other: Self) -> bool {
218        self.v.could_be_less(other.v)
219    }
220
221    fn definitely_less(self, other: Self) -> bool {
222        self.v.definitely_less(other.v)
223    }
224
225    /// A single, exactly known point: a sharp value that does not vary
226    /// with any variable. A value that is sharp but moves with a variable is
227    /// no point — taking it for one drops the gradient.
228    fn is_sharp(self) -> bool {
229        self.v.is_sharp() && self.d.iter().all(|g| g.is_zero())
230    }
231
232    fn is_exactly(self, other: Self) -> bool {
233        self.is_sharp() && other.is_sharp() && self.v.is_exactly(other.v)
234    }
235
236    /// The union of the values, with `self`'s derivatives: the width a
237    /// series' remainder adds (see [`crate::primitives::Pose::exp`]).
238    fn union(self, other: Self) -> Self {
239        Dual {
240            v: self.v.union(other.v),
241            d: self.d,
242        }
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    /// Every operation's derivative against a central difference.
251    #[test]
252    fn dual_matches_finite_differences() {
253        fn f<T: Number>(x: T, y: T) -> T {
254            let numerator = x.mul(y).add(x.sin().mul(y.cos()));
255            let denom = x.mul(x).add(T::ONE).sqrt().unwrap();
256            numerator.div(denom).unwrap().sub(y.sub(x).abs())
257        }
258        let (x, y) = (0.7, -1.3);
259        let d = f(
260            Dual::<2>::var(Scalar::from_f64(x), 0),
261            Dual::<2>::var(Scalar::from_f64(y), 1),
262        );
263        let h = 1e-6;
264        let fv = |x: f64, y: f64| f(Scalar::from_f64(x), Scalar::from_f64(y)).to_f64();
265        let dx = (fv(x + h, y) - fv(x - h, y)) / (2.0 * h);
266        let dy = (fv(x, y + h) - fv(x, y - h)) / (2.0 * h);
267        assert!((d.v.to_f64() - fv(x, y)).abs() < 1e-15);
268        assert!((d.d[0].to_f64() - dx).abs() < 1e-8, "{} vs {dx}", d.d[0]);
269        assert!((d.d[1].to_f64() - dy).abs() < 1e-8, "{} vs {dy}", d.d[1]);
270    }
271}