MultiSensorEKF.h

#pragma once
 
#include "../navtypes.h"
#include "../utils/math.h"
#include "KalmanFilterBase.h"
#include "StateSpaceUtil.h"
 
#include <array>
#include <functional>
#include <loguru.hpp>
 
#include <Eigen/Core>
#include <Eigen/LU>
 
namespace filters {

class Output {
public:
    // (state, outputnoise) -> output
    using outputfunc_t =
        std::function<Eigen::VectorXd(const Eigen::VectorXd&, const Eigen::VectorXd&)>;
 
    const int stateDim;
    const int outputDim;
    const int outputNoiseDim;
    const outputfunc_t outputFunc;
    const statespace::NoiseCovMatX covMat;

    Output(int stateDim, int outputDim, int outputNoiseDim, const outputfunc_t& outputFunc,
           const statespace::NoiseCovMatX& covMat)
        : stateDim(stateDim), outputDim(outputDim), outputNoiseDim(outputNoiseDim),
          outputFunc(outputFunc), covMat(covMat) {
        CHECK_F(covMat.stateDim == stateDim && covMat.size == outputNoiseDim &&
                covMat.paramDim == outputDim);
    }

    Eigen::MatrixXd outputFuncJacobianX(const Eigen::VectorXd& state) {
        CHECK_EQ_F(state.size(), stateDim);
        Eigen::VectorXd outputNoise = Eigen::VectorXd::Zero(outputNoiseDim);
        if (outFuncJacobianX) {
            return this->outFuncJacobianX(state, outputNoise);
        } else {
            auto func = [&](const Eigen::VectorXd& s) { return outputFunc(s, outputNoise); };
            return util::numericalJacobian(func, state, outputDim);
        }
    }

    Eigen::MatrixXd outputFuncJacobianV(const Eigen::VectorXd& state) {
        CHECK_EQ_F(state.size(), stateDim);
        if (outFuncJacobianV) {
            Eigen::VectorXd outputNoise = Eigen::VectorXd::Zero(outputNoiseDim);
            return this->outFuncJacobianV(state, outputNoise);
        } else {
            auto func = [&](const Eigen::VectorXd& outputNoise) {
                return outputFunc(state, outputNoise);
            };
            Eigen::VectorXd v = Eigen::VectorXd::Zero(outputNoiseDim);
            return util::numericalJacobian(func, v, outputDim);
        }
    }

    Eigen::VectorXd getOutput(const Eigen::VectorXd& state) {
        CHECK_EQ_F(state.size(), stateDim);
        Eigen::VectorXd outputNoise = Eigen::VectorXd::Zero(outputNoiseDim);
        return outputFunc(state, outputNoise);
    }

    void setOutputFuncJacobianX(
        const std::function<Eigen::MatrixXd(const Eigen::VectorXd&, const Eigen::VectorXd&)>&
            jacobianX) {
        outFuncJacobianX = jacobianX;
    }

    void setOutputFuncJacobianV(
        const std::function<Eigen::MatrixXd(const Eigen::VectorXd&, const Eigen::VectorXd&)>&
            jacobianV) {
        outFuncJacobianV = jacobianV;
    }
 
private:
    std::function<Eigen::MatrixXd(const Eigen::VectorXd&, const Eigen::VectorXd&)>
        outFuncJacobianX;
 
    std::function<Eigen::MatrixXd(const Eigen::VectorXd&, const Eigen::VectorXd&)>
        outFuncJacobianV;
};
 
template <int stateDim, int inputDim, int processNoiseDim, int numOutputs>
class MultiSensorEKF : public KalmanFilterBase<stateDim, inputDim> {
public:
    using state_t = navtypes::Vectord<stateDim>;
    using input_t = navtypes::Vectord<inputDim>;
    using processnoise_t = navtypes::Vectord<processNoiseDim>;
 
    using statefunc_t =
        std::function<state_t(const state_t&, const input_t&, const processnoise_t&)>;
 
    MultiSensorEKF(
        const statefunc_t& stateFunc,
        const statespace::NoiseCovMat<stateDim, processNoiseDim, inputDim>& processNoise,
        double dt, const std::array<Output, numOutputs>& outputs)
        : stateFunc(stateFunc), Q(processNoise), dt(dt), outputs(outputs) {}
 
    void predict(const input_t& input) override {
        navtypes::Matrixd<stateDim, stateDim> F = stateFuncJacobianX(this->xHat, input);
        navtypes::Matrixd<processNoiseDim, processNoiseDim> processNoise =
            Q.get(this->xHat, input);
 
        navtypes::Matrixd<stateDim, processNoiseDim> L = stateFuncJacobianW(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(int outputIdx, const Eigen::VectorXd& measurement) {
        Output& output = outputs[outputIdx];
 
        Eigen::MatrixXd H = output.outputFuncJacobianX(this->xHat);
        Eigen::MatrixXd outputNoise = output.covMat.get(this->xHat, measurement);
        Eigen::MatrixXd M = output.outputFuncJacobianV(this->xHat);
        Eigen::MatrixXd S = H * this->P * H.transpose() +
                            M * outputNoise * M.transpose(); // residual covariance
        // near-optimal kalman gain
        Eigen::Matrix<double, stateDim, Eigen::Dynamic> K =
            S.transpose().colPivHouseholderQr().solve(H * this->P.transpose()).transpose();
        Eigen::VectorXd y =
            measurement -
            output.outputFunc(this->xHat, Eigen::VectorXd::Zero(
                                              output.outputNoiseDim)); // measurement residual
 
        this->xHat = this->xHat + K * y;
        this->P = (Eigen::Matrix<double, stateDim, stateDim>::Identity() - K * H) * this->P;
    }
 
    template <int outputIdx>
    void correct(const Eigen::VectorXd& measurement) {
        static_assert(0 <= outputIdx && outputIdx < numOutputs);
        correct(outputIdx, measurement);
    }
 
    void setStateFuncJacobianX(
        const std::function<navtypes::Matrixd<stateDim, stateDim>(
            const state_t&, const input_t&, const processnoise_t&)>& stateFuncJacobianX) {
        this->stateFuncJacobianXSol = stateFuncJacobianX;
    }
 
    void setStateFuncJacobianW(
        const std::function<navtypes::Matrixd<stateDim, processNoiseDim>(
            const state_t&, const input_t&, const processnoise_t&)>& stateFuncJacobianW) {
        this->stateFuncJacobianWSol = stateFuncJacobianW;
    }
 
    navtypes::Matrixd<stateDim, stateDim> stateFuncJacobianX(const state_t& x,
                                                             const input_t& u) {
        processnoise_t w = processnoise_t::Zero();
        if (stateFuncJacobianXSol) {
            return stateFuncJacobianXSol(x, u, w);
        } else {
            auto func = [&](const state_t& state) { return stateFunc(state, u, w); };
            return util::numericalJacobian(func, x, stateDim);
        }
    }
 
    navtypes::Matrixd<stateDim, processNoiseDim> stateFuncJacobianW(const state_t& x,
                                                                    const input_t& u) {
        processnoise_t w = processnoise_t::Zero();
        if (stateFuncJacobianWSol) {
            return stateFuncJacobianWSol(x, u, w);
        } else {
            auto func = [&](const processnoise_t& noise) { return stateFunc(x, u, noise); };
            return util::numericalJacobian(func, w, stateDim);
        }
    }
 
private:
    statefunc_t stateFunc;
    statespace::NoiseCovMat<stateDim, processNoiseDim, inputDim> Q;
    double dt;
    std::array<Output, numOutputs> outputs;
 
    std::function<navtypes::Matrixd<stateDim, stateDim>(const state_t&, const input_t&,
                                                        const processnoise_t&)>
        stateFuncJacobianXSol;
 
    std::function<navtypes::Matrixd<stateDim, processNoiseDim>(const state_t&, const input_t&,
                                                               const processnoise_t&)>
        stateFuncJacobianWSol;
};
 
} // namespace filters