treams_core/numerics/
jet.rs

1//! First-order forward-mode jets: a value with its derivatives along `N` directions,
2//! propagated through numerical kernels by local chain rules. `N = 0` compiles out all
3//! derivative storage and arithmetic in forward evaluations. treams-rs extension.
4#![allow(clippy::suspicious_arithmetic_impl, clippy::indexing_slicing)] // Product and quotient rules; indices run over the `N` derivative slots.
5
6use std::{
7    iter::Sum,
8    ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign},
9};
10
11use crate::{Complex, numerics::ratio};
12
13/// A complex value and its first derivatives with respect to `N` inputs.
14///
15/// Kernels that are generic in `N` serve forward evaluation and differentiation alike.
16/// With `N = 0` the derivative array is empty, so a forward evaluation stores and
17/// computes no derivatives. Operations with a real or complex scalar skip the zero
18/// derivative terms of a constant jet. Adding or subtracting a scalar still gives the
19/// value a constant jet gives, signed zeros included: the sign of a zero imaginary part
20/// selects the side of a branch cut downstream. `special::RadialJet` differs: it holds
21/// the first two derivatives in one variable.
22#[derive(Clone, Copy, Debug)]
23pub(crate) struct Jet<const N: usize> {
24    /// The function value.
25    pub(crate) value: Complex,
26    /// The derivative with respect to each of the `N` inputs.
27    pub(crate) derivative: [Complex; N],
28}
29impl<const N: usize> Jet<N> {
30    /// A value that depends on no input: every derivative is zero.
31    pub(crate) fn constant(value: impl Into<Complex>) -> Self {
32        Self {
33            value: value.into(),
34            derivative: [Complex::default(); N],
35        }
36    }
37    /// Input number `index` itself: derivative one in that direction, zero in the others.
38    /// An `index` of `N` or more gives a constant.
39    pub(crate) fn variable(value: impl Into<Complex>, index: usize) -> Self {
40        let mut result = Self::constant(value);
41        if let Some(d) = result.derivative.get_mut(index) {
42            *d = Complex::new(1.0, 0.0);
43        }
44        result
45    }
46    /// The chain rule: the jet of `f(self)` from `value = f(x)` and `derivative = f'(x)`
47    /// at `x = self.value`. Each derivative of `self` is multiplied by `f'(x)`.
48    pub(crate) fn chain(self, value: Complex, derivative: Complex) -> Self {
49        Self {
50            value,
51            derivative: self.derivative.map(|g| g * derivative),
52        }
53    }
54    /// Integer power `self^n`.
55    pub(crate) fn powi(self, n: i32) -> Self {
56        if n == 0 {
57            return Self::constant(1.0);
58        }
59        if N == 0 {
60            return Self::constant(self.value.powi(n));
61        }
62        self.chain(self.value.powi(n), f64::from(n) * self.value.powi(n - 1))
63    }
64    /// Complex exponential.
65    pub(crate) fn exp(self) -> Self {
66        let value = self.value.exp();
67        self.chain(value, value)
68    }
69    /// Principal square root, with the branch of `numerics::complex_sqrt`.
70    pub(crate) fn sqrt(self) -> Self {
71        let root = crate::numerics::complex_sqrt(self.value);
72        self.chain(root, 0.5 / root)
73    }
74    /// The largest modulus of the value and the derivatives.
75    pub(crate) fn norm(self) -> f64 {
76        self.derivative
77            .iter()
78            .map(|g| g.norm())
79            .fold(self.value.norm(), f64::max)
80    }
81    /// Whether the value and every derivative are finite.
82    pub(crate) fn finite(self) -> bool {
83        crate::numerics::finite(self.value)
84            && self.derivative.iter().all(|&g| crate::numerics::finite(g))
85    }
86    /// Whether the value and every derivative are zero, without the moduli of `norm`.
87    pub(crate) fn is_zero(self) -> bool {
88        let zero = Complex::default();
89        self.value == zero && self.derivative.iter().all(|&g| g == zero)
90    }
91}
92impl Jet<3> {
93    /// The chain rule through three intermediate variables: `self` holds `f` and its
94    /// derivatives with respect to `x_b`, `inner[b]` holds `x_b` and its derivatives with
95    /// respect to `N` inputs. Returns the value of `f` with the derivatives
96    /// `sum_b df/dx_b dx_b/dt_a`, summed in the order `b = 0, 1, 2`.
97    pub(crate) fn compose<const N: usize>(self, inner: [Jet<N>; 3]) -> Jet<N> {
98        Jet {
99            value: self.value,
100            derivative: std::array::from_fn(|a| {
101                (0..3)
102                    .map(|b| self.derivative[b] * inner[b].derivative[a])
103                    .sum()
104            }),
105        }
106    }
107}
108impl<const N: usize> Default for Jet<N> {
109    fn default() -> Self {
110        Self::constant(0.0)
111    }
112}
113impl<const N: usize> Add for Jet<N> {
114    type Output = Self;
115    fn add(self, rhs: Self) -> Self {
116        Self {
117            value: self.value + rhs.value,
118            derivative: std::array::from_fn(|i| self.derivative[i] + rhs.derivative[i]),
119        }
120    }
121}
122impl<const N: usize> Sub for Jet<N> {
123    type Output = Self;
124    fn sub(self, rhs: Self) -> Self {
125        Self {
126            value: self.value - rhs.value,
127            derivative: std::array::from_fn(|i| self.derivative[i] - rhs.derivative[i]),
128        }
129    }
130}
131impl<const N: usize> Mul for Jet<N> {
132    type Output = Self;
133    fn mul(self, rhs: Self) -> Self {
134        Self {
135            value: self.value * rhs.value,
136            derivative: std::array::from_fn(|i| {
137                self.derivative[i] * rhs.value + self.value * rhs.derivative[i]
138            }),
139        }
140    }
141}
142impl<const N: usize> Div for Jet<N> {
143    type Output = Self;
144    fn div(self, rhs: Self) -> Self {
145        let value = ratio(self.value, rhs.value);
146        Self {
147            value,
148            derivative: std::array::from_fn(|i| {
149                ratio(self.derivative[i] - value * rhs.derivative[i], rhs.value)
150            }),
151        }
152    }
153}
154impl<const N: usize> Neg for Jet<N> {
155    type Output = Self;
156    fn neg(self) -> Self {
157        Self {
158            value: -self.value,
159            derivative: self.derivative.map(|g| -g),
160        }
161    }
162}
163impl<const N: usize> Sum for Jet<N> {
164    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
165        iter.fold(Self::default(), Add::add)
166    }
167}
168// A scalar acts as a constant jet without evaluating its zero derivative terms. Values
169// combine with the scalar as a complex number, exactly as with a constant jet:
170// num-complex's real-scalar addition keeps a -0.0 imaginary part that a constant's
171// +0.0 would turn into +0.0, and that sign selects branch sides downstream.
172macro_rules! scalar_ops {
173    ($scalar:ty) => {
174        impl<const N: usize> Add<$scalar> for Jet<N> {
175            type Output = Self;
176            fn add(self, rhs: $scalar) -> Self {
177                Self {
178                    value: self.value + Complex::from(rhs),
179                    derivative: self.derivative,
180                }
181            }
182        }
183        impl<const N: usize> Add<Jet<N>> for $scalar {
184            type Output = Jet<N>;
185            fn add(self, rhs: Jet<N>) -> Jet<N> {
186                Jet {
187                    value: Complex::from(self) + rhs.value,
188                    derivative: rhs.derivative,
189                }
190            }
191        }
192        impl<const N: usize> Sub<$scalar> for Jet<N> {
193            type Output = Self;
194            fn sub(self, rhs: $scalar) -> Self {
195                Self {
196                    value: self.value - Complex::from(rhs),
197                    derivative: self.derivative,
198                }
199            }
200        }
201        impl<const N: usize> Sub<Jet<N>> for $scalar {
202            type Output = Jet<N>;
203            fn sub(self, rhs: Jet<N>) -> Jet<N> {
204                Jet {
205                    value: Complex::from(self) - rhs.value,
206                    derivative: rhs.derivative.map(|g| -g),
207                }
208            }
209        }
210        impl<const N: usize> Mul<$scalar> for Jet<N> {
211            type Output = Self;
212            fn mul(self, rhs: $scalar) -> Self {
213                Self {
214                    value: self.value * rhs,
215                    derivative: self.derivative.map(|g| g * rhs),
216                }
217            }
218        }
219        impl<const N: usize> Mul<Jet<N>> for $scalar {
220            type Output = Jet<N>;
221            fn mul(self, rhs: Jet<N>) -> Jet<N> {
222                rhs * self
223            }
224        }
225        impl<const N: usize> Div<Jet<N>> for $scalar {
226            type Output = Jet<N>;
227            fn div(self, rhs: Jet<N>) -> Jet<N> {
228                Jet::constant(self) / rhs
229            }
230        }
231    };
232}
233macro_rules! assignments {
234    ($trait:ident, $method:ident, $optrait:ident, $method2:ident) => {
235        impl<const N: usize, Rhs> $trait<Rhs> for Jet<N>
236        where
237            Self: $optrait<Rhs, Output = Self>,
238        {
239            fn $method(&mut self, rhs: Rhs) {
240                *self = self.$method2(rhs);
241            }
242        }
243    };
244}
245scalar_ops!(f64);
246scalar_ops!(Complex);
247// Real scaling must stay a real division: complex division can square a
248// subnormal denominator to zero when normalizing a nearly axial wavevector.
249impl<const N: usize> Div<f64> for Jet<N> {
250    type Output = Self;
251    fn div(self, rhs: f64) -> Self {
252        Self {
253            value: self.value / rhs,
254            derivative: self.derivative.map(|g| g / rhs),
255        }
256    }
257}
258impl<const N: usize> Div<Complex> for Jet<N> {
259    type Output = Self;
260    fn div(self, rhs: Complex) -> Self {
261        Self {
262            value: ratio(self.value, rhs),
263            derivative: self.derivative.map(|g| ratio(g, rhs)),
264        }
265    }
266}
267assignments!(AddAssign, add_assign, Add, add);
268assignments!(SubAssign, sub_assign, Sub, sub);
269assignments!(MulAssign, mul_assign, Mul, mul);
270
271/// Algebraic identities of jet arithmetic and the chain rule against closed forms.
272#[cfg(test)]
273mod tests {
274    use std::f64::consts::PI;
275
276    use proptest::{prelude::*, test_runner::TestCaseError};
277
278    use super::Jet;
279    use crate::{
280        Complex,
281        numerics::complex_sqrt,
282        test_support::{ALGEBRA_CASES, DEFAULT_CASES, complex, log_polar, prop_assert_close},
283    };
284
285    /// Two-direction jets with value modulus in [0.1, 10] and derivatives in [-5, 5]^2.
286    fn jet() -> impl Strategy<Value = Jet<2>> {
287        (log_polar(-1.0..1.0, -PI..PI), complex(5.0), complex(5.0)).prop_map(|(value, a, b)| Jet {
288            value,
289            derivative: [a, b],
290        })
291    }
292
293    proptest! {
294        #![proptest_config(ProptestConfig::with_cases(ALGEBRA_CASES))]
295
296        #[test]
297        fn scalar_operations_equal_constant_jets(
298            a in jet(), b in jet(), re in -100.0_f64..100.0, im in -100.0_f64..100.0,
299        ) {
300            check_scalar_operations(a, b, re, Complex::new(re, im))?;
301        }
302
303        #[test]
304        fn jet_operations_invert(a in jet(), b in jet(), n in -4_i32..=6) {
305            check_inverses(a, b, n)?;
306        }
307
308        #[test]
309        fn chain_rule_matches_closed_form(
310            z in log_polar(-0.52..0.48, -PI..PI), w in complex(2.0), n in -4_i32..=4,
311        ) {
312            check_chain_rule(z, w, n)?;
313        }
314    }
315
316    proptest! {
317        #![proptest_config(ProptestConfig::with_cases(DEFAULT_CASES))]
318
319        #[test]
320        fn local_quotient_common_scale(
321            x in -0.3_f64..0.3, exponent in prop_oneof![Just(-300), Just(0), Just(300)],
322        ) {
323            check_local_quotient_common_scale(Complex::new(0.7, x), 10_f64.powi(exponent))?;
324        }
325    }
326
327    fn assert_equal(
328        actual: Jet<2>,
329        expected: Jet<2>,
330        operation: &str,
331    ) -> Result<(), TestCaseError> {
332        prop_assert_eq!(actual.value, expected.value, "{}", operation);
333        prop_assert_eq!(actual.derivative, expected.derivative, "{}", operation);
334        Ok(())
335    }
336
337    /// Real and complex scalars act as constant jets, from either side, in the binary
338    /// and assignment forms; sums and linear combinations are exact.
339    fn check_scalar_operations(
340        a: Jet<2>,
341        b: Jet<2>,
342        real: f64,
343        scalar: Complex,
344    ) -> Result<(), TestCaseError> {
345        let (r, c) = (Jet::constant(real), Jet::constant(scalar));
346        for (actual, expected, operation) in [
347            (a + real, a + r, "jet + real"),
348            (real + a, r + a, "real + jet"),
349            (a - real, a - r, "jet - real"),
350            (real - a, r - a, "real - jet"),
351            (a * real, a * r, "jet * real"),
352            (real * a, r * a, "real * jet"),
353            (a / real, a / r, "jet / real"),
354            (a + scalar, a + c, "jet + complex"),
355            (scalar + a, c + a, "complex + jet"),
356            (a - scalar, a - c, "jet - complex"),
357            (scalar - a, c - a, "complex - jet"),
358            (a * scalar, a * c, "jet * complex"),
359            (scalar * a, c * a, "complex * jet"),
360            (a / scalar, a / c, "jet / complex"),
361            ([a, b].into_iter().sum(), a + b, "sum"),
362        ] {
363            assert_equal(actual, expected, operation)?;
364        }
365        let mut assigned = a;
366        assigned *= real;
367        assigned += scalar;
368        assigned -= b;
369        assert_equal(assigned, a * real + scalar - b, "assignments")?;
370        let x = Jet::<2>::variable(a.value, 0);
371        let y = Jet::<2>::variable(b.value, 1);
372        let combination = scalar * x + y;
373        prop_assert_eq!(combination.derivative, [scalar, Complex::new(1.0, 0.0)]);
374        Ok(())
375    }
376
377    /// Division undoes multiplication, the square root squares back off the branch cut,
378    /// `exp(a) exp(-a)` is the constant 1, and integer powers are repeated products.
379    fn check_inverses(a: Jet<2>, b: Jet<2>, n: i32) -> Result<(), TestCaseError> {
380        let one = Jet::<2>::constant(1.0);
381        // Rounding in the quotient rule grows with |b'| / |b| relative to a.
382        let scale = a.norm() * (1.0 + b.norm() / b.value.norm());
383        prop_assert_close!(((a * b) / b).value, a.value, 1e-14 * scale);
384        prop_assert_close!(((a * b) / b).derivative, a.derivative, 1e-14 * scale);
385        let root = if a.value.re > 0.0 { a } else { -a }.sqrt();
386        let square = root * root;
387        let expected = if a.value.re > 0.0 { a } else { -a };
388        prop_assert_close!(square.value, expected.value, 1e-14 * a.norm());
389        prop_assert_close!(square.derivative, expected.derivative, 1e-14 * a.norm());
390        let product = a.exp() * (-a).exp();
391        prop_assert_close!(product.value, one.value, 1e-14);
392        prop_assert_close!(product.derivative, one.derivative, 1e-14 * a.norm());
393        let repeated = (0..n.abs()).fold(one, |power, _| power * a);
394        let expected = if n < 0 { 1.0 / repeated } else { repeated };
395        let power = a.powi(n);
396        let scale = expected.norm() * f64::from(1 + n.abs());
397        prop_assert_close!(power.value, expected.value, 1e-14 * scale, "n = {}", n);
398        prop_assert_close!(
399            power.derivative,
400            expected.derivative,
401            1e-14 * scale,
402            "n = {}",
403            n
404        );
405        Ok(())
406    }
407
408    /// The derivatives of `sqrt(x^n e^x / (1 + x^2)) - 3/x + y x` in both variables equal
409    /// their closed forms; the square root's branch follows `complex_sqrt`.
410    fn check_chain_rule(z: Complex, w: Complex, n: i32) -> Result<(), TestCaseError> {
411        let i = Complex::i();
412        prop_assume!((z - i).norm() > 0.2 && (z + i).norm() > 0.2);
413        let (x, y) = (Jet::<2>::variable(z, 0), Jet::<2>::variable(w, 1));
414        let f = ((x.powi(n) * x.exp()) / (1.0 + x * x)).sqrt() - 3.0 / x + y * x;
415        let g = z.powi(n) * z.exp() / (1.0 + z * z);
416        let root = complex_sqrt(g);
417        let dg = g * (f64::from(n) / z + 1.0 - 2.0 * z / (1.0 + z * z));
418        let expected = [dg / (2.0 * root) + 3.0 / (z * z) + w, z];
419        let scale = (dg / root).norm() + 3.0 / z.norm_sqr() + w.norm() + z.norm();
420        prop_assert_close!(f.value, root - 3.0 / z + w * z, 1e-13 * scale);
421        prop_assert_close!(f.derivative, expected, 1e-13 * scale);
422        Ok(())
423    }
424
425    /// The quotient rule is exact when numerator and denominator share an extreme scale.
426    fn check_local_quotient_common_scale(a: Complex, scale: f64) -> Result<(), TestCaseError> {
427        let b = Complex::new(1.3, 0.4);
428        let result = Jet::<2>::variable(a * scale, 0) / Jet::<2>::variable(b * scale, 1);
429        prop_assert_close!(result.value, a / b, 1e-14);
430        prop_assert_close!(
431            result.derivative[0] * scale,
432            Complex::new(1.0, 0.0) / b,
433            1e-14
434        );
435        prop_assert_close!(result.derivative[1] * scale, -a / b / b, 1e-14);
436        Ok(())
437    }
438}
treams_core/numerics/jet.rs

treams_core/numerics/
jet.rs
