rustyqlib/core/fd_solvers/
axis_operator.rs

1//! Per-axis tridiagonal operators on tensor-product grids — the building
2//! block for 1-D, 2-D and 3-D finite-difference schemes.
3//!
4//! A discretized diffusion operator splits into one tridiagonal operator
5//! per spatial axis (plus an explicit mixed-derivative part). Each
6//! [`AxisOperator`] stores per-node coefficient triples, so coefficients
7//! may vary over the whole grid — exactly what local vol (1-D) or Heston
8//! (2-D in spot x variance) discretizations produce. The ADI time
9//! steppers in [`adi`](super::adi) consume these directly.
10
11use super::tridiagonal::thomas_algorithm;
12
13/// A tensor-product grid: node counts per axis, flattened row-major (the
14/// last axis is contiguous). Designed and tested for 1 to 3 dimensions.
15#[derive(Debug, Clone)]
16pub struct TensorGrid {
17    dims: Vec<usize>,
18}
19
20impl TensorGrid {
21    pub fn new(dims: &[usize]) -> Self {
22        assert!(!dims.is_empty() && dims.len() <= 3, "1 to 3 axes supported");
23        assert!(dims.iter().all(|&n| n >= 1), "every axis needs at least one node");
24        Self { dims: dims.to_vec() }
25    }
26
27    /// Total number of nodes.
28    pub fn len(&self) -> usize {
29        self.dims.iter().product()
30    }
31
32    pub fn is_empty(&self) -> bool {
33        false // dims are validated >= 1
34    }
35
36    pub fn ndim(&self) -> usize {
37        self.dims.len()
38    }
39
40    pub fn dims(&self) -> &[usize] {
41        &self.dims
42    }
43
44    /// Flat-index stride of each axis (row-major).
45    pub fn strides(&self) -> Vec<usize> {
46        let mut s = vec![1; self.dims.len()];
47        for k in (0..self.dims.len().saturating_sub(1)).rev() {
48            s[k] = s[k + 1] * self.dims[k + 1];
49        }
50        s
51    }
52
53    /// Flat index of a multi-index.
54    pub fn index(&self, idx: &[usize]) -> usize {
55        assert_eq!(idx.len(), self.dims.len());
56        idx.iter().zip(self.strides()).map(|(i, s)| i * s).sum()
57    }
58
59    /// Flat indices of the first node of every grid line along `axis`.
60    fn line_starts(&self, axis: usize) -> Vec<usize> {
61        let stride = self.strides()[axis];
62        let n = self.dims[axis];
63        (0..self.len()).filter(|i| (i / stride) % n == 0).collect()
64    }
65}
66
67/// A tridiagonal operator along one axis of a [`TensorGrid`], with
68/// per-node coefficients: row `i` of `A u` reads
69/// `sub[i] * u[i - stride] + diag[i] * u[i] + sup[i] * u[i + stride]`.
70///
71/// `sub` must be zero on the first plane of the axis and `sup` on the
72/// last (there is no neighbor there); boundary conditions are whatever
73/// the boundary rows encode — an all-zero row holds the boundary value
74/// fixed through both explicit application and implicit solves.
75#[derive(Debug, Clone)]
76pub struct AxisOperator {
77    pub axis: usize,
78    pub sub: Vec<f64>,
79    pub diag: Vec<f64>,
80    pub sup: Vec<f64>,
81}
82
83impl AxisOperator {
84    /// An all-zero operator along `axis` (a starting point to fill in).
85    pub fn zero(grid: &TensorGrid, axis: usize) -> Self {
86        assert!(axis < grid.ndim());
87        let n = grid.len();
88        Self { axis, sub: vec![0.0; n], diag: vec![0.0; n], sup: vec![0.0; n] }
89    }
90
91    /// `A u`.
92    pub fn apply(&self, grid: &TensorGrid, u: &[f64]) -> Vec<f64> {
93        let stride = grid.strides()[self.axis];
94        let n = grid.dims()[self.axis];
95        assert_eq!(u.len(), grid.len());
96        (0..grid.len())
97            .map(|i| {
98                let j = (i / stride) % n;
99                let mut v = self.diag[i] * u[i];
100                if j > 0 {
101                    v += self.sub[i] * u[i - stride];
102                }
103                if j < n - 1 {
104                    v += self.sup[i] * u[i + stride];
105                }
106                v
107            })
108            .collect()
109    }
110
111    /// Solve `(I - c A) x = rhs`, line by line with the Thomas algorithm —
112    /// the implicit stage of theta and ADI schemes.
113    pub fn solve_shifted(&self, grid: &TensorGrid, c: f64, rhs: &[f64]) -> Vec<f64> {
114        let stride = grid.strides()[self.axis];
115        let n = grid.dims()[self.axis];
116        assert_eq!(rhs.len(), grid.len());
117        let mut x = rhs.to_vec();
118        if n == 1 {
119            for (xi, &r) in x.iter_mut().zip(rhs) {
120                *xi = r / (1.0 - c * self.diag[0]);
121            }
122            return x;
123        }
124        let mut a = vec![0.0; n - 1];
125        let mut b = vec![0.0; n];
126        let mut cc = vec![0.0; n - 1];
127        let mut d = vec![0.0; n];
128        for start in grid.line_starts(self.axis) {
129            for j in 0..n {
130                let i = start + j * stride;
131                b[j] = 1.0 - c * self.diag[i];
132                d[j] = rhs[i];
133                if j > 0 {
134                    a[j - 1] = -c * self.sub[i];
135                }
136                if j < n - 1 {
137                    cc[j] = -c * self.sup[i];
138                }
139            }
140            let line = thomas_algorithm(&a, &b, &cc, &d);
141            for (j, v) in line.into_iter().enumerate() {
142                x[start + j * stride] = v;
143            }
144        }
145        x
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    /// 1-D Laplacian with Dirichlet (zero) boundary rows on [0, 1].
154    pub fn laplacian_1d(grid: &TensorGrid, axis: usize, h: f64) -> AxisOperator {
155        let mut op = AxisOperator::zero(grid, axis);
156        let stride = grid.strides()[axis];
157        let n = grid.dims()[axis];
158        for i in 0..grid.len() {
159            let j = (i / stride) % n;
160            if j > 0 && j < n - 1 {
161                op.sub[i] = 1.0 / (h * h);
162                op.diag[i] = -2.0 / (h * h);
163                op.sup[i] = 1.0 / (h * h);
164            }
165        }
166        op
167    }
168
169    #[test]
170    fn strides_and_indexing_are_row_major() {
171        let g = TensorGrid::new(&[3, 4, 5]);
172        assert_eq!(g.strides(), vec![20, 5, 1]);
173        assert_eq!(g.index(&[1, 2, 3]), 33);
174        assert_eq!(g.len(), 60);
175    }
176
177    #[test]
178    fn apply_matches_dense_stencil_2d() {
179        // second difference along axis 1 of a 2-D grid, checked by hand
180        let g = TensorGrid::new(&[2, 4]);
181        let h = 1.0;
182        let op = laplacian_1d(&g, 1, h);
183        let u: Vec<f64> = (0..8).map(|i| (i * i) as f64).collect();
184        let au = op.apply(&g, &u);
185        // row 0: nodes 0..4 with u = [0,1,4,9]: interior j=1 -> 0-2+4=2, j=2 -> 1-8+9=2
186        assert_eq!(&au[0..4], &[0.0, 2.0, 2.0, 0.0]);
187        // row 1: u = [16,25,36,49]: j=1 -> 16-50+36=2, j=2 -> 25-72+49=2
188        assert_eq!(&au[4..8], &[0.0, 2.0, 2.0, 0.0]);
189    }
190
191    #[test]
192    fn solve_shifted_inverts_apply() {
193        // x solves (I - cA) x = rhs  <=>  rhs = x - c A x
194        let g = TensorGrid::new(&[3, 5, 4]);
195        let op = laplacian_1d(&g, 1, 0.25);
196        let rhs: Vec<f64> = (0..g.len()).map(|i| ((i % 7) as f64) - 3.0).collect();
197        let c = 0.37;
198        let x = op.solve_shifted(&g, c, &rhs);
199        let ax = op.apply(&g, &x);
200        for i in 0..g.len() {
201            let back = x[i] - c * ax[i];
202            assert!((back - rhs[i]).abs() < 1e-10, "node {i}: {back} vs {}", rhs[i]);
203        }
204    }
205}
rustyqlib/core/fd_solvers/axis_operator.rs

rustyqlib/core/fd_solvers/
axis_operator.rs
