numdiff 0.7.3

use crate::constants::CBRT_EPS;
use linalg_traits::Vector;

/// Jacobian of a multivariate, vector-valued function using the central difference approximation.
///
/// # Arguments
///
/// * `f` - Multivariate, vector-valued function, $\mathbf{f}:\mathbb{R}^{n}\to\mathbb{R}^{m}$.
/// * `x0` - Evaluation point, $\mathbf{x}_{0}\in\mathbb{R}^{n}$.
/// * `h` - Relative step size, $h\in\mathbb{R}$. Defaults to [`CBRT_EPS`].
///
/// # Returns
///
/// Jacobian of $\mathbf{f}$ with respect to $\mathbf{x}$, evaluated at $\mathbf{x}=\mathbf{x}_{0}$.
///
/// $$\mathbf{J}(\mathbf{x}\_{0})={\frac{\partial\mathbf{f}}{\partial\mathbf{x}}\bigg\rvert_{\mathbf{x}=\mathbf{x}_{0}}}\in\mathbb{R}^{m\times n}$$
///
/// # Note
///
/// This function performs $2n$ evaluations of $f(x)$.
///
/// # Warning
///
/// This function will always return a dynamically-sized matrix, even if the function `f` uses
/// statically-sized vectors. This is to avoid needing to pass a const generic to this function to
/// define the number of rows ($m$) of the Jacobian. Instead, the number of rows is determined at
/// runtime.
///
/// # Examples
///
/// ## Basic Example
///
/// Approximate the Jacobian of
///
/// $$
/// \mathbf{f}(\mathbf{x})=
/// \begin{bmatrix}
///     x_{0} \\\\
///     5x_{2} \\\\
///     4x_{1}^{2}-2x_{2} \\\\
///     x_{2}\sin{x_{0}}
/// \end{bmatrix}
/// $$
///
/// at $\mathbf{x}=(5,6,7)^{T}$, and compare the result to the true result of
///
/// $$
/// \mathbf{J}\left((5,6,7)^{T}\right)=
/// \begin{bmatrix}
///     1 & 0 & 0 \\\\
///     0 & 0 & 5 \\\\
///     0 & 48 & -2 \\\\
///     7\cos{(5)} & 0 & \sin{(5)}
/// \end{bmatrix}
/// $$
///
/// #### Using standard vectors
///
/// ```
/// use numtest::*;
///
/// use linalg_traits::{Mat, Matrix};
/// use numdiff::central_difference::jacobian;
///
/// // Define the function, f(x).
/// let f = |x: &Vec<f64>| {
///     vec![
///         x[0],
///         5.0 * x[2],
///         4.0 * x[1].powi(2) - 2.0 * x[2],
///         x[2] * x[0].sin(),
///     ]
/// };
///
/// // Define the evaluation point.
/// let x0 = vec![5.0, 6.0, 7.0];
///
/// // Approximate the Jacobian of f(x) at the evaluation point.
/// let jac: Mat<f64> = jacobian(&f, &x0, None);
///
/// // True Jacobian of f(x) at the evaluation point.
/// let jac_true: Mat<f64> = Mat::from_row_slice(
///     4,
///     3,
///     &[
///         1.0,
///         0.0,
///         0.0,
///         0.0,
///         0.0,
///         5.0,
///         0.0,
///         48.0,
///         -2.0,
///         7.0 * 5.0_f64.cos(),
///         0.0,
///         5.0_f64.sin(),
///     ],
/// );
///
/// // Check the accuracy of the Jacobian approximation.
/// assert_arrays_equal_to_decimal!(jac, jac_true, 9);
/// ```
///
/// #### Using other vector types
///
/// We can also use other types of vectors, such as `nalgebra::SVector`, `nalgebra::DVector`,
/// `ndarray::Array1`, `faer::Mat`, or any other type of vector that implements the
/// `linalg_traits::Vector` trait.
///
/// ```
/// use faer::Mat as FMat;
/// use linalg_traits::{Mat, Matrix, Vector};
/// use nalgebra::{dvector, DMatrix, DVector, SVector};
/// use ndarray::{array, Array1, Array2};
/// use numtest::*;
///
/// use numdiff::central_difference::jacobian;
///
/// let jac_true_row_major: Mat<f64> = Mat::from_row_slice(
///     4,
///     3,
///     &[
///         1.0,
///         0.0,
///         0.0,
///         0.0,
///         0.0,
///         5.0,
///         0.0,
///         48.0,
///         -2.0,
///         7.0 * 5.0_f64.cos(),
///         0.0,
///         5.0_f64.sin(),
///     ],
/// );
/// let jac_true_col_major: DMatrix<f64> = DMatrix::from_row_slice(
///     4,
///     3,
///     &[
///         1.0,
///         0.0,
///         0.0,
///         0.0,
///         0.0,
///         5.0,
///         0.0,
///         48.0,
///         -2.0,
///         7.0 * 5.0_f64.cos(),
///         0.0,
///         5.0_f64.sin(),
///     ],
/// );
///
/// // nalgebra::DVector
/// let f_dvector = |x: &DVector<f64>| {
///     dvector![
///         x[0],
///         5.0 * x[2],
///         4.0 * x[1].powi(2) - 2.0 * x[2],
///         x[2] * x[0].sin(),
///     ]
/// };
/// let x0_dvector: DVector<f64> = dvector![5.0, 6.0, 7.0];
/// let jac_dvector: DMatrix<f64> = jacobian(&f_dvector, &x0_dvector, None);
/// assert_arrays_equal_to_decimal!(jac_dvector, jac_true_col_major, 9);
///
/// // nalgebra::SVector
/// let f_svector = |x: &SVector<f64, 3>| {
///     SVector::<f64, 4>::from_row_slice(&[
///         x[0],
///         5.0 * x[2],
///         4.0 * x[1].powi(2) - 2.0 * x[2],
///         x[2] * x[0].sin(),
///     ])
/// };
/// let x0_svector: SVector<f64, 3> = SVector::from_row_slice(&[5.0, 6.0, 7.0]);
/// let jac_svector: DMatrix<f64> = jacobian(&f_svector, &x0_svector, None);
/// assert_arrays_equal_to_decimal!(jac_svector, jac_true_col_major, 9);
///
/// // ndarray::Array1
/// let f_array1 = |x: &Array1<f64>| {
///     array![
///         x[0],
///         5.0 * x[2],
///         4.0 * x[1].powi(2) - 2.0 * x[2],
///         x[2] * x[0].sin(),
///     ]
/// };
/// let x0_array1: Array1<f64> = array![5.0, 6.0, 7.0];
/// let jac_array1: Array2<f64> = jacobian(&f_array1, &x0_array1, None);
/// assert_arrays_equal_to_decimal!(jac_array1, jac_true_row_major, 9);
///
/// // faer::Mat
/// let f_mat = |x: &FMat<f64>| {
///     FMat::from_slice(
///         &[
///             x[(0, 0)],
///             5.0 * x[(2, 0)],
///             4.0 * x[(1, 0)].powi(2) - 2.0 * x[(2, 0)],
///             x[(2, 0)] * x[(0, 0)].sin(),
///         ]
///     )
/// };
/// let x0_mat: FMat<f64> = FMat::from_slice(&[5.0, 6.0, 7.0]);
/// let jac_mat: FMat<f64> = jacobian(&f_mat, &x0_mat, None);
/// assert_arrays_equal_to_decimal!(jac_mat.as_row_slice(), jac_true_row_major, 9);
/// ```
///
/// #### Modifying the relative step size
///
/// We can also modify the relative step size. Choosing a coarser relative step size, we get a worse
/// approximation.
///
/// ```
/// use linalg_traits::{Mat, Matrix};
/// use numtest::*;
///
/// use numdiff::central_difference::jacobian;
///
/// let f = |x: &Vec<f64>| {
///     vec![
///         x[0],
///         5.0 * x[2],
///         4.0 * x[1].powi(2) - 2.0 * x[2],
///         x[2] * x[0].sin(),
///     ]
/// };
/// let x0 = vec![5.0, 6.0, 7.0];
///
/// let jac: Mat<f64> = jacobian(&f, &x0, Some(0.001));
/// let jac_true: Mat<f64> = Mat::from_row_slice(
///     4,
///     3,
///     &[
///         1.0,
///         0.0,
///         0.0,
///         0.0,
///         0.0,
///         5.0,
///         0.0,
///         48.0,
///         -2.0,
///         7.0 * 5.0_f64.cos(),
///         0.0,
///         5.0_f64.sin(),
///     ],
/// );
///
/// assert_arrays_equal_to_decimal!(jac, jac_true, 5);
/// ```
///
/// ## Example Passing Runtime Parameters
///
/// Approximate the Jacobian of a parameterized system
///
/// $$
/// \mathbf{f}(\mathbf{x})=
/// \begin{bmatrix}
/// ax_{0}^{2}+bx_{1} \\\\
/// cx_{0}+dx_{1}^{2}
/// \end{bmatrix}
/// $$
///
/// where $a$, $b$, $c$, and $d$ are runtime parameters. Compare the result against the true
/// Jacobian of
///
/// $$
/// \mathbf{J}=
/// \begin{bmatrix}
/// 2ax_{0} & b \\\\
/// c & 2dx_{1}
/// \end{bmatrix}
/// $$
///
/// ```
/// use linalg_traits::{Mat, Matrix};
/// use numtest::*;
///
/// use numdiff::central_difference::jacobian;
///
/// // Runtime parameters.
/// let a = 1.5;
/// let b = 2.0;
/// let c = -0.8;
/// let d = 3.0;
///
/// // Define the parameterized function.
/// fn f_param(x: &Vec<f64>, a: f64, b: f64, c: f64, d: f64) -> Vec<f64> {
///     vec![
///         a * x[0].powi(2) + b * x[1],
///         c * x[0] + d * x[1].powi(2)
///     ]
/// }
///
/// // Wrap the parameterized function with a closure that captures the parameters.
/// let f = |x: &Vec<f64>| f_param(x, a, b, c, d);
///
/// // Evaluation point.
/// let x0 = vec![1.0, -0.5];
///
/// // True Jacobian function.
/// let jac_true = Mat::from_row_slice(2, 2, &[
///     2.0 * a * x0[0], b,
///     c, 2.0 * d * x0[1]
/// ]);
///
/// // Approximate the Jacobian and compare with true Jacobian.
/// let jac_eval: Mat<f64> = jacobian(&f, &x0, None);
/// assert_arrays_equal_to_decimal!(jac_eval, jac_true, 9);
/// ```
pub fn jacobian<V, U>(f: &impl Fn(&V) -> U, x0: &V, h: Option<f64>) -> V::DMatrixMxN
where
    V: Vector<f64>,
    U: Vector<f64>,
{
    // Copy the evaluation point so that we may modify it.
    let mut x0 = x0.clone();

    // Default the relative step size to h = ε¹ᐟ³ if not specified.
    let h = h.unwrap_or(*CBRT_EPS);

    // Variable to store the absolute step size in the kth direction.
    let mut dxk: f64;

    // Variable to store the original value of the evaluation point in the kth direction.
    let mut x0k: f64;

    // Variable to store the partial derivative of f with respect to xₖ.
    let mut dfk;

    // Variable to store the function evaluation from a forward step in the kth direction.
    let mut f1;

    // Variable to store the function evaluation from a backward step in the kth direction.
    let mut f2;

    // Original value of the evaluation point in the 0th direction.
    x0k = x0.vget(0);

    // Absolute step size in the 0th direction.
    dxk = h * (1.0 + x0.vget(0).abs());

    // Step forward in the 0th direction.
    x0.vset(0, x0.vget(0) + dxk);
    f1 = f(&x0);

    // Step backward in the 0th direction.
    x0.vset(0, x0k - dxk);
    f2 = f(&x0);

    // Reset the evaluation point.
    x0.vset(0, x0k);

    // Partial derivative of f with respect to x₀ (0th column of the Jacobian).
    dfk = f1.sub(&f2).div(2.0 * dxk);

    // Determine the size of the Jacobian.
    let m = dfk.len();
    let n = x0.len();

    // Initialize a matrix of zeros to store the Jacobian.
    let mut jac = x0.new_dmatrix_m_by_n(m);

    // Store the partial derivative of f with respect to x₀ in the 0th column of the Jacobian.
    for i in 0..m {
        jac[(i, 0)] = dfk.vget(i);
    }

    // Evaluate the remaining columns of the Jacobian.
    for k in 1..n {
        // Original value of the evaluation point in the kth direction.
        x0k = x0.vget(k);

        // Absolute step size in the kth direction.
        dxk = h * (1.0 + x0k.abs());

        // Step forward in the kth direction.
        x0.vset(k, x0k + dxk);
        f1 = f(&x0);

        // Step backward in the kth direction.
        x0.vset(k, x0k - dxk);
        f2 = f(&x0);

        // Reset the evaluation point.
        x0.vset(k, x0k);

        // Partial derivative of f with respect to xₖ.
        dfk = f1.sub(&f2).div(2.0 * dxk);

        // Store the partial derivative of f with respect to xₖ in the kth column of the Jacobian.
        for i in 0..m {
            jac[(i, k)] = dfk.vget(i);
        }
    }

    jac
}

#[cfg(test)]
mod tests {
    use super::*;
    use linalg_traits::{Mat, Matrix};
    use nalgebra::{DMatrix, DVector, SMatrix, SVector, dvector};
    use ndarray::{Array1, Array2, array};
    use numtest::*;

    #[test]
    fn test_jacobian_1() {
        let f = |x: &Vec<f64>| vec![x[0].powi(2)];
        let x0 = vec![2.0];
        let jac = |x: &Vec<f64>| Mat::from_row_slice(1, 1, &[2.0 * x[0]]);
        assert_arrays_equal_to_decimal!(jacobian(&f, &x0, None), jac(&x0), 11);
    }

    #[test]
    fn test_jacobian_2() {
        let f = |x: &Array1<f64>| array![x[0].powi(2), x[0].powi(3)];
        let x0 = array![2.0];
        let jac = |x: &Array1<f64>| {
            Array2::<f64>::from_row_slice(2, 1, &[2.0 * x[0], 3.0 * x[0].powi(2)])
        };
        assert_arrays_equal_to_decimal!(jacobian(&f, &x0, None), jac(&x0), 9);
    }

    #[test]
    fn test_jacobian_3() {
        let f = |x: &DVector<f64>| dvector![x[0].powi(2) + x[1].powi(3)];
        let x0 = dvector![1.0, 2.0];
        let jac = |x: &DVector<f64>| {
            <DMatrix<f64> as Matrix<f64>>::from_row_slice(1, 2, &[2.0 * x[0], 3.0 * x[1].powi(2)])
        };
        assert_arrays_equal_to_decimal!(jacobian(&f, &x0, None), jac(&x0), 9);
    }

    #[test]
    fn test_jacobian_4() {
        let f =
            |x: &SVector<f64, 2>| SVector::<f64, 2>::from_row_slice(&[x[0].powi(2), x[1].powi(3)]);
        let x0 = SVector::<f64, 2>::from_row_slice(&[1.0, 2.0]);
        let jac = |x: &SVector<f64, 2>| {
            DMatrix::<f64>::from_row_slice(2, 2, &[2.0 * x[0], 0.0, 0.0, 3.0 * x[1].powi(2)])
        };
        assert_arrays_equal_to_decimal!(jacobian(&f, &x0, None), jac(&x0), 9);
    }

    #[test]
    fn test_jacobian_5() {
        let f = |x: &SVector<f64, 3>| {
            SVector::<f64, 4>::from_row_slice(&[
                x[0],
                5.0 * x[2],
                4.0 * x[1].powi(2) - 2.0 * x[2],
                x[2] * x[0].sin(),
            ])
        };
        let x0: SVector<f64, 3> = SVector::from_row_slice(&[5.0, 6.0, 7.0]);
        let jac = |x: &SVector<f64, 3>| {
            SMatrix::<f64, 4, 3>::from_row_slice(&[
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                5.0,
                0.0,
                8.0 * x[1],
                -2.0,
                x[2] * x[0].cos(),
                0.0,
                x[0].sin(),
            ])
        };
        assert_arrays_equal_to_decimal!(jacobian(&f, &x0, None), jac(&x0), 9);
    }
}