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


use ndarray::*;
use ndarray_linalg::*;
use num_traits::Float;

use super::traits::*;

/// Jacobian operator using numerical-differentiation
pub struct Jacobian<'a, A, S, D, TEO>
    where A: Scalar,
          S: Data<Elem = A>,
          D: Dimension,
          TEO: 'a
{
    f: &'a TEO,
    x: ArrayBase<S, D>,
    fx: ArrayBase<S, D>,
    alpha: A::Real,
}

pub fn jacobian<'a, A, S, D, TEO>(f: &'a TEO,
                                  x: ArrayBase<S, D>,
                                  alpha: A::Real)
                                  -> Jacobian<'a, A, S, D, TEO>
    where A: Scalar,
          S: DataClone<Elem = A> + DataMut,
          D: Dimension,
          TEO: TimeEvolutionBase<S, D>
{
    let mut fx = x.clone();
    f.iterate(&mut fx);
    Jacobian {
        f: f,
        x: x,
        fx: fx,
        alpha: alpha,
    }
}

impl<'j, A, S, Sr, D, TEO> OperatorMut<Sr, D> for Jacobian<'j, A, S, D, TEO>
    where A: Scalar,
          S: Data<Elem = A>,
          Sr: DataMut<Elem = A>,
          D: Dimension,
          TEO: TimeEvolutionBase<Sr, D>
{
    fn op_mut<'a>(&self, dx: &'a mut ArrayBase<Sr, D>) -> &'a mut ArrayBase<Sr, D> {
        let dx_nrm = dx.norm_l2().max(self.alpha);
        let n = self.alpha / dx_nrm;
        Zip::from(&mut *dx)
            .and(&self.x)
            .apply(|dx, &x| { *dx = x + dx.mul_real(n); });
        let x_dx = self.f.iterate(dx);
        Zip::from(&mut *x_dx)
            .and(&self.fx)
            .apply(|x_dx, &fx| { *x_dx = (*x_dx - fx).div_real(n); });
        x_dx
    }
}

impl<'j, A, S, Sr, D, TEO> OperatorInto<Sr, D> for Jacobian<'j, A, S, D, TEO>
    where A: Scalar,
          S: Data<Elem = A>,
          Sr: DataMut<Elem = A>,
          D: Dimension,
          TEO: TimeEvolutionBase<Sr, D>
{
    fn op_into(&self, mut dx: ArrayBase<Sr, D>) -> ArrayBase<Sr, D> {
        self.op_mut(&mut dx);
        dx
    }
}

impl<'j, A, Si, S, D, TEO> Operator<A, Si, D> for Jacobian<'j, A, S, D, TEO>
    where A: Scalar,
          S: Data<Elem = A>,
          Si: Data<Elem = A>,
          D: Dimension,
          TEO: TimeEvolutionBase<OwnedRepr<A>, D>
{
    fn op(&self, dx: &ArrayBase<Si, D>) -> Array<A, D> {
        let dx = replicate(dx);
        self.op_into(dx)
    }
}