Program Listing for File NonlinearSolver.h

Program Listing for File NonlinearSolver.h#

Return to documentation for file (include/Karana/Math/NonlinearSolver.h)

/*
 * Copyright (c) 2024-2026 Karana Dynamics Pty Ltd. All rights reserved.
 *
 * NOTICE TO USER:
 *
 * This source code and/or documentation (the "Licensed Materials") is
 * the confidential and proprietary information of Karana Dynamics Inc.
 * Use of these Licensed Materials is governed by the terms and conditions
 * of a separate software license agreement between Karana Dynamics and the
 * Licensee ("License Agreement"). Unless expressly permitted under that
 * agreement, any reproduction, modification, distribution, or disclosure
 * of the Licensed Materials, in whole or in part, to any third party
 * without the prior written consent of Karana Dynamics is strictly prohibited.
 *
 * THE LICENSED MATERIALS ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND.
 * KARANA DYNAMICS DISCLAIMS ALL WARRANTIES, EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO WARRANTIES OF MERCHANTABILITY, NON-INFRINGEMENT, AND
 * FITNESS FOR A PARTICULAR PURPOSE.
 *
 * IN NO EVENT SHALL KARANA DYNAMICS BE LIABLE FOR ANY DAMAGES WHATSOEVER,
 * INCLUDING BUT NOT LIMITED TO LOSS OF PROFITS, DATA, OR USE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGES, WHETHER IN CONTRACT, TORT,
 * OR OTHERWISE ARISING OUT OF OR IN CONNECTION WITH THE LICENSED MATERIALS.
 *
 * U.S. Government End Users: The Licensed Materials are a "commercial item"
 * as defined at 48 C.F.R. 2.101, and are provided to the U.S. Government
 * only as a commercial end item under the terms of this license.
 *
 * Any use of the Licensed Materials in individual or commercial software must
 * include, in the user documentation and internal source code comments,
 * this Notice, Disclaimer, and U.S. Government Use Provision.
 */

/**
 * @file
 * @brief Contains the declarations for the NonLinearSolver class and associated Eigen functors.
 */

#pragma once

#include "Karana/KCore/LockingBase.h"
#include "Karana/Math/Defs.h"
#include <functional>
#include <memory>
#include <unsupported/Eigen/NonLinearOptimization>
#include <unsupported/Eigen/NumericalDiff>
#include <variant>

namespace Karana::Math {
    namespace kc = Karana::Core;

    using cost_fn = std::function<void(const Vec &, Vec &)>; ///< Nonlinear function to solve.
    using jac_fn =
        std::function<void(const Vec &, Mat &)>; ///< Jacobian for the nonlinear function.

    /**
     * @brief Generic functor base class required by Eigen's nonlinear solver.
     * @tparam Scalar_ Scalar type.
     * @tparam NX Number of inputs at compile time.
     * @tparam NY Number of values at compile time.
     */
    template <typename Scalar_ = double, int NX = Eigen::Dynamic, int NY = Eigen::Dynamic>
    struct Functor {
        /// Scalar type used by the functor.
        using Scalar = Scalar_;
        enum {
            InputsAtCompileTime = NX, // NOLINT(readability-identifier-naming)
            ValuesAtCompileTime = NY  // NOLINT(readability-identifier-naming)
        };

        /// Input type used by the functor. It has type Scalar and is size NX x 1
        using InputType = Eigen::Matrix<Scalar, InputsAtCompileTime, 1>;

        /// Value (output) type used by the functor. It has type Scalar and is size NY x 1
        using ValueType = Eigen::Matrix<Scalar, ValuesAtCompileTime, 1>;

        /// Jacobian type used by the functor. It has type Scalar and is size NY x Nx
        using JacobianType = Eigen::Matrix<Scalar, ValuesAtCompileTime, InputsAtCompileTime>;

        /// Internal variable that holds the number of inputs
        const int _inputs; // NOLINT(readability-identifier-naming)

        /// Internal variable that holds the number of values (outputs)
        const int _values; // NOLINT(readability-identifier-naming)

        /**
         * @brief Constructor for the Functor.
         * @param inputs Number of inputs.
         * @param values Number of outputs.
         */
        Functor(int inputs, int values)
            : _inputs(inputs)
            , _values(values) {}

        /**
         * @brief Get the number of inputs.
         * @return Number of inputs.
         */
        int inputs() const { return _inputs; }

        /**
         * @brief Get the number of outputs.
         * @return Number of outputs.
         */
        int values() const { return _values; }
    };

    /**
     * @brief Functor for dynamic Eigen matrices with user-supplied cost and Jacobian functions.
     */
    struct DynamicFunctor : Functor<double, Eigen::Dynamic, Eigen::Dynamic> {
        /// Alias to the base functor
        using Base = Functor<double, Eigen::Dynamic, Eigen::Dynamic>;

        /// The cost function
        cost_fn _f; // NOLINT(readability-identifier-naming)

        /// The Jacobian function
        jac_fn _j; // NOLINT(readability-identifier-naming)

        /**
         * @brief Constructor for DynamicFunctor.
         * @param dim1 Number of inputs.
         * @param dim2 Number of outputs.
         * @param f Cost function.
         * @param j Jacobian function.
         */
        DynamicFunctor(int dim1, int dim2, cost_fn f, jac_fn j)
            : Base(dim1, dim2)
            , _f(f)
            , _j(j) {}

        /**
         * @brief Evaluate the cost function.
         * @param x Input vector.
         * @param fvec Output vector.
         * @return Status code (0 for success).
         */
        int operator()(const Vec &x, Vec &fvec) const {
            _f(x, fvec);
            return 0;
        }

        /**
         * @brief Evaluate the Jacobian.
         * @param x Input vector.
         * @param fjac Output Jacobian matrix.
         * @return Status code (0 for success).
         */
        int df(const Vec &x, Base::JacobianType &fjac) const {
            // We must do a copy here, since we use row-major and Eigen's non-linear solvers use
            // column major. If in the future they support row-major order in the non-linear
            // solvers, then modify this to remove the copy.
            Mat dark(fjac.rows(), fjac.cols());
            _j(x, dark);
            fjac = dark;
            return 0;
        }
    };

    /**
     * @brief Functor for dynamic Eigen matrices with no Jacobian provided.
     */
    struct DynamicFunctorNoJac : Functor<double, Eigen::Dynamic, Eigen::Dynamic> {
        /// Alias to the base functor
        using Base = Functor<double, Eigen::Dynamic, Eigen::Dynamic>;

        /// The cost function
        cost_fn _f; // NOLINT(readability-identifier-naming)

        /// The Jacobian function
        jac_fn _j; // NOLINT(readability-identifier-naming)

        /**
         * @brief Constructor for DynamicFunctorNoJac.
         * @param dim1 Number of inputs.
         * @param dim2 Number of outputs.
         * @param f Cost function.
         */
        DynamicFunctorNoJac(int dim1, int dim2, cost_fn f)
            : Base(dim1, dim2)
            , _f(f) {}

        /**
         * @brief Evaluate the cost function.
         * @param x Input vector.
         * @param fvec Output vector.
         * @return Status code (0 for success).
         */
        int operator()(const Vec &x, Vec &fvec) const {
            _f(x, fvec);
            return 0;
        }
    };

    /**
     * @brief Numerical difference of the non-jacobian Function. Here, Jacobians are calculated
     *        using Eigen's NumericalDiff package.
     */
    using DynamicFunctorNumJac = Eigen::NumericalDiff<DynamicFunctorNoJac>;

    /// The functor types that we accept for the NonlinearSolver.
    using functor_type = std::variant<DynamicFunctor, DynamicFunctorNumJac>;

    /// Combinations of the the solver types and functors for the for the Nonlinear solver
    using solver_type = std::variant<LevenbergMarquardt<DynamicFunctor>,
                                     LevenbergMarquardt<DynamicFunctorNumJac>,
                                     HybridNonLinearSolver<DynamicFunctor>,
                                     HybridNonLinearSolver<DynamicFunctorNumJac>>;

    /// Solver types used by the Nonlinear solver
    enum SolverType { LEVENBERG_MARQUARDT, HYBRID_NON_LINEAR_SOLVER };

    /**
     * @brief NonlinearSolver class manages a nonlinear problem and associated solver.
     */
    class NonlinearSolver : public kc::LockingBase {
      public:
        /**
         * @brief NonlinearSolver constructor. This is a container that holds the information needed
         * to solve a nonlinear problem. However, the solver itself is not added during
         * construction. You must either pass in a solver with setSolver or create one with
         * createSolver.
         *
         * @param name The name of the nonlinear solver.
         * @param input_dim The number of inputs.
         * @param value_dim The number of values (outputs).
         * @param f The cost function. This should take in a vector of inputs and output a
         * vector of outputs.
         * @param j The jacobian function. This is optional, if not specified, the jacobian
         * will be computed by forward numerical differentiation.
         */
        NonlinearSolver(
            std::string_view name, int input_dim, int value_dim, cost_fn f, jac_fn j = nullptr);

        /**
         * @brief NonlinearSolver constructor. This is a container that holds the information needed
         * to solve a nonlinear problem. However, the solver itself is not added during
         * construction. You must either pass in a solver with setSolver or create one with
         * createSolver.
         *
         * @param name The name of the Nonlinear solver.
         * @param input_dim The number of inputs.
         * @param value_dim The number of values (outputs).
         * @param f The cost function. This should take in a vector of inputs and output a
         * vector of outputs.
         * @param j The jacobian function. This is optional, if not specified, the jacobian
         * will be computed by forward numerical differentiation.
         * @return A ks_ptr to a new Nonlinear solver instance.
         */
        static kc::ks_ptr<NonlinearSolver>
        create(std::string_view name, int input_dim, int value_dim, cost_fn f, jac_fn j = nullptr);

        /**
         * @brief Get the cost function.
         * @returns The cost function.
         */
        cost_fn getF();

        /**
         * @brief Set the cost function. After setting the function, you will need to re-create the
         *        solver if one has already been set. Otherwise, the solver will still retain the
         * old cost function.
         * @param f The cost function.
         */
        void setF(cost_fn f);

        /**
         * @brief Get the jacobian function.
         * @returns The jacobian function.
         */
        jac_fn getJac();

        /**
         * @brief Set the jacobian function. After setting the jacobian.function, you will need to
         * re-create the solver if one has already been set. Otherwise, the solver will still retain
         * the old cost function.
         * @param jac The jacobian function.
         */
        void setJac(jac_fn jac);

        /**
         * @brief Get the non-linear solver.
         * @returns The current non-linear solver.
         */
        kc::ks_ptr<solver_type> getSolver();

        /**
         * @brief Set the solver.
         * @param solver The solver.
         */
        void setSolver(const kc::ks_ptr<solver_type> &solver);

        /**
         * @brief Create a solver of the given type. This will return a copy of the shared pointer
         * to the solver as well as set the solver for this NonlinearSolver.
         * @param solver The solver type to create.
         * @return The solver that was created.
         */
        kc::ks_ptr<solver_type> createSolver(const SolverType &solver);

        /**
         * @brief Get the current Jacobian matrix being used by the solver.
         *
         * @return The current Jacobian matrix being used by the solver.
         */
        Mat getCachedJacobian();

        /**
         * @brief Solve the nonlinear system.
         *
         * @param x The initial guess.
         * @return The norm squared of the final cost error. */
        double solve(Vec &x);

        /**
         * @brief The status of the nonlinear solver.
         */
        std::variant<Eigen::LevenbergMarquardtSpace::Status,
                     Eigen::HybridNonLinearSolverSpace::Status>
            status;

      protected:
        /**
         * @brief Internal method to ensure the solver is ready.
         */
        void _makeHealthy() final;

      private:
        /**
         * @brief Internal method to construct the functor.
         */
        void _makeFunctor();

        int _input_dim;
        int _value_dim;
        cost_fn _f;
        jac_fn _j;
        kc::ks_ptr<solver_type> _solver = nullptr;
        std::unique_ptr<functor_type> _func;

        /**
         * @brief Internal helper to create a solver with an analytic Jacobian functor.
         * @param solver Solver type.
         * @param func Functor with Jacobian.
         */
        void _createSolver(SolverType solver, DynamicFunctor &func);

        /**
         * @brief Internal helper to create a solver with numerical Jacobian functor.
         * @param solver Solver type.
         * @param func Functor without Jacobian.
         */
        void _createSolver(SolverType solver, DynamicFunctorNumJac &func);

        /**
         * @brief Internal solver execution using Levenberg-Marquardt with analytic Jacobian.
         * @param s Solver instance.
         * @param x Input/output vector.
         */
        void _solve(Eigen::LevenbergMarquardt<DynamicFunctor> &s, Vec &x);

        /**
         * @brief Internal solver execution using Levenberg-Marquardt with numeric Jacobian.
         * @param s Solver instance.
         * @param x Input/output vector.
         */
        void _solve(Eigen::LevenbergMarquardt<DynamicFunctorNumJac> &s, Vec &x);

        /**
         * @brief Internal solver execution using HybridNonLinearSolver with analytic Jacobian.
         * @param s Solver instance.
         * @param x Input/output vector.
         */
        void _solve(Eigen::HybridNonLinearSolver<DynamicFunctor> &s, Vec &x);

        /**
         * @brief Internal solver execution using HybridNonLinearSolver with numeric Jacobian.
         * @param s Solver instance.
         * @param x Input/output vector.
         */
        void _solve(Eigen::HybridNonLinearSolver<DynamicFunctorNumJac> &s, Vec &x);
    };
} // namespace Karana::Math