Skip to main content

qualia_core_db/solvers/calculus/
differential.rs

1//! Numerical and automatic differential calculus.
2
3use super::analysis::{AnalysisError, Complex64, LinearMap, Vector};
4
5#[repr(C)]
6#[derive(Debug, Clone, Copy, PartialEq)]
7pub struct DerivativeEstimate {
8    pub value: f64,
9    pub absolute_error: f64,
10    pub step: f64,
11}
12
13pub fn adaptive_central_difference<F>(
14    function: F,
15    x: f64,
16) -> Result<DerivativeEstimate, AnalysisError>
17where
18    F: Fn(f64) -> f64,
19{
20    if !x.is_finite() {
21        return Err(AnalysisError::NonFinite);
22    }
23    let h = f64::EPSILON.cbrt() * (1.0 + x.abs());
24    let coarse = central_difference(&function, x, h)?;
25    let fine = central_difference(&function, x, h * 0.5)?;
26    Ok(DerivativeEstimate {
27        value: fine + (fine - coarse) / 3.0,
28        absolute_error: (fine - coarse).abs() / 3.0,
29        step: h * 0.5,
30    })
31}
32
33fn central_difference<F>(function: &F, x: f64, h: f64) -> Result<f64, AnalysisError>
34where
35    F: Fn(f64) -> f64,
36{
37    let high = function(x + h);
38    let low = function(x - h);
39    if !high.is_finite() || !low.is_finite() {
40        return Err(AnalysisError::InvalidDomain);
41    }
42    Ok((high - low) / (2.0 * h))
43}
44
45pub fn complex_step_derivative<F>(function: F, x: f64) -> Result<DerivativeEstimate, AnalysisError>
46where
47    F: Fn(Complex64) -> Complex64,
48{
49    if !x.is_finite() {
50        return Err(AnalysisError::NonFinite);
51    }
52    let h = 1e-20;
53    let value = function(Complex64::new(x, h));
54    if !value.re.is_finite() || !value.im.is_finite() {
55        return Err(AnalysisError::InvalidDomain);
56    }
57    Ok(DerivativeEstimate {
58        value: value.im / h,
59        absolute_error: h,
60        step: h,
61    })
62}
63
64pub fn jacobian<const M: usize, const N: usize, F>(
65    function: F,
66    point: Vector<N>,
67) -> Result<LinearMap<M, N>, AnalysisError>
68where
69    F: Fn(Vector<N>) -> Vector<M>,
70{
71    point.validate()?;
72    let mut coefficients = [[0.0; N]; M];
73    for column in 0..N {
74        let h = f64::EPSILON.cbrt() * (1.0 + point.data[column].abs());
75        let mut high = point;
76        let mut low = point;
77        high.data[column] += h;
78        low.data[column] -= h;
79        let high_value = function(high);
80        let low_value = function(low);
81        high_value.validate()?;
82        low_value.validate()?;
83        for row in 0..M {
84            coefficients[row][column] = (high_value.data[row] - low_value.data[row]) / (2.0 * h);
85        }
86    }
87    Ok(LinearMap::new(coefficients))
88}
89
90pub fn jvp<const M: usize, const N: usize>(
91    derivative: &LinearMap<M, N>,
92    direction: Vector<N>,
93) -> Result<Vector<M>, AnalysisError> {
94    derivative.apply(direction)
95}
96
97pub fn vjp<const M: usize, const N: usize>(
98    derivative: &LinearMap<M, N>,
99    cotangent: Vector<M>,
100) -> Result<Vector<N>, AnalysisError> {
101    derivative.transpose().apply(cotangent)
102}
103
104#[repr(C)]
105#[derive(Debug, Clone, Copy, PartialEq)]
106pub struct Dual<const N: usize> {
107    pub value: f64,
108    pub derivative: [f64; N],
109}
110
111impl<const N: usize> Dual<N> {
112    pub const fn constant(value: f64) -> Self {
113        Self {
114            value,
115            derivative: [0.0; N],
116        }
117    }
118
119    pub fn variable(value: f64, index: usize) -> Result<Self, AnalysisError> {
120        if index >= N || !value.is_finite() {
121            return Err(AnalysisError::InvalidDomain);
122        }
123        let mut derivative = [0.0; N];
124        derivative[index] = 1.0;
125        Ok(Self { value, derivative })
126    }
127
128    pub fn add(self, other: Self) -> Self {
129        let mut derivative = [0.0; N];
130        for (index, value) in derivative.iter_mut().enumerate() {
131            *value = self.derivative[index] + other.derivative[index];
132        }
133        Self {
134            value: self.value + other.value,
135            derivative,
136        }
137    }
138
139    pub fn mul(self, other: Self) -> Self {
140        let mut derivative = [0.0; N];
141        for (index, value) in derivative.iter_mut().enumerate() {
142            *value = self.derivative[index] * other.value + self.value * other.derivative[index];
143        }
144        Self {
145            value: self.value * other.value,
146            derivative,
147        }
148    }
149
150    pub fn sin(self) -> Self {
151        let scale = self.value.cos();
152        let mut derivative = self.derivative;
153        for value in &mut derivative {
154            *value *= scale;
155        }
156        Self {
157            value: self.value.sin(),
158            derivative,
159        }
160    }
161
162    pub fn exp(self) -> Self {
163        let value = self.value.exp();
164        let mut derivative = self.derivative;
165        for component in &mut derivative {
166            *component *= value;
167        }
168        Self { value, derivative }
169    }
170}
171
172pub fn hessian<const N: usize, F>(
173    function: F,
174    point: Vector<N>,
175) -> Result<LinearMap<N, N>, AnalysisError>
176where
177    F: Fn(Vector<N>) -> f64,
178{
179    point.validate()?;
180    let mut coefficients = [[0.0; N]; N];
181    let base = function(point);
182    if !base.is_finite() {
183        return Err(AnalysisError::InvalidDomain);
184    }
185    for row in 0..N {
186        let hr = f64::EPSILON.powf(0.25) * (1.0 + point.data[row].abs());
187        for column in row..N {
188            let hc = f64::EPSILON.powf(0.25) * (1.0 + point.data[column].abs());
189            let value = if row == column {
190                let mut high = point;
191                let mut low = point;
192                high.data[row] += hr;
193                low.data[row] -= hr;
194                (function(high) - 2.0 * base + function(low)) / (hr * hr)
195            } else {
196                let mut pp = point;
197                let mut pm = point;
198                let mut mp = point;
199                let mut mm = point;
200                pp.data[row] += hr;
201                pp.data[column] += hc;
202                pm.data[row] += hr;
203                pm.data[column] -= hc;
204                mp.data[row] -= hr;
205                mp.data[column] += hc;
206                mm.data[row] -= hr;
207                mm.data[column] -= hc;
208                (function(pp) - function(pm) - function(mp) + function(mm)) / (4.0 * hr * hc)
209            };
210            if !value.is_finite() {
211                return Err(AnalysisError::InvalidDomain);
212            }
213            coefficients[row][column] = value;
214            coefficients[column][row] = value;
215        }
216    }
217    Ok(LinearMap::new(coefficients))
218}
219
220#[repr(C)]
221#[derive(Debug, Clone, Copy, PartialEq)]
222pub struct NewtonReport<const N: usize> {
223    pub solution: Vector<N>,
224    pub iterations: u32,
225    pub residual_norm: f64,
226    pub step_norm: f64,
227}
228
229pub fn damped_newton<const N: usize, F>(
230    function: F,
231    initial: Vector<N>,
232    tolerance: f64,
233    max_iterations: u32,
234) -> Result<NewtonReport<N>, AnalysisError>
235where
236    F: Fn(Vector<N>) -> Vector<N>,
237{
238    if !tolerance.is_finite() || tolerance <= 0.0 {
239        return Err(AnalysisError::InvalidDomain);
240    }
241    initial.validate()?;
242    let mut point = initial;
243    let mut last_step = 0.0;
244    for iteration in 0..=max_iterations {
245        let residual = function(point);
246        residual.validate()?;
247        let residual_norm = residual.norm()?;
248        if residual_norm <= tolerance {
249            return Ok(NewtonReport {
250                solution: point,
251                iterations: iteration,
252                residual_norm,
253                step_norm: last_step,
254            });
255        }
256        if iteration == max_iterations {
257            return Err(AnalysisError::IterationLimit {
258                residual: residual_norm,
259            });
260        }
261        let derivative = jacobian(&function, point)?;
262        let step = derivative.solve(residual.scale(-1.0)?)?;
263        last_step = step.norm()?;
264
265        let mut damping = 1.0;
266        let mut accepted = false;
267        for _ in 0..20 {
268            let candidate = point.add(step.scale(damping)?);
269            let candidate_norm = function(candidate).norm()?;
270            if candidate_norm < residual_norm {
271                point = candidate;
272                accepted = true;
273                break;
274            }
275            damping *= 0.5;
276        }
277        if !accepted {
278            return Err(AnalysisError::IterationLimit {
279                residual: residual_norm,
280            });
281        }
282    }
283    unreachable!()
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    #[test]
291    fn adaptive_and_complex_step_derivatives_match_analytic_values() {
292        let central = adaptive_central_difference(|x| x.sin(), 0.4).unwrap();
293        assert!((central.value - 0.4_f64.cos()).abs() < 1e-10);
294
295        let complex = complex_step_derivative(|z| z.exp(), 0.4).unwrap();
296        assert!((complex.value - 0.4_f64.exp()).abs() < 1e-14);
297    }
298
299    #[test]
300    fn jvp_vjp_duality_holds() {
301        let derivative = LinearMap::new([[1.0, 2.0], [-3.0, 4.0], [0.5, -2.0]]);
302        let direction = Vector::new([2.0, -1.0]);
303        let cotangent = Vector::new([3.0, 0.25, -4.0]);
304        let left = cotangent.dot(jvp(&derivative, direction).unwrap()).unwrap();
305        let right = vjp(&derivative, cotangent).unwrap().dot(direction).unwrap();
306        assert!((left - right).abs() < 1e-14);
307    }
308
309    #[test]
310    fn forward_dual_gradient_matches_analytic_gradient() {
311        let x = Dual::<2>::variable(2.0, 0).unwrap();
312        let y = Dual::<2>::variable(0.3, 1).unwrap();
313        let value = x.mul(x).add(y.sin());
314        assert!((value.value - (4.0 + 0.3_f64.sin())).abs() < 1e-14);
315        assert!((value.derivative[0] - 4.0).abs() < 1e-14);
316        assert!((value.derivative[1] - 0.3_f64.cos()).abs() < 1e-14);
317    }
318
319    #[test]
320    fn hessian_and_damped_newton_match_quadratic_oracle() {
321        let point = Vector::new([1.2, -0.7]);
322        let matrix = hessian(
323            |x: Vector<2>| {
324                3.0 * x.data[0] * x.data[0]
325                    + 2.0 * x.data[0] * x.data[1]
326                    + 4.0 * x.data[1] * x.data[1]
327            },
328            point,
329        )
330        .unwrap();
331        assert!((matrix.coefficients[0][0] - 6.0).abs() < 1e-6);
332        assert!((matrix.coefficients[0][1] - 2.0).abs() < 1e-6);
333        assert!((matrix.coefficients[1][1] - 8.0).abs() < 1e-6);
334
335        let report = damped_newton(
336            |x: Vector<2>| Vector::new([x.data[0] * x.data[0] - 2.0, x.data[1] - 3.0]),
337            Vector::new([1.0, 0.0]),
338            1e-12,
339            20,
340        )
341        .unwrap();
342        assert!((report.solution.data[0] - 2.0_f64.sqrt()).abs() < 1e-12);
343        assert!((report.solution.data[1] - 3.0).abs() < 1e-12);
344    }
345}