rustyqlib/core/aad/
tape.rs

1//! The AAD tape (Wengert list) and the backward sweep.
2//!
3//! Every arithmetic operation on [`Var`](super::var::Var) records one
4//! node holding the indices of its (up to two) parents and the local
5//! partial derivatives with respect to them. The backward sweep walks
6//! the tape once in reverse, accumulating adjoints — so the gradient of
7//! one scalar output with respect to **every** input costs one forward
8//! evaluation plus one reverse pass, independent of the number of
9//! inputs. That O(1) property is the whole point for Greeks.
10
11use std::cell::RefCell;
12
13/// Sentinel parent index for leaf slots.
14const NONE: usize = usize::MAX;
15
16#[derive(Debug, Clone, Copy)]
17pub(crate) struct Node {
18    pub(crate) parents: [usize; 2],
19    pub(crate) partials: [f64; 2],
20}
21
22/// The recording tape. Create one per differentiated computation (or
23/// [`clear`](Tape::clear) between computations, e.g. per Monte Carlo
24/// path).
25#[derive(Debug, Default)]
26pub struct Tape {
27    pub(crate) nodes: RefCell<Vec<Node>>,
28}
29
30impl Tape {
31    pub fn new() -> Tape {
32        Tape::default()
33    }
34
35    /// A new independent input variable.
36    pub fn var(&self, value: f64) -> super::var::Var<'_> {
37        let idx = self.push0();
38        super::var::Var { tape: self, idx, val: value }
39    }
40
41    /// Number of recorded nodes.
42    pub fn len(&self) -> usize {
43        self.nodes.borrow().len()
44    }
45
46    pub fn is_empty(&self) -> bool {
47        self.nodes.borrow().is_empty()
48    }
49
50    /// Drop all recorded nodes (existing `Var`s become invalid).
51    pub fn clear(&self) {
52        self.nodes.borrow_mut().clear();
53    }
54
55    pub(crate) fn push0(&self) -> usize {
56        let mut nodes = self.nodes.borrow_mut();
57        nodes.push(Node { parents: [NONE, NONE], partials: [0.0, 0.0] });
58        nodes.len() - 1
59    }
60
61    pub(crate) fn push1(&self, parent: usize, partial: f64) -> usize {
62        let mut nodes = self.nodes.borrow_mut();
63        nodes.push(Node { parents: [parent, NONE], partials: [partial, 0.0] });
64        nodes.len() - 1
65    }
66
67    pub(crate) fn push2(&self, p0: usize, w0: f64, p1: usize, w1: f64) -> usize {
68        let mut nodes = self.nodes.borrow_mut();
69        nodes.push(Node { parents: [p0, p1], partials: [w0, w1] });
70        nodes.len() - 1
71    }
72
73    /// Backward sweep from the node `output`: returns the adjoint of
74    /// every node, i.e. `d output / d node`.
75    pub(crate) fn backward(&self, output: usize) -> Vec<f64> {
76        let nodes = self.nodes.borrow();
77        let mut adjoint = vec![0.0; nodes.len()];
78        adjoint[output] = 1.0;
79        for i in (0..=output).rev() {
80            let a = adjoint[i];
81            if a == 0.0 {
82                continue;
83            }
84            let node = nodes[i];
85            for slot in 0..2 {
86                let p = node.parents[slot];
87                if p != NONE {
88                    adjoint[p] += node.partials[slot] * a;
89                }
90            }
91        }
92        adjoint
93    }
94}
95
96/// The result of a backward sweep: query with the input `Var`s.
97#[derive(Debug, Clone)]
98pub struct Gradients {
99    pub(crate) adjoints: Vec<f64>,
100}
101
102impl Gradients {
103    /// `d output / d v`.
104    pub fn wrt(&self, v: super::var::Var<'_>) -> f64 {
105        self.adjoints[v.idx]
106    }
107}
rustyqlib/core/aad/tape.rs

rustyqlib/core/aad/
tape.rs
