NonlinearLeastSquares.hpp

Go to the documentation of this file.

/******************************************************************************
* Copyright (c) ATTX LLC 2024. All Rights Reserved.
*
* This software and associated documentation (the "Software") are the 
* proprietary and confidential information of ATTX, LLC. The Software is 
* furnished under a license agreement between ATTX and the user organization  
* and may be used or copied only in accordance with the terms of the agreement.
* Refer to 'license/attx_license.adoc' for standard license terms.
*
* EXPORT CONTROL NOTICE: THIS SOFTWARE MAY INCLUDE CONTENT CONTROLLED UNDER THE
* INTERNATIONAL TRAFFIC IN ARMS REGULATIONS (ITAR) OR THE EXPORT ADMINISTRATION 
* REGULATIONS (EAR99). No part of the Software may be used, reproduced, or 
* transmitted in any form or by any means, for any purpose, without the express 
* written permission of ATTX, LLC.
******************************************************************************/
 
#ifndef GNC_NAVIGATION_EKF_NONLINEAR_LEAST_SQUARES_HPP
#define GNC_NAVIGATION_EKF_NONLINEAR_LEAST_SQUARES_HPP
 
#include "core/Matrix.hpp"
#include "core/CartesianVector.hpp"
#include "core/matrixmath.hpp"
#include "core/vectormath.hpp"
#include "gncutils/dynamics/Rates.hpp"
#include "gncutils/integrator/Integrator.hpp"
#include "gncutils/EKF/Measurements.hpp"
#include "gncutils/EKF/EkfMeasurementUpdate.hpp"
#include "gncutils/EKF/KFCore/linalg.h"
 
namespace warpos {

    template <uint32 N, uint32 M, uint32 O>
    class NonlinearLeastSquares {
    public:
        NonlinearLeastSquares(Integrator<N + N*N>& integ, Measurements<N, M, O> &measurement_ref);

        int16 solve(const std::array<floating_point, N> &initial_guess,
                    const clockwerk::Matrix<N, N> &P0,
                    const clockwerk::Matrix<M, M> &R, 
                    const std::array<floating_point, M> *measurements_pointer,
                    const floating_point *time_pointer, 
                    const std::array<floating_point, O> *obs_state_pointer,
                    const uint32 num_measurements, 
                    const uint32 max_iterations,
                    const floating_point tolerance, 
                    const floating_point max_step_size, 
                    clockwerk::Matrix<N, N> *covariance_estimate,
                    std::array<floating_point, N> *state_estimate);
    private:
        std::array<floating_point, N + N*N> _full_state_initial;
        std::array<floating_point, N + N*N> _full_state_post;
 
        floating_point _max_shift;
        uint32 _num_iter;
        int16 _error;
        uint32 _meas_idx;
 
        // Time values for measurement times
        floating_point t_start;
        floating_point t_current;
        floating_point t_next;
 
        // STM
        clockwerk::Matrix<N, N> _phi;
        clockwerk::Matrix<N, N> _phi_transpose;
 
        // Information matrix and vector
        clockwerk::Matrix<N, N> _tmp_chol;
        clockwerk::Matrix<N, N> _tmp_chol_inv;
        clockwerk::Matrix<N, N> _I_matrix;
        clockwerk::Matrix<N, N> _P0_inv;
        clockwerk::CartesianVector<N> _N_vector;
        clockwerk::CartesianVector<N> _dx_minus;
        std::array<floating_point, N> _dx_plus;
 
        // Current state estimate
        std::array<floating_point, N> _x0;
        std::array<floating_point, N> _x_prop;
 
        // Measurement noise covariance inverse
        clockwerk::Matrix<M, M> _R_inv;
 
        // Residual storage
        std::array<floating_point, M> _residual;
 
        // H matrix storage
        std::array<floating_point, M*N> _tmp_H_array;
        clockwerk::Matrix<M, N> _H;
        clockwerk::Matrix<N, M> _H_transpose;
 
        // Temporary storage for Cholesky
        floating_point _I_flat[N * N];
        floating_point _y_temp[N];
 
        // Temporary calculation matrices
        clockwerk::Matrix<N, N> _tmp_n_n_1;
        clockwerk::Matrix<N, N> _tmp_n_n_2;
        clockwerk::Matrix<N, M> _tmp_n_m_1;
        clockwerk::Matrix<M, N> _tmp_m_n_1;
        clockwerk::Matrix<M, M> _tmp_m_m_1;
        clockwerk::Matrix<N, 1> _tmp_n_1;
        clockwerk::Matrix<M, 1> _tmp_m_1;

        Integrator<N + N*N>& _integrator;
        
        Measurements<N, M, O>& _measurement;
        
        EkfMeasurementUpdate<N, M, O> _ekf_meas_update;
    }; 
 
    template <uint32 N, uint32 M, uint32 O>
    NonlinearLeastSquares<N, M, O>::NonlinearLeastSquares(Integrator<N + N*N>& integ, Measurements<N, M, O> &measurement_ref) 
        : _integrator(integ), _measurement(measurement_ref), _ekf_meas_update(measurement_ref) {}
 
    template<uint32 N, uint32 M, uint32 O>
    int16 NonlinearLeastSquares<N, M, O>::solve(
        const std::array<floating_point, N> &initial_guess,
        const clockwerk::Matrix<N, N> &P0, 
        const clockwerk::Matrix<M, M> &R,
        const std::array<floating_point, M> *measurements_pointer,
        const floating_point *time_pointer, 
        const std::array<floating_point, O> *obs_state_pointer,
        const uint32 num_measurements, 
        const uint32 max_iterations,
        const floating_point tolerance, 
        const floating_point max_step_size, 
        clockwerk::Matrix<N, N> *covariance_estimate,
        std::array<floating_point, N> *state_estimate) {
        
        // Verify output pointers
        if(state_estimate == nullptr || covariance_estimate == nullptr) {
            return ERROR_NULLPTR;
        }
 
        // Initialize convergence parameters
        _max_shift = tolerance + 1.0;
        _num_iter = 0;
        _error = NO_ERROR;
        
        // Initialize state estimate with initial guess
        for(uint32 i = 0; i < N; i++) {
            _x0[i] = initial_guess[i];
        }
        
        // Initialize dx_minus to zero
        for(uint32 i = 0; i < N; i++) {
            _dx_minus[i] = 0.0;
        }
        
        // Set R matrix in EKF measurement update
        _ekf_meas_update.setMeasurementNoiseMatrix(R);
 
        // Compute R inverse for use in iterations
        _error = R.inverse(_R_inv);
        if(_error) {
            return _error;
        }
 
        // Compute P0 inverse for use in iterations
        _error = P0.inverse(_P0_inv);
        if(_error) {
            return _error;
        }
        
        // Outer loop - iterate until convergence or max iterations
        while(_max_shift > tolerance && _num_iter < max_iterations) {
 
            // I = inv(P0) on first iteration
            _error = P0.chol(_tmp_chol);
            if(_error) {return _error;}
            _error = _tmp_chol.inverse(_tmp_chol_inv);
            if(_error) {return _error;}
 
            // Set I and N
            _I_matrix = _tmp_chol_inv*_tmp_chol_inv.transpose();
            _N_vector = _I_matrix*_dx_minus;
            
            // Inner loop - propagate through all measurements
            for(_meas_idx = 0; _meas_idx < num_measurements; _meas_idx++) {
                
                // Set up full state vector [x; phi(:)]
                if(_meas_idx == 0) {
                    // Initialize phi to identity on first measurement
                    unsigned int one_idx = N;
                    for(unsigned int i = 0; i < N + N*N; i++) {
                        if(i < N) {
                            _full_state_initial[i] = _x_prop[i];
                        } else if(i == one_idx) {
                            _full_state_initial[i] = 1.0;
                            one_idx = one_idx + N + 1;
                        } else {
                            _full_state_initial[i] = 0.0;
                        }
                        _full_state_post[i] = _full_state_initial[i];   // Set post to initial in case we don't propagate
                    }
                } else {
                    // For subsequent measurements, use accumulated state and phi from previous step
                    for(unsigned int i = 0; i < N + N*N; i++) {
                        _full_state_initial[i] = _full_state_post[i];
                    }
                }
                
                // Propagate state and STM from previous measurement time to current measurement time
                if (_meas_idx == 0) {
                    t_start = time_pointer[0];
                } else {
                    t_start = time_pointer[_meas_idx - 1];
                }
 
                // Propagate with max step size constraint
                t_current = t_start;
                while(t_current < time_pointer[_meas_idx]) {
                    if (time_pointer[_meas_idx] - t_current > max_step_size) {
                        t_next = t_current + max_step_size;
                    } else {
                        t_next = time_pointer[_meas_idx];
                    }
                    
                    _error = _integrator.step(t_current, t_next, _full_state_initial, _full_state_post);
                    if(_error > 0) {
                        return _error;
                    }
                    
                    // Update for next sub-step
                    for(unsigned int i = 0; i < N + N*N; i++) {
                        _full_state_initial[i] = _full_state_post[i];
                    }
                    t_current = t_next;
                }
                
                // Extract propagated state
                for(uint32 i = 0; i < N; i++) {
                    _x_prop[i] = _full_state_post[i];
                }
                
                // Extract STM (phi)
                _phi.setFromArray(&_full_state_post[N]);
                _phi.transpose(_phi_transpose);
                
                // Use EKF measurement update to generate residual
                _error = _ekf_meas_update.generateResidual(
                    time_pointer[_meas_idx], 
                    _x_prop, 
                    obs_state_pointer[_meas_idx],
                    measurements_pointer[_meas_idx],
                    &_residual
                );
                if(_error > 0) {
                    return _error;
                }
                
                // Get H matrix (measurement Jacobian) from measurement reference
                _error = _measurement.calculateMeasurementsMatrix(
                    time_pointer[_meas_idx], 
                    _x_prop, 
                    obs_state_pointer[_meas_idx], 
                    _tmp_H_array.data()
                );
                if(_error) {
                    return _error;
                }
                _H.setFromArray(&_tmp_H_array[0]);
                _H.transpose(_H_transpose);
                
                // Accumulate information matrix: I = I + phi^T * H^T * R^{-1} * H * phi
                // Step 1: H * phi -> _tmp_m_n_1
                clockwerk::multiply(_H, _phi, _tmp_m_n_1);
                
                // Step 2: R^{-1} * (H * phi) -> _tmp_m_n_1 (overwrite)
                clockwerk::multiply(_R_inv, _tmp_m_n_1, _tmp_m_n_1);
                
                // Step 3: H^T * R^{-1} * H * phi -> _tmp_n_n_1
                clockwerk::multiply(_H_transpose, _tmp_m_n_1, _tmp_n_n_1);
                
                // Step 4: phi^T * H^T * R^{-1} * H * phi -> _tmp_n_n_1 (overwrite)
                clockwerk::multiply(_phi_transpose, _tmp_n_n_1, _tmp_n_n_1);
                
                // Step 5: I = I + result
                clockwerk::eAdd(_I_matrix, _tmp_n_n_1, _I_matrix);
                
                // Accumulate information vector: N = N + phi^T * H^T * R^{-1} * r
                // Step 1: R^{-1} * r -> _tmp_m_1
                _tmp_m_1.setFromArray(&_residual[0]);
                clockwerk::multiply(_R_inv, _tmp_m_1, _tmp_m_1);
                
                // Step 2: H^T * R^{-1} * r -> _tmp_n_1
                clockwerk::multiply(_H_transpose, _tmp_m_1, _tmp_n_1);
                
                // Step 3: phi^T * H^T * R^{-1} * r -> _tmp_n_1
                clockwerk::multiply(_phi_transpose, _tmp_n_1, _tmp_n_1);
                
                // Step 4: N = N + result
                for(uint32 i = 0; i < N; i++) {
                    _N_vector[i] += _tmp_n_1.get(i, 0);
                }
 
            }
 
            _error = _I_matrix.chol(_tmp_chol);
            if(_error) {return _error;}
            _error = _tmp_chol.inverse(_tmp_chol_inv);
            if(_error) {return _error;}
            _tmp_n_n_1 = _tmp_chol_inv*_tmp_chol_inv.transpose();
 
            // Compute dx_plus = P * N
            for(uint32 i = 0; i < N; i++) {
                _dx_plus[i] = 0.0;
                for(uint32 j = 0; j < N; j++) {
                    _dx_plus[i] += _tmp_n_n_1[i][j]*_N_vector[j];
                }
            }
            
            // Update initial guess: x0 = initial_guess + dx_plus
            for(uint32 i = 0; i < N; i++) {
                _x0[i] = initial_guess[i] + _dx_plus[i];
            }
            
            // Update dx_minus: dx_minus = dx_minus + dx_plus
            for(uint32 i = 0; i < N; i++) {
                _dx_minus[i] = _dx_minus[i] + _dx_plus[i];
            }
            
            // Calculate maximum shift for convergence check
            _max_shift = 0.0;
            for(uint32 i = 0; i < N; i++) {
                floating_point abs_shift = _dx_plus[i] >= 0.0 ? _dx_plus[i] : -_dx_plus[i];
                if(abs_shift > _max_shift) {
                    _max_shift = abs_shift;
                }
            }
            
            _num_iter++;
        }
 
        // Set output state estimate
        for(uint32 i = 0; i < N; i++) {
            (*state_estimate)[i] = _x0[i]-initial_guess[i];
        }
        
        // Set output covariance (final P = inv(I))
        _error = _I_matrix.chol(_tmp_chol);
        if(_error) {return _error;}
        _error = _tmp_chol.inverse(_tmp_chol_inv);
        if(_error) {return _error;}
        *covariance_estimate = _tmp_chol_inv*_tmp_chol_inv.transpose();
 
        return NO_ERROR;
    }
 
}
 
#endif