text: #include "common/math/levenberg_marquardt.h"
  
  #include <stdbool.h>
  #include <stdio.h>
  #include <string.h>
  
  #include "common/math/macros.h"
  #include "common/math/mat.h"
  #include "common/math/vec.h"
  
  // FORWARD DECLARATIONS
  ////////////////////////////////////////////////////////////////////////
  static bool checkRelativeStepSize(const float *step, const float *state,
                                    size_t dim, float relative_error_threshold);
  
  static bool computeResidualAndGradients(ResidualAndJacobianFunction func,
                                          const float *state, const void *f_data,
                                          float *jacobian,
                                          float gradient_threshold,
                                          size_t state_dim, size_t meas_dim,
                                          float *residual, float *gradient,
                                          float *hessian);
  
  static bool computeStep(const float *gradient, float *hessian, float *L,
                          float damping_factor, size_t dim, float *step);
  
  const static float kEps = 1e-10f;
  
  // FUNCTION IMPLEMENTATIONS
  ////////////////////////////////////////////////////////////////////////
  void lmSolverInit(struct LmSolver *solver, const struct LmParams *params,
                    ResidualAndJacobianFunction func) {
    ASSERT_NOT_NULL(solver);
    ASSERT_NOT_NULL(params);
    ASSERT_NOT_NULL(func);
    memset(solver, 0, sizeof(struct LmSolver));
    memcpy(&solver->params, params, sizeof(struct LmParams));
    solver->func = func;
    solver->num_iter = 0;
  }
  
  void lmSolverDestroy(struct LmSolver *solver) {
    (void)solver;
  }
  
  void lmSolverSetData(struct LmSolver *solver, struct LmData *data) {
    ASSERT_NOT_NULL(solver);
    ASSERT_NOT_NULL(data);
    solver->data = data;
  }
  
  enum LmStatus lmSolverSolve(struct LmSolver *solver, const float *initial_state,
                              void *f_data, size_t state_dim, size_t meas_dim,
                              float *state) {
    // Initialize parameters.
    float damping_factor = 0.0f;
    float v = 2.0f;
  
    // Check dimensions.
    if (meas_dim > MAX_LM_MEAS_DIMENSION || state_dim > MAX_LM_STATE_DIMENSION) {
      return INVALID_DATA_DIMENSIONS;
    }
  
    // Check pointers (note that f_data can be null if no additional data is
    // required by the error function).
    ASSERT_NOT_NULL(solver);
    ASSERT_NOT_NULL(initial_state);
    ASSERT_NOT_NULL(state);
    ASSERT_NOT_NULL(solver->data);
  
    // Allocate memory for intermediate variables.
    float state_new[MAX_LM_STATE_DIMENSION];
    struct LmData *data = solver->data;
  
    // state = initial_state, num_iter = 0
    memcpy(state, initial_state, sizeof(float) * state_dim);
    solver->num_iter = 0;
  
    // Compute initial cost function gradient and return if already sufficiently
    // small to satisfy solution.
    if (computeResidualAndGradients(solver->func, state, f_data, data->temp,
                                    solver->params.gradient_threshold, state_dim,
                                    meas_dim, data->residual,
                                    data->gradient,
                                    data->hessian)) {
      return GRADIENT_SUFFICIENTLY_SMALL;
    }
  
    // Initialize damping parameter.
    damping_factor = solver->params.initial_u_scale *
        matMaxDiagonalElement(data->hessian, state_dim);
  
    // Iterate solution.
    for (solver->num_iter = 0;
         solver->num_iter < solver->params.max_iterations;
         ++solver->num_iter) {
  
      // Compute new solver step.
      if (!computeStep(data->gradient, data->hessian, data->temp, damping_factor,
                       state_dim, data->step)) {
        return CHOLESKY_FAIL;
      }
  
      // If the new step is already sufficiently small, we have a solution.
      if (checkRelativeStepSize(data->step, state, state_dim,
                                solver->params.relative_step_threshold)) {
        return RELATIVE_STEP_SUFFICIENTLY_SMALL;
      }
  
      // state_new = state + step.
      vecAdd(state_new, state, data->step, state_dim);
  
      // Compute new cost function residual.
      solver->func(state_new, f_data, data->residual_new, NULL);
  
      // Compute ratio of expected to actual cost function gain for this step.
      const float gain_ratio = computeGainRatio(data->residual,
                                                data->residual_new,
                                                data->step, data->gradient,
                                                damping_factor, state_dim,
                                                meas_dim);
  
      // If gain ratio is positive, the step size is good, otherwise adjust
      // damping factor and compute a new step.
      if (gain_ratio > 0.0f) {
        // Set state to new state vector: state = state_new.
        memcpy(state, state_new, sizeof(float) * state_dim);
  
        // Check if cost function gradient is now sufficiently small,
        // in which case we have a local solution.
        if (computeResidualAndGradients(solver->func, state, f_data, data->temp,
                                        solver->params.gradient_threshold,
                                        state_dim, meas_dim, data->residual,
                                        data->gradient, data->hessian)) {
          return GRADIENT_SUFFICIENTLY_SMALL;
        }
  
        // Update damping factor based on gain ratio.
        // Note, this update logic comes from Equation 2.21 in the following:
        // [Madsen, Kaj, Hans Bruun Nielsen, and Ole Tingleff.
        // "Methods for non-linear least squares problems." (2004)].
        const float tmp = 2.f * gain_ratio - 1.f;
        damping_factor *= NANO_MAX(0.33333f, 1.f - tmp * tmp * tmp);
        v = 2.f;
      } else {
        // Update damping factor and try again.
        damping_factor *= v;
        v *= 2.f;
      }
    }
  
    return HIT_MAX_ITERATIONS;
  }
  
  float computeGainRatio(const float *residual, const float *residual_new,
                         const float *step, const float *gradient,
                         float damping_factor, size_t state_dim,
                         size_t meas_dim) {
    // Compute true_gain = residual' residual - residual_new' residual_new.
    const float true_gain = vecDot(residual, residual, meas_dim)
        - vecDot(residual_new, residual_new, meas_dim);
  
    // predicted gain = 0.5 * step' * (damping_factor * step + gradient).
    float tmp[MAX_LM_STATE_DIMENSION];
    vecScalarMul(tmp, step, damping_factor, state_dim);
    vecAddInPlace(tmp, gradient, state_dim);
    const float predicted_gain = 0.5f * vecDot(step, tmp, state_dim);
  
    // Check that we don't divide by zero! If denominator is too small,
    // set gain_ratio = 1 to use the current step.
    if (predicted_gain < kEps) {
      return 1.f;
    }
  
    return true_gain / predicted_gain;
  }
  
  /*
   * Tests if a solution is found based on the size of the step relative to the
   * current state magnitude. Returns true if a solution is found.
   *
   * TODO(dvitus): consider optimization of this function to use squared norm
   * rather than norm for relative error computation to avoid square root.
   */
  bool checkRelativeStepSize(const float *step, const float *state,
                             size_t dim, float relative_error_threshold) {
    // r = eps * (||x|| + eps)
    const float relative_error = relative_error_threshold *
        (vecNorm(state, dim) + relative_error_threshold);
  
    // solved if ||step|| <= r
    // use squared version of this compare to avoid square root.
    return (vecNormSquared(step, dim) <= relative_error * relative_error);
  }
  
  /*
   * Computes the residual, f(x), as well as the gradient and hessian of the cost
   * function for the given state.
   *
   * Returns a boolean indicating if the computed gradient is sufficiently small
   * to indicate that a solution has been found.
   *
   * INPUTS:
   * state: state estimate (x) for which to compute the gradient & hessian.
   * f_data: pointer to parameter data needed for the residual or jacobian.
   * jacobian: pointer to temporary memory for storing jacobian.
   *           Must be at least MAX_LM_STATE_DIMENSION * MAX_LM_MEAS_DIMENSION.
   * gradient_threshold: if gradient is below this threshold, function returns 1.
   *
   * OUTPUTS:
   * residual: f(x).
   * gradient: - J' f(x), where J = df(x)/dx
   * hessian: df^2(x)/dx^2 = J' J
   */
  bool computeResidualAndGradients(ResidualAndJacobianFunction func,
                                   const float *state, const void *f_data,
                                   float *jacobian, float gradient_threshold,
                                   size_t state_dim, size_t meas_dim,
                                   float *residual, float *gradient,
                                   float *hessian) {
    // Compute residual and Jacobian.
    ASSERT_NOT_NULL(state);
    ASSERT_NOT_NULL(residual);
    ASSERT_NOT_NULL(gradient);
    ASSERT_NOT_NULL(hessian);
    func(state, f_data, residual, jacobian);
  
    // Compute the cost function hessian = jacobian' jacobian and
    // gradient = -jacobian' residual
    matTransposeMultiplyMat(hessian, jacobian, meas_dim, state_dim);
    matTransposeMultiplyVec(gradient, jacobian, residual, meas_dim, state_dim);
    vecScalarMulInPlace(gradient, -1.f, state_dim);
  
    // Check if solution is found (cost function gradient is sufficiently small).
    return (vecMaxAbsoluteValue(gradient, state_dim) < gradient_threshold);
  }
  
  /*
   * Computes the Levenberg-Marquardt solver step to satisfy the following:
   *    (J'J + uI) * step = - J' f
   *
   * INPUTS:
   * gradient:  -J'f
   * hessian:  J'J
   * L: temp memory of at least MAX_LM_STATE_DIMENSION * MAX_LM_STATE_DIMENSION.
   * damping_factor: u
   * dim: state dimension
   *
   * OUTPUTS:
   * step: solution to the above equation.
   * Function returns false if the solution fails (due to cholesky failure),
   * otherwise returns true.
   *
   * Note that the hessian is modified in this function in order to reduce
   * local memory requirements.
   */
  bool computeStep(const float *gradient, float *hessian, float *L,
                   float damping_factor, size_t dim, float *step) {
  
    // 1) A = hessian + damping_factor * Identity.
    matAddConstantDiagonal(hessian, damping_factor, dim);
  
    // 2) Solve A * step = gradient for step.
    // a) compute cholesky decomposition of A = L L^T.
    if (!matCholeskyDecomposition(L, hessian, dim)) {
      return false;
    }
  
    // b) solve for step via back-solve.
    return matLinearSolveCholesky(step, L, gradient, dim);
  }
  