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 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 MultidualType = return_type_t<Function>;
 
            constexpr size_t N = MultidualType::size_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 f(arg).Dual();
        }
 
 
        template <
            typename Function,
            enable_scalar_field<Function> = true
        >
        inline auto gradient(Function f) {
 
            constexpr size_t N = return_type_t<Function>::size_argument;
 
            return [f](const vec<real, N>& x) {
                return gradient(f, x);
            };
        }
 
 
        template <
            typename Function, typename Vector = vec<real>,
            enable_vector_field<Function> = true,
            enable_vector<Vector> = true
        >
        inline real divergence(Function V, const Vector& x) {
 
            using MultidualT = return_type_t<Function>;
            const size_t N = MultidualT::size_argument;
 
            // Sum the diagonal elements of the jacobian
            const vec<multidual<N>, N> res = V(multidual<N>::make_argument(x));
 
            real div = 0.0;
            for (unsigned int i = 0; i < res.size(); ++i) {
                
                // dreal numbers initialized to a scalar
                // have an empty dual part, check the size before summing
                if (res[i].v.size() > i)
                    div += res[i].Dual()[i];
            }
 
            return div;
        }
 
 
        template <
            typename Function,
            enable_scalar_field<Function> = true
        >
        inline auto divergence(Function V) {
 
            constexpr size_t N = return_type_t<Function>::size_argument;
            using Vector = vec<real, N>;
 
            return [V](const Vector& x) {
                return divergence(V, x);
            };
        }
 
 
        template <
            typename MultidualFunction, typename Vector,
            enable_vector<Vector> = true,
            enable_vector_field<MultidualFunction> = true
        >
        inline auto jacobian(MultidualFunction f, const Vector& x) {
 
            constexpr size_t M = return_type_t<MultidualFunction>::size_argument;
            constexpr size_t N = _internal::func_helper<MultidualFunction>::first_arg_type::size_argument;
 
            const vec<multidual<N>, N> 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) {
 
                const unsigned int dual_size = res[j].v.size();
 
                if (dual_size > J.cols()) {
                    TH_MATH_ERROR("autodiff::jacobian", dual_size, MathError::InvalidArgument);
                    return algebra::mat_error(J);
                }
 
                unsigned int i = 0;
                for (; i < dual_size; ++i)
                    J(j, i) = res[j].v[i];
 
                for (; i < J.cols(); ++i)
                    J(j, i) = 0.0;
            }
 
            return J;
        }
 
 
        template<
            typename MultidualFunction,
            enable_vector_field<MultidualFunction> = true
        >
        inline auto jacobian(MultidualFunction f) {
 
            constexpr size_t N = return_type_t<MultidualFunction>::size_argument;
            using Vector = vec<real, N>;
 
            return [f](const Vector& x) {
                return jacobian(f, x);
            };
        }
 
 
        template <
            typename MultidualFunction, typename Vector,
            enable_vector<Vector> = true,
            enable_vector_field<MultidualFunction> = true
        >
        inline auto curl(MultidualFunction f, const Vector& x) {
 
            Vector res;
            res.resize(3);
 
            if(x.size() != 3) {
                TH_MATH_ERROR("th::curl", x.size(), MathError::InvalidArgument);
                return algebra::vec_error(res);
            }
 
            constexpr size_t N = return_type_t<MultidualFunction>::size_argument;
            const mat<real, N, N> J = jacobian(f, x);
 
            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<
            typename MultidualFunction,
            enable_vector_field<MultidualFunction> = true
        >
        inline auto curl(MultidualFunction f) {
 
            constexpr size_t N = return_type_t<MultidualFunction>::size_argument;
            using Vector = vec<real, N>;
 
            return [f](const Vector& x) {
                return curl(f, x);
            };
        }
 
 
        template <
            typename Dual2Function, typename Vector,
            enable_vector<Vector> = true
        >
        inline real laplacian(Dual2Function f, const Vector& x) {
 
            // Extract size template argument of the input vector
            constexpr size_t N = _internal::func_helper<Dual2Function>
                ::first_arg_type::size_argument;
 
            vec<dual2, N> d;
            d.resize(x.size());
 
            for (unsigned int i = 0; i < d.size(); ++i)
                d[i].a = x[i];
 
            // Accumulate second derivatives
            real res = 0;
            for (unsigned int i = 0; i < d.size(); ++i) {
                d[i].b = 1;
                res += f(d).Dual2();
                d[i].b = 0;
            }
 
            return res;
        }
 
 
        template <typename Dual2Function>
        inline auto laplacian(Dual2Function f) {
 
            constexpr size_t N = _internal::func_helper<Dual2Function>
                ::first_arg_type::size_argument;
 
            return [f](const vec<real, N>& x) {
                return laplacian(f, x);
            };
        }
 
    }
}
 
#endif