rexafs 0.1.0

Rust-powered X-ray absorption spectroscopy analysis and EXAFS fitting
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

use levenberg_marquardt::{LeastSquaresProblem, LevenbergMarquardt};
use nalgebra::{DMatrix, DVector, Dyn, Owned};

const EPS_F64: f64 = f64::EPSILON;

/// Update the function value at `x[idx]` and return the value.
pub fn mod_and_calc_nalgebra_f64<T>(
    x: &mut DVector<f64>,
    f: &dyn Fn(&DVector<f64>) -> T,
    idx: usize,
    y: f64,
) -> T {
    let xtmp = x[idx];
    x[idx] = xtmp + y;
    let fx1 = (f)(x);
    x[idx] = xtmp;
    fx1
}

/// Calculation of the Jacobian matrix by forward difference.
pub fn forward_jacobian_nalgebra_f64(
    x: &DVector<f64>,
    fs: &dyn Fn(&DVector<f64>) -> DVector<f64>,
) -> DMatrix<f64> {
    let fx = (fs)(x);
    let mut xt = x.clone();
    let mut jac = DMatrix::zeros(fx.len(), x.len());
    for i in 0..x.len() {
        let fx1 = mod_and_calc_nalgebra_f64(&mut xt, fs, i, EPS_F64.sqrt());
        jac.set_column(i, &((fx1 - &fx) / EPS_F64.sqrt()));
    }
    jac
}

/// Calculation of the Jacobian matrix by central difference.
pub fn center_jacobian_nalgebra_f64(
    x: &DVector<f64>,
    fs: &dyn Fn(&DVector<f64>) -> DVector<f64>,
) -> DMatrix<f64> {
    let fx = (fs)(x);
    let mut xt = x.clone();
    let mut jac = DMatrix::zeros(fx.len(), x.len());
    for i in 0..x.len() {
        let fx1 = mod_and_calc_nalgebra_f64(&mut xt, fs, i, EPS_F64.sqrt());
        let fx2 = mod_and_calc_nalgebra_f64(&mut xt, fs, i, -EPS_F64.sqrt());
        jac.set_column(i, &((fx1 - fx2) / (2.0 * EPS_F64.sqrt())));
    }
    jac
}

/// Calculation of the approximate Hessian matrix.
/// The Hessian matrix is calculated as the product of the Jacobian matrix and its transpose.
/// This approximation is valid only when the residual function near the minimum.
pub fn approx_hessian_nalgebra_f64(
    x: &DVector<f64>,
    fs: &dyn Fn(&DVector<f64>) -> DVector<f64>,
) -> DMatrix<f64> {
    let jac = forward_jacobian_nalgebra_f64(x, fs);
    jac.transpose() * &jac
}

/// Calculation of the approximate covariance matrix.
/// The covariance matrix is calculated as the inverse of the approximate Hessian matrix.
/// This approximation is valid only when the residual function near the minimum.
pub fn approx_covariance_matrix_nalgebra_f64(
    x: &DVector<f64>,
    fs: &dyn Fn(&DVector<f64>) -> DVector<f64>,
) -> Option<DMatrix<f64>> {
    let hess = approx_hessian_nalgebra_f64(x, fs);
    hess.try_inverse()
}

/// Trait for Levenberg-Marquardt parameters.
/// It implements the Jacobian matrix, the Hessian matrix, and the covariance matrix.
pub trait LMParameters<T> {
    fn jacobian(&self, f: T) -> DMatrix<f64>;
    fn hessian(&self, f: T) -> DMatrix<f64>;
    fn covariance(&self, f: T) -> Option<DMatrix<f64>>;
}

impl LMParameters<&dyn Fn(&DVector<f64>) -> DVector<f64>> for DVector<f64> {
    fn jacobian(&self, f: &dyn Fn(&DVector<f64>) -> DVector<f64>) -> DMatrix<f64> {
        forward_jacobian_nalgebra_f64(self, f)
    }

    fn hessian(&self, f: &dyn Fn(&DVector<f64>) -> DVector<f64>) -> DMatrix<f64> {
        approx_hessian_nalgebra_f64(self, f)
    }

    fn covariance(&self, f: &dyn Fn(&DVector<f64>) -> DVector<f64>) -> Option<DMatrix<f64>> {
        approx_covariance_matrix_nalgebra_f64(self, f)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::xafs::tests::PARAM_LOADTXT;
    use crate::xafs::tests::TEST_TOL;
    use crate::xafs::tests::TOP_DIR;
    use approx::assert_abs_diff_eq;

    const NUM_DIFF_TOL: f64 = 1e-6;

    fn residuals(p: &DVector<f64>) -> DVector<f64> {
        let mut res = DVector::zeros(p.len());

        res.iter_mut()
            .enumerate()
            .for_each(|(i, x)| *x = p[i] - (i as f64 + 1.0));

        res
    }

    #[test]
    fn test_forward_jacobian_nalgebra_f64() {
        let x = DVector::from_vec(vec![1.0, 2.0, 3.0]);
        let fs = |x: &DVector<f64>| {
            let mut y = DVector::zeros(x.len());
            y[0] = 2.0 * x[0] + 3.0 * x[1] + 4.0 * x[2];
            y[1] = 3.0 * x[1] + 4.0 * x[2];
            y[2] = 4.0 * x[2] + 5.0 + x[1].exp();
            y
        };
        let jac = forward_jacobian_nalgebra_f64(&x, &fs);
        let jac_ref = DMatrix::from_vec(
            3,
            3,
            vec![2.0, 0.0, 0.0, 3.0, 3.0, 2.0_f64.exp(), 4.0, 4.0, 4.0],
        );

        assert_abs_diff_eq!(jac, jac_ref, epsilon = NUM_DIFF_TOL);
    }

    #[test]
    fn test_center_jacobian_nalgebra_f64() {
        let x = DVector::from_vec(vec![1.0, 2.0, 3.0]);
        let fs = |x: &DVector<f64>| {
            let mut y = DVector::zeros(x.len());
            y[0] = 2.0 * x[0] + 3.0 * x[1] + 4.0 * x[2];
            y[1] = 3.0 * x[1] + 4.0 * x[2];
            y[2] = 4.0 * x[2] + 5.0 + x[1].exp();
            y
        };
        let jac = center_jacobian_nalgebra_f64(&x, &fs);
        let jac_ref = DMatrix::from_vec(
            3,
            3,
            vec![2.0, 0.0, 0.0, 3.0, 3.0, 2.0_f64.exp(), 4.0, 4.0, 4.0],
        );

        assert_abs_diff_eq!(jac, jac_ref, epsilon = NUM_DIFF_TOL);
    }
}