tiny_cost_function.h

Go to the documentation of this file.

/* Brush
copyright 2020 William La Cava
license: GNU/GPL v3
 
Code below heavily inspired by heal-research/operon, Copyright 2019-2022 Heal Research
*/
 
#ifndef TINY_OPTIMIZER
#define TINY_OPTIMIZER
 
 
namespace Brush {
using Scalar = float;
// this cost function is adapted to work with both solvers from Ceres: the normal one and the tiny solver
// for this, a number of template parameters are necessary:
// - the CostFunctor is the actual functor for computing the residuals
// - the Dual type represents a dual number, the user can specify the type for the Scalar part (float, double) and the Stride (Ceres-specific)
// - the StorageOrder specifies the format of the jacobian (row-major for the big Ceres solver, column-major for the tiny solver)
 
// TODO: eliminate this and use the Ceres tiny solver autodiff function instead:
// https://github.com/ceres-solver/ceres-solver/blob/caf614a6c1ac1717be606c37fe434391edb2f417/include/ceres/tiny_solver_autodiff_function.h
namespace detail {
    template<typename CostFunctor, typename Dual, typename Scalar, int JacobianLayout = Eigen::ColMajor>
    inline auto Autodiff(CostFunctor const& function, Scalar const* parameters, Scalar* residuals, Scalar* jacobian) -> bool
    {
        static_assert(std::is_convertible_v<typename Dual::Scalar, Scalar>, "The chosen Jet and Scalar types are not compatible.");
        static_assert(std::is_convertible_v<Scalar, typename Dual::Scalar>, "The chosen Jet and Scalar types are not compatible.");
 
        // EXPECT(parameters != nullptr);
        // EXPECT(residuals != nullptr || jacobian != nullptr);
 
        if (jacobian == nullptr) {
            return function(parameters, residuals);
        }
 
        std::vector<Dual> inputs(function.NumParameters());
        for (size_t i = 0; i < inputs.size(); ++i) {
            inputs[i].a = parameters[i];
            inputs[i].v.setZero();
        }
        std::vector<Dual> outputs(function.NumResiduals());
 
        static auto constexpr D{Dual::DIMENSION};
        Eigen::Map<Eigen::Matrix<Scalar, -1, -1, JacobianLayout>> jmap(jacobian, outputs.size(), inputs.size());
 
        for (int s = 0; s < inputs.size(); s += D) {
            int r = std::min(static_cast<int>(inputs.size()), s + D); // remaining parameters
 
            for (int i = s; i < r; ++i) {
                inputs[i].v[i - s] = 1.0;
            }
 
            if (!function(inputs.data(), outputs.data())) {
                return false;
            }
 
            for (int i = s; i < r; ++i) {
                inputs[i].v[i - s] = 0.0;
            }
 
            // fill in the jacobian trying to exploit its layout for efficiency
            if constexpr (JacobianLayout == Eigen::ColMajor) {
                for (int i = s; i < r; ++i) {
                    std::transform(outputs.cbegin(), outputs.cend(), jmap.col(i).data(), [&](auto const& jet) { return jet.v[i - s]; });
                }
            } else {
                for (auto i = 0; i < outputs.size(); ++i) {
                    std::copy_n(outputs[i].v.data(), r - s, jmap.row(i).data() + s);
                }
            }
        }
        if (residuals != nullptr) {
            std::transform(std::cbegin(outputs), std::cend(outputs), residuals, [](auto const& jet) { return jet.a; });
        }
        return true;
    }
} // namespace detail
 
template <typename CostFunctor, typename DualType=fJet, typename ScalarType=float, 
    int StorageOrder = Eigen::ColMajor>
struct TinyCostFunction {
    static constexpr int Stride = DualType::DIMENSION;
    static constexpr int Storage = StorageOrder;
    using Scalar = ScalarType;
 
    enum {
        NUM_RESIDUALS = Eigen::Dynamic,  // NOLINT
        NUM_PARAMETERS = Eigen::Dynamic, // NOLINT
    };
 
    explicit TinyCostFunction(CostFunctor const& functor)
        : functor_(functor)
    {
    }
 
    auto Evaluate(Scalar const* parameters, Scalar* residuals, Scalar* jacobian) const -> bool
    {
        return detail::Autodiff<CostFunctor, DualType, ScalarType, StorageOrder>(functor_, parameters, residuals, jacobian);
    }
 
    // ceres solver - jacobian must be in row-major format
    // ceres tiny solver - jacobian must be in col-major format
    auto operator()(Scalar const* parameters, Scalar* residuals, Scalar* jacobian) const -> bool
    {
        return Evaluate(parameters, residuals, jacobian);
    }
 
    [[nodiscard]] auto NumResiduals() const -> int { return functor_.NumResiduals(); }
    [[nodiscard]] auto NumParameters() const -> int { return functor_.NumParameters(); }
 
    // required by Eigen::LevenbergMarquardt
    using JacobianType = Eigen::Matrix<Brush::Scalar, -1, -1>;
    using QRSolver     = Eigen::ColPivHouseholderQR<JacobianType>;
 
    // there is no real documentation but looking at Eigen unit tests, these functions should return zero
    // see: https://gitlab.com/libeigen/eigen/-/blob/master/unsupported/test/NonLinearOptimization.cpp
    auto operator()(Eigen::Matrix<Scalar, -1, 1> const& input, Eigen::Matrix<Scalar, -1, 1> &residual) -> int
    {
        Evaluate(input.data(), residual.data(), nullptr);
        return 0;
    }
 
    auto df(Eigen::Matrix<Scalar, -1, 1> const& input, Eigen::Matrix<Scalar, -1, -1> &jacobian) -> int // NOLINT
    {
        static_assert(StorageOrder == Eigen::ColMajor, "Eigen::LevenbergMarquardt requires the Jacobian to be stored in column-major format.");
        Evaluate(input.data(), nullptr, jacobian.data());
        return 0;
    }
 
    [[nodiscard]] auto values() const -> int { return NumResiduals(); }  // NOLINT
    [[nodiscard]] auto inputs() const -> int { return NumParameters(); } // NOLINT
 
private:
    CostFunctor functor_;
};
} // namespace Brush
 
#endif