269 lines
10 KiB
C
269 lines
10 KiB
C
|
#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"
|
||
|
|
#include "chre/util/nanoapp/assert.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) {
|
||
|
|
CHRE_ASSERT_NOT_NULL(solver);
|
||
|
|
CHRE_ASSERT_NOT_NULL(params);
|
||
|
|
CHRE_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 lmSolverSetData(struct LmSolver *solver, struct LmData *data) {
|
||
|
|
CHRE_ASSERT_NOT_NULL(solver);
|
||
|
|
CHRE_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).
|
||
|
|
CHRE_ASSERT_NOT_NULL(solver);
|
||
|
|
CHRE_ASSERT_NOT_NULL(initial_state);
|
||
|
|
CHRE_ASSERT_NOT_NULL(state);
|
||
|
|
CHRE_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.
|
||
|
|
CHRE_ASSERT_NOT_NULL(state);
|
||
|
|
CHRE_ASSERT_NOT_NULL(residual);
|
||
|
|
CHRE_ASSERT_NOT_NULL(gradient);
|
||
|
|
CHRE_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);
|
||
|
|
}
|