rustyqlib/core/aad/
var.rs

1//! The differentiable value type: `Var` records every operation on the
2//! tape via operator overloading, so pricing code written over `Var`
3//! looks like ordinary arithmetic.
4
5use std::ops::{Add, Div, Mul, Neg, Sub};
6
7use crate::core::utils::{norm_pdf, norm_cdf};
8
9use super::tape::{Gradients, Tape};
10
11/// A value recorded on an AAD [`Tape`]. Copyable and cheap; all
12/// arithmetic allocates one tape node.
13#[derive(Debug, Clone, Copy)]
14pub struct Var<'a> {
15    pub(crate) tape: &'a Tape,
16    pub(crate) idx: usize,
17    pub(crate) val: f64,
18}
19
20impl<'a> Var<'a> {
21    pub fn value(self) -> f64 {
22        self.val
23    }
24
25    /// Backward sweep: the gradient of `self` with respect to every
26    /// variable on the tape (query via [`Gradients::wrt`]).
27    pub fn grad(self) -> Gradients {
28        Gradients { adjoints: self.tape.backward(self.idx) }
29    }
30
31    fn unary(self, val: f64, partial: f64) -> Var<'a> {
32        Var { tape: self.tape, idx: self.tape.push1(self.idx, partial), val }
33    }
34
35    pub fn exp(self) -> Var<'a> {
36        let e = self.val.exp();
37        self.unary(e, e)
38    }
39
40    pub fn ln(self) -> Var<'a> {
41        self.unary(self.val.ln(), 1.0 / self.val)
42    }
43
44    pub fn sqrt(self) -> Var<'a> {
45        let s = self.val.sqrt();
46        self.unary(s, 0.5 / s)
47    }
48
49    pub fn powf(self, n: f64) -> Var<'a> {
50        self.unary(self.val.powf(n), n * self.val.powf(n - 1.0))
51    }
52
53    pub fn sin(self) -> Var<'a> {
54        self.unary(self.val.sin(), self.val.cos())
55    }
56
57    pub fn cos(self) -> Var<'a> {
58        self.unary(self.val.cos(), -self.val.sin())
59    }
60
61    /// Standard normal CDF (derivative: the density).
62    pub fn norm_cdf(self) -> Var<'a> {
63        self.unary(norm_cdf(self.val), norm_pdf(self.val))
64    }
65
66    /// `max(self, other)` with the one-sided subgradient at ties.
67    pub fn max(self, other: Var<'a>) -> Var<'a> {
68        if self.val >= other.val {
69            Var {
70                tape: self.tape,
71                idx: self.tape.push2(self.idx, 1.0, other.idx, 0.0),
72                val: self.val,
73            }
74        } else {
75            Var {
76                tape: self.tape,
77                idx: self.tape.push2(self.idx, 0.0, other.idx, 1.0),
78                val: other.val,
79            }
80        }
81    }
82
83    /// `min(self, other)` with the one-sided subgradient at ties.
84    pub fn min(self, other: Var<'a>) -> Var<'a> {
85        if self.val <= other.val {
86            Var {
87                tape: self.tape,
88                idx: self.tape.push2(self.idx, 1.0, other.idx, 0.0),
89                val: self.val,
90            }
91        } else {
92            Var {
93                tape: self.tape,
94                idx: self.tape.push2(self.idx, 0.0, other.idx, 1.0),
95                val: other.val,
96            }
97        }
98    }
99
100    /// `max(self, constant)` — the positive-part operator for payoffs
101    /// (`x.maxf(0.0)`), differentiable almost everywhere.
102    pub fn maxf(self, c: f64) -> Var<'a> {
103        if self.val >= c {
104            self.unary(self.val, 1.0)
105        } else {
106            self.unary(c, 0.0)
107        }
108    }
109}
110
111impl<'a> Add for Var<'a> {
112    type Output = Var<'a>;
113    fn add(self, rhs: Var<'a>) -> Var<'a> {
114        Var {
115            tape: self.tape,
116            idx: self.tape.push2(self.idx, 1.0, rhs.idx, 1.0),
117            val: self.val + rhs.val,
118        }
119    }
120}
121
122impl<'a> Sub for Var<'a> {
123    type Output = Var<'a>;
124    fn sub(self, rhs: Var<'a>) -> Var<'a> {
125        Var {
126            tape: self.tape,
127            idx: self.tape.push2(self.idx, 1.0, rhs.idx, -1.0),
128            val: self.val - rhs.val,
129        }
130    }
131}
132
133impl<'a> Mul for Var<'a> {
134    type Output = Var<'a>;
135    fn mul(self, rhs: Var<'a>) -> Var<'a> {
136        Var {
137            tape: self.tape,
138            idx: self.tape.push2(self.idx, rhs.val, rhs.idx, self.val),
139            val: self.val * rhs.val,
140        }
141    }
142}
143
144impl<'a> Div for Var<'a> {
145    type Output = Var<'a>;
146    fn div(self, rhs: Var<'a>) -> Var<'a> {
147        let v = self.val / rhs.val;
148        Var {
149            tape: self.tape,
150            idx: self.tape.push2(self.idx, 1.0 / rhs.val, rhs.idx, -v / rhs.val),
151            val: v,
152        }
153    }
154}
155
156impl<'a> Neg for Var<'a> {
157    type Output = Var<'a>;
158    fn neg(self) -> Var<'a> {
159        self.unary(-self.val, -1.0)
160    }
161}
162
163// mixed Var / f64 arithmetic
164impl<'a> Add<f64> for Var<'a> {
165    type Output = Var<'a>;
166    fn add(self, rhs: f64) -> Var<'a> {
167        self.unary(self.val + rhs, 1.0)
168    }
169}
170
171impl<'a> Add<Var<'a>> for f64 {
172    type Output = Var<'a>;
173    fn add(self, rhs: Var<'a>) -> Var<'a> {
174        rhs + self
175    }
176}
177
178impl<'a> Sub<f64> for Var<'a> {
179    type Output = Var<'a>;
180    fn sub(self, rhs: f64) -> Var<'a> {
181        self.unary(self.val - rhs, 1.0)
182    }
183}
184
185impl<'a> Sub<Var<'a>> for f64 {
186    type Output = Var<'a>;
187    fn sub(self, rhs: Var<'a>) -> Var<'a> {
188        rhs.unary(self - rhs.val, -1.0)
189    }
190}
191
192impl<'a> Mul<f64> for Var<'a> {
193    type Output = Var<'a>;
194    fn mul(self, rhs: f64) -> Var<'a> {
195        self.unary(self.val * rhs, rhs)
196    }
197}
198
199impl<'a> Mul<Var<'a>> for f64 {
200    type Output = Var<'a>;
201    fn mul(self, rhs: Var<'a>) -> Var<'a> {
202        rhs * self
203    }
204}
205
206impl<'a> Div<f64> for Var<'a> {
207    type Output = Var<'a>;
208    fn div(self, rhs: f64) -> Var<'a> {
209        self.unary(self.val / rhs, 1.0 / rhs)
210    }
211}
212
213impl<'a> Div<Var<'a>> for f64 {
214    type Output = Var<'a>;
215    fn div(self, rhs: Var<'a>) -> Var<'a> {
216        rhs.unary(self / rhs.val, -self / (rhs.val * rhs.val))
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn basic_rules_match_calculus() {
226        let tape = Tape::new();
227        let x = tape.var(1.3);
228        let y = tape.var(0.7);
229        // f = x*y + sin(x) + x/y
230        let f = x * y + x.sin() + x / y;
231        let g = f.grad();
232        assert!((g.wrt(x) - (0.7 + 1.3f64.cos() + 1.0 / 0.7)).abs() < 1e-14);
233        assert!((g.wrt(y) - (1.3 - 1.3 / (0.7 * 0.7))).abs() < 1e-14);
234    }
235
236    #[test]
237    fn fan_out_accumulates_adjoints() {
238        // x used twice: d(x*x)/dx = 2x
239        let tape = Tape::new();
240        let x = tape.var(3.0);
241        let g = (x * x).grad();
242        assert!((g.wrt(x) - 6.0).abs() < 1e-14);
243        // deep chain: exp(ln(sqrt(x^4))) = x^2
244        let h = x.powf(4.0).sqrt().ln().exp().grad();
245        assert!((h.wrt(x) - 6.0).abs() < 1e-12);
246    }
247
248    #[test]
249    fn composite_matches_finite_differences() {
250        let f_val = |x: f64, y: f64| ((x * y).exp() + (x / y).sqrt()).ln() * y.cos();
251        let tape = Tape::new();
252        let x = tape.var(0.8);
253        let y = tape.var(1.9);
254        let f = ((x * y).exp() + (x / y).sqrt()).ln() * y.cos();
255        assert!((f.value() - f_val(0.8, 1.9)).abs() < 1e-14);
256        let g = f.grad();
257        let h = 1e-6;
258        let fd_x = (f_val(0.8 + h, 1.9) - f_val(0.8 - h, 1.9)) / (2.0 * h);
259        let fd_y = (f_val(0.8, 1.9 + h) - f_val(0.8, 1.9 - h)) / (2.0 * h);
260        assert!((g.wrt(x) - fd_x).abs() < 1e-8, "{} vs {fd_x}", g.wrt(x));
261        assert!((g.wrt(y) - fd_y).abs() < 1e-8, "{} vs {fd_y}", g.wrt(y));
262    }
263
264    #[test]
265    fn positive_part_has_the_indicator_derivative() {
266        let tape = Tape::new();
267        let x = tape.var(2.0);
268        let up = (x - 1.0).maxf(0.0).grad();
269        assert!((up.wrt(x) - 1.0).abs() < 1e-14);
270        let x2 = tape.var(0.5);
271        let down = (x2 - 1.0).maxf(0.0).grad();
272        assert_eq!(down.wrt(x2), 0.0);
273        // two-variable max routes the adjoint to the winner
274        let a = tape.var(3.0);
275        let b = tape.var(4.0);
276        let g = (a.max(b) * 2.0).grad();
277        assert_eq!(g.wrt(a), 0.0);
278        assert!((g.wrt(b) - 2.0).abs() < 1e-14);
279    }
280
281    #[test]
282    fn norm_cdf_differentiates_to_the_density() {
283        let tape = Tape::new();
284        let x = tape.var(0.37);
285        let g = x.norm_cdf().grad();
286        assert!((g.wrt(x) - crate::core::utils::norm_pdf(0.37)).abs() < 1e-14);
287    }
288}
rustyqlib/core/aad/var.rs

rustyqlib/core/aad/
var.rs
