autodiff.h

Go to the documentation of this file.

 
 
#ifndef THEORETICA_AUTODIFF_H
#define THEORETICA_AUTODIFF_H
 
#include "dual.h"
#include "dual2.h"
#include "multidual.h"
#include "../algebra/vec.h"
#include "../algebra/mat.h"
#include "../core/error.h"
#include "../core/core_traits.h"
#include "./autodiff_types.h"
 
#include <functional>
 
 
namespace theoretica {
 
    namespace autodiff {
 
 
        // Univariate automatic differentiation
 
 
        template <
            typename DualFunction = std::function<dual(dual)>,
            enable_dual_func<DualFunction> = true
        >
        inline real deriv(DualFunction f, real x) {
            return f(dual(x, 1.0)).Dual();
        }
 
 
        template <
            typename DualFunction = std::function<dual(dual)>,
            enable_dual_func<DualFunction> = true
        >
        inline auto deriv(DualFunction f) {
 
            return [f](real x) -> real {
                return autodiff::deriv(f, x);
            };
        }
 
 
        template <
            typename Dual2Function = std::function<dual2(dual2)>,
            enable_dual2_func<Dual2Function> = true
        >
        inline real deriv2(Dual2Function f, real x) {
            return f(dual2(x, 1.0, 0.0)).Dual2();
        }
 
 
        template <
            typename Dual2Function = std::function<dual2(dual2)>,
            enable_dual2_func<Dual2Function> = true
        >
        inline auto deriv2(Dual2Function f) {
 
            return [f](real x) -> real {
                return autodiff::deriv2(f, x);
            };
        }
 
 
        // Differential operators
 
 
        template <
            typename MultidualType,
            typename Vector = vec<real>
        >
        inline auto make_autodiff_arg(const Vector& x) {
 
            constexpr size_t N = MultidualType::vector_argument;
            vec<MultidualType, N> arg;
            arg.resize(x.size());
 
            // Iterate over each element, setting the dual part to
            // the i-th element of the canonical base
            for (unsigned int i = 0; i < x.size(); ++i) {
                arg[i] = MultidualType(
                    x[i], vec<real, N>::euclidean_base(i, x.size())
                );
            }
 
            return arg;
        }
 
 
        template <
            typename Function, typename Vector = vec<real>,
            enable_scalar_field<Function> = true,
            enable_vector<Vector> = true
        >
        inline auto gradient(Function f, const Vector& x) {
 
            // Extract the return type of the function
            using R = return_type_t<Function>;
 
            return f(make_autodiff_arg<R>(x)).Dual();
        }
 
 
        template <
            typename Function,
            enable_scalar_field<Function> = true
        >
        inline auto gradient(Function f) {
 
            constexpr size_t N = return_type_t<Function>::vector_argument;
 
            return [f](vec<real, N> x) {
                return gradient(f, x);
            };
        }
 
 
        template <
            typename Function, typename Vector = vec<real>,
            enable_scalar_field<Function> = true,
            enable_vector<Vector> = true
        >
        inline real divergence(Function f, const Vector& x) {
 
            using MultidualT = return_type_t<Function>;
            MultidualT d = f(MultidualT::make_argument(x));
 
            real div = 0;
            for (unsigned int i = 0; i < d.v.size(); ++i)
                div += d.v[i];
 
            return div;
        }
 
 
        template <
            typename Function,
            enable_scalar_field<Function> = true
        >
        inline auto divergence(Function f) {
 
            constexpr size_t N = return_type_t<Function>::vector_argument;
 
            return [f](vec<real, N> x) {
                return divergence(f, x);
            };
        }
 
 
        template<unsigned int N = 0, unsigned int M = 0>
        inline mat<real, M, N> jacobian(
            vec<multidual<N>, M>(*f)(vec<multidual<N>, N>), const vec<real, N>& x) {
 
            vec<multidual<N>, M> res = f(multidual<N>::make_argument(x));
            
            // Construct the jacobian matrix
            mat<real, M, N> J;
            J.resize(res.size(), x.size());
 
            for (unsigned int j = 0; j < J.rows(); ++j)
                for (unsigned int i = 0; i < res[j].v.size(); ++i)
                    J(j, i) = res[j].v[i];
 
            return J;
        }
 
 
        template<unsigned int N = 0, unsigned int M = 0>
        inline auto jacobian(
            vec<multidual<N>, M>(*f)(vec<multidual<N>, N>)) {
 
            return [f](vec<real, N> x) {
                return jacobian(f, x);
            };
        }
 
 
        template<unsigned int N = 0>
        inline vec<real, N> curl(
            vec<multidual<N>, N>(*f)(vec<multidual<N>, N>), const vec<real, N>& x) {
 
            if(x.size() != 3) {
                TH_MATH_ERROR("th::curl", x.size(), INVALID_ARGUMENT);
                return vec<real, N>(nan(), x.size());
            }
 
            mat<real, N, N> J = jacobian<N, N>(f, x);
            J.resize(3, 3);
 
            vec<real, N> res;
            res.resize(3);
 
            res[0] = J(2, 1) - J(1, 2);
            res[1] = J(0, 2) - J(2, 0);
            res[2] = J(1, 0) - J(0, 1);
 
            return res;
        }
 
 
        template<unsigned int N = 0>
        inline auto curl(vec<multidual<N>, N>(*f)(vec<multidual<N>, N>)) {
 
            return [f](vec<real, N> x) {
                return curl(f, x);
            };
        }
 
 
        template<unsigned int N = 0>
        inline vec<real, N> directional_derivative(multidual<N>(*f)(vec<multidual<N>, N>),
            const vec<real, N>& x, const vec<real, N>& v) {
 
            return v * dot(v, gradient(f, x));
        }
 
 
        template<unsigned int N = 0>
        inline auto directional_derivative(
            multidual<N>(*f)(vec<multidual<N>, N>), const vec<real, N>& v) {
 
            return [f, v](vec<real, N> x) {
                return directional_derivative(f, x, v);
            };
        }
 
 
        template<unsigned int N = 0>
        inline real laplacian(dual2(*f)(vec<dual2, N>), const vec<real, N>& x) {
 
            real res = 0;
            vec<dual2, N> d;
            d.resize(x.size());
 
            for (unsigned int i = 0; i < x.size(); ++i)
                d[i].a = x[i];
 
            for (unsigned int i = 0; i < x.size(); ++i) {
                d[i].b = 1;
                res += f(d).Dual2();
                d[i].b = 0;
            }
 
            return res;
        }
 
 
        template<unsigned int N = 0>
        inline auto laplacian(dual2(*f)(vec<dual2, N>)) {
 
            return [f](vec<real, N> x) {
                return laplacian(f, x);
            };
        }
 
 
        template<unsigned int N = 0>
        inline real sturm_liouville(
            multidual<N>(*f)(vec<multidual<N>, N>),
            multidual<N>(*H)(vec<multidual<N>, N>),
            vec<real, N> eta) {
 
            return gradient(f, eta)
                 * mat<real, N, N>::symplectic(eta.size(), eta.size())
                 * gradient(H, eta);
        }
 
    }
}
 
#endif