diff options
Diffstat (limited to 'firmware/os/algos/common/math/levenberg_marquardt.c')
| -rw-r--r-- | firmware/os/algos/common/math/levenberg_marquardt.c | 270 |
1 files changed, 270 insertions, 0 deletions
diff --git a/firmware/os/algos/common/math/levenberg_marquardt.c b/firmware/os/algos/common/math/levenberg_marquardt.c new file mode 100644 index 00000000..9c179ac0 --- /dev/null +++ b/firmware/os/algos/common/math/levenberg_marquardt.c @@ -0,0 +1,270 @@ +#include "common/math/levenberg_marquardt.h" + +#include <stdbool.h> +#include <stdio.h> +#include <string.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); +} |