ExtendedKalmanFilter.h

#pragma once
 
#include "KalmanFilterBase.h"
#include "StateSpaceUtil.h"
 
#include <Eigen/Core>
#include <Eigen/LU>
 
namespace filters {

template <int stateDim, int inputDim, int outputDim, int processNoiseDim, int outputNoiseDim>
class ExtendedKalmanFilter : public KalmanFilterBase<stateDim, inputDim> {
public:
    using state_t = Eigen::Matrix<double, stateDim, 1>;
    using output_t = Eigen::Matrix<double, outputDim, 1>;
    using input_t = Eigen::Matrix<double, inputDim, 1>;
    using statefunc_t = std::function<state_t(
        const state_t&, const input_t&, const Eigen::Matrix<double, processNoiseDim, 1>&)>;
    using outputfunc_t = std::function<output_t(
        const state_t&, const Eigen::Matrix<double, outputNoiseDim, 1>&)>;
    using processnoise_t = Eigen::Matrix<double, processNoiseDim, 1>;
    using outputnoise_t = Eigen::Matrix<double, outputNoiseDim, 1>;

    ExtendedKalmanFilter(const statefunc_t& stateFunc, const outputfunc_t& outputFunc,
                         const statespace::NoiseCovMat<stateDim, processNoiseDim, inputDim>& processNoise,
                         const statespace::NoiseCovMat<stateDim, outputNoiseDim, outputDim>& outputNoise,
                         double dt)
        : stateFunc(stateFunc), outputFunc(outputFunc), Q(processNoise), R(outputNoise),
          dt(dt) {}
 
    void predict(const Eigen::Matrix<double, inputDim, 1>& input) override {
        Eigen::Matrix<double, stateDim, stateDim> F =
            getStateFuncJacobianX(stateFunc, this->xHat, input);
        Eigen::Matrix<double, processNoiseDim, processNoiseDim> processNoise =
            Q.get(this->xHat, input);
 
        Eigen::Matrix<double, stateDim, processNoiseDim> L =
            getStateFuncJacobianW(stateFunc, this->xHat, input);
 
        this->xHat = stateFunc(this->xHat, input, processnoise_t::Zero());
        this->P = F * this->P * F.transpose() + L * processNoise * L.transpose();
    }

    void correct(const Eigen::Matrix<double, outputDim, 1>& measurement) {
        Eigen::Matrix<double, outputDim, stateDim> H =
            getOutputFuncJacobianX(outputFunc, this->xHat); // output matrix
 
        Eigen::Matrix<double, outputNoiseDim, outputNoiseDim> outputNoise =
            R.get(this->xHat, measurement);
        Eigen::Matrix<double, outputDim, outputNoiseDim> M =
            getOutputFuncJacobianV(outputFunc, this->xHat);
        Eigen::Matrix<double, outputDim, outputDim> S =
            H * this->P * H.transpose() +
            M * outputNoise * M.transpose(); // residual covariance
        // near-optimal kalman gain
        Eigen::Matrix<double, stateDim, outputDim> K =
            S.transpose().colPivHouseholderQr().solve(H * this->P.transpose()).transpose();
        output_t y = measurement -
                     outputFunc(this->xHat, outputnoise_t::Zero()); // measurement residual
 
        this->xHat = this->xHat + K * y;
        this->P = (Eigen::Matrix<double, stateDim, stateDim>::Identity() - K * H) * this->P;
    }

    std::function<Eigen::Matrix<double, stateDim, stateDim>(const state_t&, const input_t&,
                                                            const processnoise_t&)>
        stateFuncJacobianX;
    std::function<Eigen::Matrix<double, stateDim, processNoiseDim>(
        const state_t&, const input_t&, const processnoise_t&)>
        stateFuncJacobianW;
    std::function<Eigen::Matrix<double, outputDim, stateDim>(const state_t&,
                                                             const outputnoise_t&)>
        outputFuncJacobianX;
    std::function<Eigen::Matrix<double, outputDim, outputNoiseDim>(const state_t&,
                                                                   const outputnoise_t&)>
        outputFuncJacobianV;
 
private:
    statefunc_t stateFunc;
    outputfunc_t outputFunc;
    // process noise and measurement noise covariance matrices
    statespace::NoiseCovMat<stateDim, processNoiseDim, inputDim> Q;
    statespace::NoiseCovMat<stateDim, outputNoiseDim, outputDim> R;
    double dt;
    static constexpr double epsilon = 1e-5;
 
    Eigen::Matrix<double, stateDim, stateDim>
    getStateFuncJacobianX(const statefunc_t& f, const state_t& x, const input_t& u) const {
        processnoise_t w = processnoise_t::Zero();
        // If we have an analytic solution, use that
        if (stateFuncJacobianX) {
            return stateFuncJacobianX(x, u, w);
        } else {
            Eigen::Matrix<double, stateDim, stateDim> jacobian;
            for (int i = 0; i < stateDim; i++) {
                state_t delta = state_t::Zero();
                delta[i] = epsilon;
                state_t derivative = (f(x + delta, u, w) - f(x - delta, u, w)) / (2 * epsilon);
                jacobian.col(i) = derivative;
            }
            return jacobian;
        }
    }
 
    Eigen::Matrix<double, stateDim, processNoiseDim>
    getStateFuncJacobianW(const statefunc_t& f, const state_t& x, const input_t& u) const {
        processnoise_t w = processnoise_t::Zero();
        // If we have an analytic solution, use that
        if (stateFuncJacobianW) {
            return stateFuncJacobianW(x, u, w);
        } else {
            Eigen::Matrix<double, stateDim, processNoiseDim> jacobian;
            for (int i = 0; i < processNoiseDim; i++) {
                processnoise_t delta = processnoise_t::Zero();
                delta[i] = epsilon;
                state_t derivative = (f(x, u, w + delta) - f(x, u, w - delta)) / (2 * epsilon);
                jacobian.col(i) = derivative;
            }
            return jacobian;
        }
    }
 
    Eigen::Matrix<double, outputDim, stateDim> getOutputFuncJacobianX(const outputfunc_t& f,
                                                                      const state_t& x) const {
        outputnoise_t v = outputnoise_t::Zero();
        // If we have an analytic solution, use that
        if (outputFuncJacobianX) {
            return outputFuncJacobianX(x, v);
        } else {
            Eigen::Matrix<double, outputDim, stateDim> jacobian;
            for (int i = 0; i < stateDim; i++) {
                state_t delta = Eigen::Matrix<double, stateDim, 1>::Zero();
                delta[i] = epsilon;
                output_t derivative = (f(x + delta, v) - f(x - delta, v)) / (2 * epsilon);
                jacobian.col(i) = derivative;
            }
            return jacobian;
        }
    }
 
    Eigen::Matrix<double, outputDim, outputNoiseDim>
    getOutputFuncJacobianV(const outputfunc_t& f, const state_t& x) const {
        outputnoise_t v = outputnoise_t::Zero();
        // If we have an analytic solution, use that
        if (outputFuncJacobianV) {
            return outputFuncJacobianV(x, v);
        } else {
            Eigen::Matrix<double, outputDim, outputNoiseDim> jacobian;
            for (int i = 0; i < outputNoiseDim; i++) {
                outputnoise_t delta = outputnoise_t::Zero();
                delta[i] = epsilon;
                output_t derivative = (f(x, v + delta) - f(x, v - delta)) / (2 * epsilon);
                jacobian.col(i) = derivative;
            }
            return jacobian;
        }
    }
};
 
} // namespace filters