finitediff/vec/
hessian.rs

1// Copyright 2018-2024 argmin developers
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use std::ops::AddAssign;
9
10use anyhow::{anyhow, Error};
11use num::{Float, FromPrimitive};
12
13use crate::utils::{mod_and_calc, restore_symmetry_vec, KV};
14
15use super::{CostFn, GradientFn};
16
17pub fn forward_hessian_vec<F>(x: &Vec<F>, grad: GradientFn<'_, F>) -> Result<Vec<Vec<F>>, Error>
18where
19    F: Float + FromPrimitive,
20{
21    let eps_sqrt = F::epsilon().sqrt();
22    let fx = (grad)(x)?;
23    let mut xt = x.clone();
24    let out: Vec<Vec<F>> = (0..x.len())
25        .map(|i| {
26            let fx1 = mod_and_calc(&mut xt, grad, i, eps_sqrt)?;
27            Ok(fx1
28                .iter()
29                .zip(fx.iter())
30                .map(|(&a, &b)| (a - b) / eps_sqrt)
31                .collect::<Vec<F>>())
32        })
33        .collect::<Result<_, Error>>()?;
34
35    // restore symmetry
36    Ok(restore_symmetry_vec(out))
37}
38
39pub fn central_hessian_vec<F>(x: &[F], grad: GradientFn<'_, F>) -> Result<Vec<Vec<F>>, Error>
40where
41    F: Float + FromPrimitive,
42{
43    let eps_cbrt = F::epsilon().cbrt();
44    let mut xt = x.to_owned();
45    let out: Vec<Vec<F>> = (0..x.len())
46        .map(|i| {
47            let fx1 = mod_and_calc(&mut xt, grad, i, eps_cbrt)?;
48            let fx2 = mod_and_calc(&mut xt, grad, i, -eps_cbrt)?;
49            Ok(fx1
50                .iter()
51                .zip(fx2.iter())
52                .map(|(&a, &b)| (a - b) / (F::from_f64(2.0).unwrap() * eps_cbrt))
53                .collect::<Vec<F>>())
54        })
55        .collect::<Result<_, Error>>()?;
56
57    // restore symmetry
58    Ok(restore_symmetry_vec(out))
59}
60
61pub fn forward_hessian_vec_prod_vec<F>(
62    x: &Vec<F>,
63    grad: GradientFn<'_, F>,
64    p: &[F],
65) -> Result<Vec<F>, Error>
66where
67    F: Float,
68{
69    let eps_sqrt = F::epsilon().sqrt();
70    let fx = (grad)(x)?;
71    let out: Vec<F> = {
72        let x1 = x
73            .iter()
74            .zip(p.iter())
75            .map(|(&xi, &pi)| xi + pi * eps_sqrt)
76            .collect();
77        let fx1 = (grad)(&x1)?;
78        fx1.iter()
79            .zip(fx.iter())
80            .map(|(&a, &b)| (a - b) / eps_sqrt)
81            .collect::<Vec<F>>()
82    };
83    Ok(out)
84}
85
86pub fn central_hessian_vec_prod_vec<F>(
87    x: &[F],
88    grad: GradientFn<'_, F>,
89    p: &[F],
90) -> Result<Vec<F>, Error>
91where
92    F: Float + FromPrimitive,
93{
94    let eps_cbrt = F::epsilon().cbrt();
95    let out: Vec<F> = {
96        // TODO: Do this in single array
97        let x1 = x
98            .iter()
99            .zip(p.iter())
100            .map(|(&xi, &pi)| xi + pi * eps_cbrt)
101            .collect();
102        let x2 = x
103            .iter()
104            .zip(p.iter())
105            .map(|(&xi, &pi)| xi - pi * eps_cbrt)
106            .collect();
107        let fx1 = (grad)(&x1)?;
108        let fx2 = (grad)(&x2)?;
109        fx1.iter()
110            .zip(fx2.iter())
111            .map(|(&a, &b)| (a - b) / (F::from_f64(2.0).unwrap() * eps_cbrt))
112            .collect::<Vec<F>>()
113    };
114    Ok(out)
115}
116
117pub fn forward_hessian_nograd_vec<F>(x: &Vec<F>, f: CostFn<'_, F>) -> Result<Vec<Vec<F>>, Error>
118where
119    F: Float + FromPrimitive + AddAssign,
120{
121    // TODO: Check why this is necessary
122    let eps_nograd = F::from_f64(2.0).unwrap() * F::epsilon();
123    let eps_sqrt_nograd = eps_nograd.sqrt();
124
125    let fx = (f)(x)?;
126    let n = x.len();
127    let mut xt = x.clone();
128
129    // Precompute f(x + sqrt(EPS) * e_i) for all i
130    let fxei: Vec<F> = (0..n)
131        .map(|i| mod_and_calc(&mut xt, f, i, eps_sqrt_nograd))
132        .collect::<Result<_, Error>>()?;
133
134    let mut out: Vec<Vec<F>> = vec![vec![F::from_f64(0.0).unwrap(); n]; n];
135
136    for i in 0..n {
137        for j in 0..=i {
138            let t = {
139                let xti = xt[i];
140                let xtj = xt[j];
141                xt[i] += eps_sqrt_nograd;
142                xt[j] += eps_sqrt_nograd;
143                let fxij = (f)(&xt)?;
144                xt[i] = xti;
145                xt[j] = xtj;
146                (fxij - fxei[i] - fxei[j] + fx) / eps_nograd
147            };
148            out[i][j] = t;
149            out[j][i] = t;
150        }
151    }
152    Ok(out)
153}
154
155pub fn forward_hessian_nograd_sparse_vec<F>(
156    x: &Vec<F>,
157    f: CostFn<'_, F>,
158    indices: Vec<[usize; 2]>,
159) -> Result<Vec<Vec<F>>, Error>
160where
161    F: Float + FromPrimitive + AddAssign,
162{
163    // TODO: Check why this is necessary
164    let eps_nograd = F::from_f64(2.0).unwrap() * F::epsilon();
165    let eps_sqrt_nograd = eps_nograd.sqrt();
166
167    let fx = (f)(x)?;
168    let n = x.len();
169    let mut xt = x.clone();
170
171    let mut idxs: Vec<usize> = indices
172        .iter()
173        .flat_map(|i| i.iter())
174        .cloned()
175        .collect::<Vec<usize>>();
176    idxs.sort();
177    idxs.dedup();
178
179    let mut fxei = KV::new(idxs.len());
180
181    for idx in idxs.iter() {
182        fxei.set(*idx, mod_and_calc(&mut xt, f, *idx, eps_sqrt_nograd)?);
183    }
184
185    let mut out: Vec<Vec<F>> = vec![vec![F::from_f64(0.0).unwrap(); n]; n];
186    for [i, j] in indices {
187        let t = {
188            let xti = xt[i];
189            let xtj = xt[j];
190            xt[i] += eps_sqrt_nograd;
191            xt[j] += eps_sqrt_nograd;
192            let fxij = (f)(&xt)?;
193            xt[i] = xti;
194            xt[j] = xtj;
195
196            let fxi = fxei.get(i).ok_or(anyhow!("Bug"))?;
197            let fxj = fxei.get(j).ok_or(anyhow!("Bug"))?;
198            (fxij - fxi - fxj + fx) / eps_nograd
199        };
200        out[i][j] = t;
201        out[j][i] = t;
202    }
203    Ok(out)
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    const COMP_ACC: f64 = 1e-6;
211
212    fn f(x: &Vec<f64>) -> Result<f64, Error> {
213        Ok(x[0] + x[1].powi(2) + x[2] * x[3].powi(2))
214    }
215
216    fn g(x: &Vec<f64>) -> Result<Vec<f64>, Error> {
217        Ok(vec![1.0, 2.0 * x[1], x[3].powi(2), 2.0 * x[3] * x[2]])
218    }
219
220    fn x() -> Vec<f64> {
221        vec![1.0f64, 1.0, 1.0, 1.0]
222    }
223
224    fn p() -> Vec<f64> {
225        vec![2.0, 3.0, 4.0, 5.0]
226    }
227
228    fn res1() -> Vec<Vec<f64>> {
229        vec![
230            vec![0.0, 0.0, 0.0, 0.0],
231            vec![0.0, 2.0, 0.0, 0.0],
232            vec![0.0, 0.0, 0.0, 2.0],
233            vec![0.0, 0.0, 2.0, 2.0],
234        ]
235    }
236
237    fn res2() -> Vec<f64> {
238        vec![0.0, 6.0, 10.0, 18.0]
239    }
240
241    #[test]
242    fn test_forward_hessian_vec_f64() {
243        let hessian = forward_hessian_vec(&x(), &g).unwrap();
244        let res = res1();
245        // println!("hessian:\n{:#?}", hessian);
246        // println!("diff:\n{:#?}", diff);
247        for i in 0..4 {
248            for j in 0..4 {
249                assert!((res[i][j] - hessian[i][j]).abs() < COMP_ACC)
250            }
251        }
252    }
253
254    #[test]
255    fn test_central_hessian_vec_f64() {
256        let hessian = central_hessian_vec(&x(), &g).unwrap();
257        let res = res1();
258        // println!("hessian:\n{:#?}", hessian);
259        // println!("diff:\n{:#?}", diff);
260        for i in 0..4 {
261            for j in 0..4 {
262                assert!((res[i][j] - hessian[i][j]).abs() < COMP_ACC)
263            }
264        }
265    }
266
267    #[test]
268    fn test_forward_hessian_vec_prod_vec_f64() {
269        let hessian = forward_hessian_vec_prod_vec(&x(), &g, &p()).unwrap();
270        let res = res2();
271        // println!("hessian:\n{:#?}", hessian);
272        // println!("diff:\n{:#?}", diff);
273        for i in 0..4 {
274            assert!((res[i] - hessian[i]).abs() < COMP_ACC)
275        }
276    }
277
278    #[test]
279    fn test_central_hessian_vec_prod_vec_f64() {
280        let hessian = central_hessian_vec_prod_vec(&x(), &g, &p()).unwrap();
281        let res = res2();
282        // println!("hessian:\n{:#?}", hessian);
283        // println!("diff:\n{:#?}", diff);
284        for i in 0..4 {
285            assert!((res[i] - hessian[i]).abs() < COMP_ACC)
286        }
287    }
288
289    #[test]
290    fn test_forward_hessian_nograd_vec_f64() {
291        let hessian = forward_hessian_nograd_vec(&x(), &f).unwrap();
292        let res = res1();
293        // println!("hessian:\n{:#?}", hessian);
294        for i in 0..4 {
295            for j in 0..4 {
296                assert!((res[i][j] - hessian[i][j]).abs() < COMP_ACC)
297            }
298        }
299    }
300
301    #[test]
302    fn test_forward_hessian_nograd_sparse_vec_f64() {
303        let indices = vec![[1, 1], [2, 3], [3, 3]];
304        let hessian = forward_hessian_nograd_sparse_vec(&x(), &f, indices).unwrap();
305        let res = res1();
306        // println!("hessian:\n{:#?}", hessian);
307        // println!("diff:\n{:#?}", diff);
308        for i in 0..4 {
309            for j in 0..4 {
310                assert!((res[i][j] - hessian[i][j]).abs() < COMP_ACC)
311            }
312        }
313    }
314}
finitediff/vec/hessian.rs

finitediff/vec/
hessian.rs
