numdiff 0.7.3

/// Get a function that returns the Jacobian of the provided multivariate, vector-valued function.
///
/// The Jacobian is computed using forward-mode automatic differentiation.
///
/// # Arguments
///
/// * `f` - Multivariate, vector-valued function, $\mathbf{f}:\mathbb{R}^{n}\to\mathbb{R}^{m}$.
/// * `func_name` - Name of the function that will return the Jacobian of
///   $\mathbf{f}(\mathbf{x})$ at any point $\mathbf{x}\in\mathbb{R}^{n}$.
/// * `param_type` (optional) - Type of the extra runtime parameter `p` that is passed to `f`.
///   Defaults to `[f64]` (implying that `f` accepts `p: &[f64]`).
///
/// # Defining `f`
///
/// The multivariate, vector-valued function `f` must have the following function signature:
///
/// ```ignore
/// // Note that `T` is a placeholder for the type of the extra runtime parameter `p`, which can be
/// // any type.
/// fn f<S: Scalar, V: Vector<S>, T>(x: &V, p: &T) -> V::DVectorT<S> {
///     // place function contents here
/// }
/// ```
///
/// For the automatic differentiation to work, `f` must be fully generic over the types of scalars
/// and vectors used. Additionally, the function must return an instance of `V::DVector` (a
/// dynamically-sized vector type that is compatible with `V`, see the
/// [`linalg-traits` docs](https://docs.rs/linalg-traits/latest/linalg_traits/trait.Vector.html#associatedtype.DVectorT)
/// for more information) since we did not want to burden the user with having to specify the size
/// of the output vector (i.e. $m$, where $\mathbf{f}:\mathbb{R}^{n}\to\mathbb{R}^{m}$) at compile
/// time, especially since users may be using this crate exclusively with dynamically-sized types.
///
/// # Warning
///
/// `f` cannot be defined as closure. It must be defined as a function.
///
/// # Warning
///
/// This Jacobian function generated by this macro 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.
///
/// # Note
///
/// The function produced by this macro will perform $n$ evaluations of $\mathbf{f}(\mathbf{x})$ to
/// evaluate its Jacobian.
///
/// # Examples
///
/// ## Basic Example
///
/// Compute 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 linalg_traits::{Mat, Matrix, Scalar, Vector};
/// use numtest::*;
///
/// use numdiff::{get_jacobian, Dual, DualVector};
///
/// // Define the function, f(x).
/// fn f<S: Scalar, V: Vector<S>>(x: &V, _p: &[f64]) -> V::DVectorT<S> {
///     V::DVectorT::from_slice(&[
///         x.vget(0),
///         x.vget(2) * S::new(5.0),
///         x.vget(1).powi(2) * S::new(4.0) - x.vget(2) * S::new(2.0),
///         x.vget(2) * x.vget(0).sin(),
///     ])
/// }
///
/// // Define the evaluation point.
/// let x0 = vec![5.0, 6.0, 7.0];
///
/// // Parameter vector (empty for this example).
/// let p = [];
///
/// // Autogenerate the function "jac" that can be used to compute the Jacobian of f(x) at any point
/// // x.
/// get_jacobian!(f, jac);
///
/// // Evaluate the Jacobian using "jac".
/// let jac_eval: Mat<f64> = jac(&x0, &p);
///
/// // True Jacobian of f(x) at the evaluation point.
/// let jac_eval_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(),
///     ],
/// );
///
/// // Verify that the Jacobian function obtained using get_jacobian! computes the Jacobian
/// // correctly.
/// assert_eq!(jac_eval, jac_eval_true);
/// ```
///
/// #### Using other vector types
///
/// The function produced by `get_jacobian!` can accept _any_ type for `x0`, as long as it
/// implements the `linalg_traits::Vector` trait.
///
/// ```
/// use faer::Mat;
/// use linalg_traits::{Scalar, Vector};
/// use nalgebra::{dvector, DMatrix, DVector, SVector};
/// use ndarray::{array, Array1, Array2};
///
/// use numdiff::{get_jacobian, Dual, DualVector};
///
/// // Define the function, f(x).
/// fn f<S: Scalar, V: Vector<S>>(x: &V, _p: &[f64]) -> V::DVectorT<S> {
///     V::DVectorT::from_slice(&[
///         x.vget(0),
///         x.vget(2) * 5.0,
///         x.vget(1).powi(2) * 4.0 - x.vget(2) * 2.0,
///         x.vget(2) * x.vget(0).sin(),
///     ])
/// }
///
/// // Parameter vector (empty for this example).
/// let p = [];
///
/// // Autogenerate the function "jac" that can be used to compute the Jacobian of f(x) at any point
/// // x.
/// get_jacobian!(f, jac);
///
/// // nalgebra::DVector
/// let x0: DVector<f64> = dvector![5.0, 6.0, 7.0];
/// let jac_eval: DMatrix<f64> = jac(&x0, &p);
///
/// // nalgebra::SVector
/// let x0: SVector<f64, 3> = SVector::from_slice(&[5.0, 6.0, 7.0]);
/// let jac_eval: DMatrix<f64> = jac(&x0, &p);
///
/// // ndarray::Array1
/// let x0: Array1<f64> = array![5.0, 6.0, 7.0];
/// let jac_eval: Array2<f64> = jac(&x0, &p);
///
/// // faer::Mat
/// let x0: Mat<f64> = Mat::from_slice(&[5.0, 6.0, 7.0]);
/// let jac_eval: Mat<f64> = jac(&x0, &p);
/// ```
///
/// ## Example Passing Runtime Parameters
///
/// Compute 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, Scalar, Vector};
/// use numtest::*;
///
/// use numdiff::{get_jacobian, Dual, DualVector};
///
/// // Define the function, f(x).
/// fn f<S: Scalar, V: Vector<S>>(x: &V, p: &[f64]) -> V::DVectorT<S> {
///     let a = S::new(p[0]);
///     let b = S::new(p[1]);
///     let c = S::new(p[2]);
///     let d = S::new(p[3]);
///     V::DVectorT::from_slice(&[
///         a * x.vget(0).powi(2) + b * x.vget(1),
///         c * x.vget(0) + d * x.vget(1).powi(2)
///     ])
/// }
///
/// // Parameter vector.
/// let a = 1.5;
/// let b = 2.0;
/// let c = -0.8;
/// let d = 3.0;
/// let p = [a, b, c, d];
///
/// // Evaluation point.
/// let x0 = vec![1.0, -0.5];
///
/// // Autogenerate the Jacobian function.
/// get_jacobian!(f, jac);
///
/// // True Jacobian function.
/// let jac_true = |x: &Vec<f64>| Mat::from_row_slice(2, 2, &[
///     2.0 * a * x[0], b,
///     c, 2.0 * d * x[1]
/// ]);
///
/// // Compute the Jacobian using both the automatically generated Jacobian function and the true
/// // Jacobian function, and compare the results.
/// let jac_eval: Mat<f64> = jac(&x0, &p);
/// let jac_eval_true: Mat<f64> = jac_true(&x0);
/// assert_eq!(jac_eval, jac_eval_true);
/// ```
///
/// ## Example Passing Custom Parameter Types
///
/// Use a custom parameter struct instead of `f64` values.
///
/// ```
/// use linalg_traits::{Mat, Matrix, Scalar, Vector};
/// use numtest::*;
///
/// use numdiff::{get_jacobian, Dual, DualVector};
///
/// struct Data {
///     a: f64,
///     b: f64,
///     c: f64,
///     d: f64,
/// }
///
/// // Define the function, f(x).
/// fn f<S: Scalar, V: Vector<S>>(x: &V, p: &Data) -> V::DVectorT<S> {
///     let a = S::new(p.a);
///     let b = S::new(p.b);
///     let c = S::new(p.c);
///     let d = S::new(p.d);
///     V::DVectorT::from_slice(&[
///         a * x.vget(0).powi(2) + b * x.vget(1),
///         c * x.vget(0) + d * x.vget(1).powi(2)
///     ])
/// }
///
/// // Runtime parameter struct.
/// let p = Data {
///     a: 1.5,
///     b: 2.0,
///     c: -0.8,
///     d: 3.0,
/// };
///
/// // Evaluation point.
/// let x0 = vec![1.0, -0.5];
///
/// // Autogenerate the Jacobian function, telling the macro to expect a runtime parameter of type
/// // &Data.
/// get_jacobian!(f, jac, Data);
///
/// // True Jacobian function.
/// let jac_true = |x: &Vec<f64>| Mat::from_row_slice(2, 2, &[
///     2.0 * p.a * x[0], p.b,
///     p.c, 2.0 * p.d * x[1]
/// ]);
///
/// // Compute the Jacobian using both the automatically generated Jacobian function and the true
/// // Jacobian function, and compare the results.
/// let jac_eval: Mat<f64> = jac(&x0, &p);
/// let jac_eval_true: Mat<f64> = jac_true(&x0);
/// assert_eq!(jac_eval, jac_eval_true);
/// ```
#[macro_export]
macro_rules! get_jacobian {
    ($f:ident, $func_name:ident) => {
        get_jacobian!($f, $func_name, [f64]);
    };
    ($f:ident, $func_name:ident, $param_type:ty) => {
        /// Jacobian of a multivariate, vector-valued function `f: ℝⁿ → ℝᵐ`.
        ///
        /// This function is generated for a specific function `f` using the
        /// `numdiff::get_jacobian!` macro.
        ///
        /// # Arguments
        ///
        /// * `x0` - Evaluation point, `x₀ ∈ ℝⁿ`.
        /// * `p` - Extra runtime parameter. This is a parameter (can be of any arbitrary type)
        ///   defined at runtime that the function may depend on but is not differentiated with
        ///   respect to.
        ///
        /// # Returns
        ///
        /// Jacobian of `f` with respect to `x`, evaluated at `x = x₀`.
        ///
        /// `J(x₀) = (∂f/∂x)|ₓ₌ₓ₀ ∈ ℝᵐˣⁿ`
        fn $func_name<S, V>(x0: &V, p: &$param_type) -> V::DMatrixMxNf64
        where
            S: Scalar,
            V: Vector<S>,
        {
            // Promote the evaluation point to a vector of dual numbers.
            let mut x0_dual = x0.clone().to_dual_vector();

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

            // Variable to store the function evaluation at the evaluation point perturbed in the
            // kth dual direction.
            let mut f_x0k;

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

            // Take a unit step forward in the 0th dual direction.
            x0_dual.vset(0, Dual::new(x0k.get_real(), 1.0));

            // Evaluate the function at the evaluation point perturbed in the 0th dual direction.
            f_x0k = $f(&x0_dual, p);

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

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

            // Initialize a matrix of zeros to store the Jacobian.
            let mut jac = x0.new_dmatrix_m_by_n_f64(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)] = f_x0k.vget(i).get_dual();
            }

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

                // Take a unit step forward in the kth dual direction.
                x0_dual.vset(k, Dual::new(x0k.get_real(), 1.0));

                // Evaluate the function at the evaluation point perturbed in the kth dual
                // direction.
                f_x0k = $f(&x0_dual, p);

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

                // 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)] = f_x0k.vget(i).get_dual();
                }
            }

            jac
        }
    };
}

#[cfg(test)]
mod tests {
    use crate::{Dual, DualVector};
    use linalg_traits::{Mat, Matrix, Scalar, Vector};
    use nalgebra::{DMatrix, DVector, SVector, dvector};
    use ndarray::{Array1, Array2, array};

    #[test]
    fn test_jacobian_1() {
        // Function to take the Jacobian of.
        fn f<S: Scalar, V: Vector<S>>(x: &V, _p: &[f64]) -> V::DVectorT<S> {
            V::DVectorT::from_slice(&[x.vget(0).powi(2)])
        }

        // Evaluation point.
        let x0 = vec![2.0];

        // Parameter vector (empty for this test).
        let p = [];

        // True Jacobian function.
        let jac = |x: &Vec<f64>| Mat::from_row_slice(1, 1, &[2.0 * x[0]]);

        // Jacobian function obtained via forward-mode automatic differentation.
        get_jacobian!(f, jac_autodiff);

        // Evaluate the Jacobian using both functions.
        let jac_eval_autodiff: Mat<f64> = jac_autodiff(&x0, &p);
        let jac_eval: Mat<f64> = jac(&x0);

        // Test autodiff Jacobian against true Jacobian.
        assert_eq!(jac_eval_autodiff, jac_eval);
    }

    #[test]
    fn test_jacobian_2() {
        // Function to take the Jacobian of.
        fn f<S: Scalar, V: Vector<S>>(x: &V, _p: &[f64]) -> V::DVectorT<S> {
            V::DVectorT::from_slice(&[x.vget(0).powi(2), x.vget(0).powi(3)])
        }

        // Evaluation point.
        let x0 = array![2.0];

        // Parameter vector (empty for this test).
        let p = [];

        // True Jacobian function.
        let jac = |x: &Array1<f64>| {
            Array2::<f64>::from_row_slice(2, 1, &[2.0 * x[0], 3.0 * x[0].powi(2)])
        };

        // Jacobian function obtained via forward-mode automatic differentation.
        get_jacobian!(f, jac_autodiff);

        // Evaluate the Jacobian using both functions.
        let jac_eval_autodiff: Array2<f64> = jac_autodiff(&x0, &p);
        let jac_eval: Array2<f64> = jac(&x0);

        // Test autodiff Jacobian against true Jacobian.
        assert_eq!(jac_eval_autodiff, jac_eval);
    }

    #[test]
    fn test_jacobian_3() {
        // Function to take the Jacobian of.
        fn f<S: Scalar, V: Vector<S>>(x: &V, _p: &[f64]) -> V::DVectorT<S> {
            V::DVectorT::from_slice(&[x.vget(0).powi(2) + x.vget(1).powi(3)])
        }

        // Evaluation point.
        let x0 = dvector![1.0, 2.0];

        // Parameter vector (empty for this test).
        let p = [];

        // True Jacobian function.
        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)])
        };

        // Jacobian function obtained via forward-mode automatic differentation.
        get_jacobian!(f, jac_autodiff);

        // Evaluate the Jacobian using both functions.
        let jac_eval_autodiff: DMatrix<f64> = jac_autodiff(&x0, &p);
        let jac_eval: DMatrix<f64> = jac(&x0);

        // Test autodiff Jacobian against true Jacobian.
        assert_eq!(jac_eval_autodiff, jac_eval);
    }

    #[test]
    fn test_jacobian_4() {
        // Function to take the Jacobian of.
        fn f<S: Scalar, V: Vector<S>>(x: &V, _p: &[f64]) -> V::DVectorT<S> {
            V::DVectorT::from_slice(&[x.vget(0).powi(2), x.vget(1).powi(3)])
        }

        // Evaluation point.
        let x0 = SVector::<f64, 2>::from_row_slice(&[1.0, 2.0]);

        // Parameter vector (empty for this test).
        let p = [];

        // True Jacobian function.
        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)])
        };

        // Jacobian function obtained via forward-mode automatic differentation.
        get_jacobian!(f, jac_autodiff);

        // Evaluate the Jacobian using both functions.
        let jac_eval_autodiff: DMatrix<f64> = jac_autodiff(&x0, &p);
        let jac_eval: DMatrix<f64> = jac(&x0);

        // Test autodiff Jacobian against true Jacobian.
        assert_eq!(jac_eval_autodiff, jac_eval);
    }

    #[test]
    fn test_jacobian_5() {
        // Function to take the Jacobian of.
        fn f<S: Scalar, V: Vector<S>>(x: &V, _p: &[f64]) -> V::DVectorT<S> {
            V::DVectorT::from_slice(&[
                x.vget(0),
                x.vget(2) * 5.0,
                x.vget(1).powi(2) * 4.0 - x.vget(2) * 2.0,
                x.vget(2) * x.vget(0).sin(),
            ])
        }

        // Evaluation point.
        let x0: SVector<f64, 3> = SVector::from_row_slice(&[5.0, 6.0, 7.0]);

        // Parameter vector (empty for this test).
        let p = [];

        // True Jacobian function.
        let jac = |x: &SVector<f64, 3>| {
            DMatrix::<f64>::from_row_slice(
                4,
                3,
                &[
                    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(),
                ],
            )
        };

        // Jacobian function obtained via forward-mode automatic differentation.
        get_jacobian!(f, jac_autodiff);

        // Evaluate the Jacobian using both functions.
        let jac_eval_autodiff: DMatrix<f64> = jac_autodiff(&x0, &p);
        let jac_eval: DMatrix<f64> = jac(&x0);

        // Test autodiff Jacobian against true Jacobian.
        assert_eq!(jac_eval_autodiff, jac_eval);
    }

    #[test]
    fn test_jacobian_6() {
        // Function to take the Jacobian of.
        #[allow(clippy::many_single_char_names)]
        fn f<S: Scalar, V: Vector<S>>(x: &V, p: &[f64]) -> V::DVectorT<S> {
            let a = S::new(p[0]);
            let b = S::new(p[1]);
            let c = S::new(p[2]);
            let d = S::new(p[3]);
            V::DVectorT::from_slice(&[
                a * (b * x.vget(0)).sin(),
                c * x.vget(0) * x.vget(1) + d * x.vget(1).cos(),
            ])
        }

        // Parameter vector.
        let p = [2.0, 0.5, 1.2, -1.5];

        // Evaluation point.
        let x0 = dvector![1.0, 0.8];

        // True Jacobian function.
        let jac = |x: &DVector<f64>, p: &[f64]| {
            DMatrix::<f64>::from_row_slice(
                2,
                2,
                &[
                    p[0] * p[1] * (p[1] * x[0]).cos(),
                    0.0,
                    p[2] * x[1],
                    p[2] * x[0] - p[3] * x[1].sin(),
                ],
            )
        };

        // Jacobian function obtained via forward-mode automatic differentiation.
        get_jacobian!(f, jac_autodiff);

        // Evaluate the Jacobian using both methods.
        let jac_eval_autodiff: DMatrix<f64> = jac_autodiff(&x0, &p);
        let jac_eval: DMatrix<f64> = jac(&x0, &p);

        // Test autodiff Jacobian against true Jacobian.
        assert_eq!(jac_eval_autodiff, jac_eval);
    }

    #[test]
    fn test_jacobian_custom_params() {
        struct Data {
            a: f64,
            b: f64,
            c: f64,
            d: f64,
        }

        // Function to take the Jacobian of.
        #[allow(clippy::many_single_char_names)]
        fn f<S: Scalar, V: Vector<S>>(x: &V, p: &Data) -> V::DVectorT<S> {
            let a = S::new(p.a);
            let b = S::new(p.b);
            let c = S::new(p.c);
            let d = S::new(p.d);
            V::DVectorT::from_slice(&[
                a * x.vget(0).powi(2) + b * x.vget(1),
                c * x.vget(0) + d * x.vget(1).powi(2),
            ])
        }

        // Runtime parameter struct.
        let p = Data {
            a: 1.5,
            b: 2.0,
            c: -0.8,
            d: 3.0,
        };

        // Evaluation point.
        let x0 = vec![1.0, -0.5];

        // Jacobian function obtained via forward-mode automatic differentiation.
        get_jacobian!(f, jac, Data);

        // True Jacobian function.
        let jac_true = |x: &Vec<f64>| {
            Mat::from_row_slice(2, 2, &[2.0 * p.a * x[0], p.b, p.c, 2.0 * p.d * x[1]])
        };

        // Evaluate the Jacobian using both functions.
        let jac_eval: Mat<f64> = jac(&x0, &p);
        let jac_eval_true: Mat<f64> = jac_true(&x0);

        // Test autodiff Jacobian against true Jacobian.
        assert_eq!(jac_eval, jac_eval_true);
    }
}