copp\path/
autodiff.rs

1//! Third-order forward-mode automatic differentiation primitives.
2//!
3//! [`Jet3`] carries value and derivatives up to the 3rd order with respect
4//! to a single scalar seed variable, and supports common arithmetic plus
5//! elementary functions ([`sin`], [`cos`], [`exp`], [`ln`], [`sqrt`], [`powi`]).
6
7use std::ops::{Add, Div, Mul, Neg, Sub};
8
9/// Third-order forward-mode automatic differentiation scalar.
10#[derive(Clone, Copy, Debug, Default)]
11pub struct Jet3 {
12    /// Function value.
13    pub v: f64,
14    /// First derivative with respect to the seed variable.
15    pub d1: f64,
16    /// Second derivative with respect to the seed variable.
17    pub d2: f64,
18    /// Third derivative with respect to the seed variable.
19    pub d3: f64,
20}
21
22impl Jet3 {
23    /// Construct a constant scalar with zero derivatives.
24    #[inline(always)]
25    pub fn constant(v: f64) -> Self {
26        Self {
27            v,
28            d1: 0.0,
29            d2: 0.0,
30            d3: 0.0,
31        }
32    }
33
34    /// Construct an independent variable seed (`d/ds = 1`).
35    #[inline(always)]
36    pub fn seed(v: f64) -> Self {
37        Self {
38            v,
39            d1: 1.0,
40            d2: 0.0,
41            d3: 0.0,
42        }
43    }
44
45    /// Apply sine and propagate derivatives up to third order.
46    #[inline(always)]
47    pub fn sin(self) -> Self {
48        let sv = self.v.sin();
49        let cv = self.v.cos();
50        let d1sq = self.d1 * self.d1;
51        let d1 = cv * self.d1;
52        let d2 = -sv * d1sq + cv * self.d2;
53        let d3 = -cv * d1sq * self.d1 - 3.0 * sv * self.d1 * self.d2 + cv * self.d3;
54        Self { v: sv, d1, d2, d3 }
55    }
56
57    /// Apply cosine and propagate derivatives up to third order.
58    #[inline(always)]
59    pub fn cos(self) -> Self {
60        let sv = self.v.sin();
61        let cv = self.v.cos();
62        let d1sq = self.d1 * self.d1;
63        let d1 = -sv * self.d1;
64        let d2 = -cv * d1sq - sv * self.d2;
65        let d3 = sv * d1sq * self.d1 - 3.0 * cv * self.d1 * self.d2 - sv * self.d3;
66        Self { v: cv, d1, d2, d3 }
67    }
68
69    /// Apply exponential and propagate derivatives up to third order.
70    #[inline(always)]
71    pub fn exp(self) -> Self {
72        let ev = self.v.exp();
73        let d1sq = self.d1 * self.d1;
74        let d1 = ev * self.d1;
75        let d2 = ev * (d1sq + self.d2);
76        let d3 = ev * (d1sq * self.d1 + 3.0 * self.d1 * self.d2 + self.d3);
77        Self { v: ev, d1, d2, d3 }
78    }
79
80    /// Apply natural logarithm and propagate derivatives up to third order.
81    #[inline(always)]
82    pub fn ln(self) -> Self {
83        let v = self.v.ln();
84        let inv_x = 1.0 / self.v;
85        let inv_x2 = inv_x * inv_x;
86        let inv_x3 = inv_x2 * inv_x;
87        let d1sq = self.d1 * self.d1;
88        let d1 = inv_x * self.d1;
89        let d2 = -inv_x2 * d1sq + inv_x * self.d2;
90        let d3 = 2.0 * inv_x3 * d1sq * self.d1 - 3.0 * inv_x2 * self.d1 * self.d2 + inv_x * self.d3;
91        Self { v, d1, d2, d3 }
92    }
93
94    /// Apply square root and propagate derivatives up to third order.
95    ///
96    /// The value component should be positive for finite real derivatives.
97    #[inline(always)]
98    pub fn sqrt(self) -> Self {
99        let sqrtv = self.v.sqrt();
100        let inv_sqrt = 1.0 / sqrtv;
101        let inv_v_sqrt = inv_sqrt / self.v;
102        let inv_v2_sqrt = inv_v_sqrt / self.v;
103        let d1sq = self.d1 * self.d1;
104        let d1 = 0.5 * inv_sqrt * self.d1;
105        let d2 = -0.25 * inv_v_sqrt * d1sq + 0.5 * inv_sqrt * self.d2;
106        let d3 = 0.375 * inv_v2_sqrt * d1sq * self.d1 - 0.75 * inv_v_sqrt * self.d1 * self.d2
107            + 0.5 * inv_sqrt * self.d3;
108        Self {
109            v: sqrtv,
110            d1,
111            d2,
112            d3,
113        }
114    }
115
116    /// Raise the value to an integer power and propagate derivatives.
117    #[inline(always)]
118    pub fn powi(self, n: i32) -> Self {
119        if n == 0 {
120            return Self::constant(1.0);
121        }
122        let nf = n as f64;
123        // Compute v^(n-3) once; derive v^(n-2), v^(n-1), v^n by repeated multiplication.
124        // This replaces 4 independent `f64::powi` calls with 1 call + 3 multiplications,
125        // avoiding redundant repeated-squaring work for the same base.
126        let vn3 = self.v.powi(n - 3);
127        let vn2 = vn3 * self.v;
128        let vn1 = vn2 * self.v;
129        let vn = vn1 * self.v;
130        let dv = nf * vn1;
131        let ddv = nf * (nf - 1.0) * vn2;
132        let dddv = nf * (nf - 1.0) * (nf - 2.0) * vn3;
133        let d1sq = self.d1 * self.d1;
134        let d1 = dv * self.d1;
135        let d2 = ddv * d1sq + dv * self.d2;
136        let d3 = dddv * d1sq * self.d1 + 3.0 * ddv * self.d1 * self.d2 + dv * self.d3;
137        Self { v: vn, d1, d2, d3 }
138    }
139
140    #[inline(always)]
141    fn inv(self) -> Self {
142        let v = 1.0 / self.v;
143        let inv_x2 = v * v;
144        let inv_x3 = inv_x2 * v;
145        let inv_x4 = inv_x3 * v;
146        let d1sq = self.d1 * self.d1;
147        let d1 = -inv_x2 * self.d1;
148        let d2 = 2.0 * inv_x3 * d1sq - inv_x2 * self.d2;
149        let d3 =
150            -6.0 * inv_x4 * d1sq * self.d1 + 6.0 * inv_x3 * self.d1 * self.d2 - inv_x2 * self.d3;
151        Self { v, d1, d2, d3 }
152    }
153}
154
155/// Function-style wrapper for [`Jet3::sin`].
156#[inline(always)]
157pub fn sin(x: Jet3) -> Jet3 {
158    x.sin()
159}
160
161/// Function-style wrapper for [`Jet3::cos`].
162#[inline(always)]
163pub fn cos(x: Jet3) -> Jet3 {
164    x.cos()
165}
166
167/// Function-style wrapper for [`Jet3::exp`].
168#[inline(always)]
169pub fn exp(x: Jet3) -> Jet3 {
170    x.exp()
171}
172
173/// Function-style wrapper for [`Jet3::ln`].
174#[inline(always)]
175pub fn ln(x: Jet3) -> Jet3 {
176    x.ln()
177}
178
179/// Function-style wrapper for [`Jet3::sqrt`].
180#[inline(always)]
181pub fn sqrt(x: Jet3) -> Jet3 {
182    x.sqrt()
183}
184
185/// Function-style wrapper for [`Jet3::powi`].
186#[inline(always)]
187pub fn powi(x: Jet3, n: i32) -> Jet3 {
188    x.powi(n)
189}
190
191impl From<f64> for Jet3 {
192    #[inline(always)]
193    fn from(value: f64) -> Self {
194        Self::constant(value)
195    }
196}
197
198impl Add for Jet3 {
199    type Output = Self;
200
201    #[inline(always)]
202    fn add(self, rhs: Self) -> Self::Output {
203        Self {
204            v: self.v + rhs.v,
205            d1: self.d1 + rhs.d1,
206            d2: self.d2 + rhs.d2,
207            d3: self.d3 + rhs.d3,
208        }
209    }
210}
211
212impl Add<f64> for Jet3 {
213    type Output = Self;
214
215    #[inline(always)]
216    fn add(self, rhs: f64) -> Self::Output {
217        Self {
218            v: self.v + rhs,
219            ..self
220        }
221    }
222}
223
224impl Add<Jet3> for f64 {
225    type Output = Jet3;
226
227    #[inline(always)]
228    fn add(self, rhs: Jet3) -> Self::Output {
229        rhs + self
230    }
231}
232
233impl Sub for Jet3 {
234    type Output = Self;
235
236    #[inline(always)]
237    fn sub(self, rhs: Self) -> Self::Output {
238        Self {
239            v: self.v - rhs.v,
240            d1: self.d1 - rhs.d1,
241            d2: self.d2 - rhs.d2,
242            d3: self.d3 - rhs.d3,
243        }
244    }
245}
246
247impl Sub<f64> for Jet3 {
248    type Output = Self;
249
250    #[inline(always)]
251    fn sub(self, rhs: f64) -> Self::Output {
252        Self {
253            v: self.v - rhs,
254            ..self
255        }
256    }
257}
258
259impl Sub<Jet3> for f64 {
260    type Output = Jet3;
261
262    #[inline(always)]
263    fn sub(self, rhs: Jet3) -> Self::Output {
264        Jet3::constant(self) - rhs
265    }
266}
267
268impl Mul for Jet3 {
269    type Output = Self;
270
271    #[inline(always)]
272    fn mul(self, rhs: Self) -> Self::Output {
273        Self {
274            v: self.v * rhs.v,
275            d1: self.d1 * rhs.v + self.v * rhs.d1,
276            d2: self.d2 * rhs.v + 2.0 * self.d1 * rhs.d1 + self.v * rhs.d2,
277            d3: self.d3 * rhs.v + 3.0 * (self.d2 * rhs.d1 + self.d1 * rhs.d2) + self.v * rhs.d3,
278        }
279    }
280}
281
282impl Mul<f64> for Jet3 {
283    type Output = Self;
284
285    #[inline(always)]
286    fn mul(self, rhs: f64) -> Self::Output {
287        Self {
288            v: self.v * rhs,
289            d1: self.d1 * rhs,
290            d2: self.d2 * rhs,
291            d3: self.d3 * rhs,
292        }
293    }
294}
295
296impl Mul<Jet3> for f64 {
297    type Output = Jet3;
298
299    #[inline(always)]
300    fn mul(self, rhs: Jet3) -> Self::Output {
301        rhs * self
302    }
303}
304
305impl Div for Jet3 {
306    type Output = Self;
307
308    #[allow(clippy::suspicious_arithmetic_impl)]
309    #[inline(always)]
310    fn div(self, rhs: Self) -> Self::Output {
311        self * rhs.inv()
312    }
313}
314
315impl Div<f64> for Jet3 {
316    type Output = Self;
317
318    #[inline(always)]
319    fn div(self, rhs: f64) -> Self::Output {
320        self * (1.0 / rhs)
321    }
322}
323
324impl Div<Jet3> for f64 {
325    type Output = Jet3;
326
327    #[inline(always)]
328    fn div(self, rhs: Jet3) -> Self::Output {
329        Jet3::constant(self) / rhs
330    }
331}
332
333impl Neg for Jet3 {
334    type Output = Self;
335
336    #[inline(always)]
337    fn neg(self) -> Self::Output {
338        Self {
339            v: -self.v,
340            d1: -self.d1,
341            d2: -self.d2,
342            d3: -self.d3,
343        }
344    }
345}
copp\path/autodiff.rs

copp\path/
autodiff.rs
