kalman_filter.h

#pragma once
 
#include "core/math/eigen_interface.h"
 
#include "core/math/systems/linear_system.h"

template <int STATES, int INPUTS, int OUTPUTS> class KalmanFilter {
  public:
    using StateVector = EVec<STATES>;
    using InputVector = EVec<INPUTS>;
    using OutputVector = EVec<OUTPUTS>;
 
    using StateMatrix = EMat<STATES, STATES>;
    using InputMatrix = EMat<STATES, INPUTS>;

    KalmanFilter(
      LinearSystem<STATES, INPUTS, OUTPUTS> &plant, const StateVector &state_stddevs,
      const OutputVector &measurement_stddevs
    )
        : KalmanFilter(plant.A(), plant.B(), plant.C(), plant.D(), state_stddevs, measurement_stddevs) {}

    KalmanFilter(
      const StateMatrix &A, const InputMatrix &B, const EMat<OUTPUTS, STATES> &C, const EMat<OUTPUTS, INPUTS> &D,
      const StateVector &state_stddevs, const OutputVector &measurement_stddevs
    ) {
        A_ = A;
        B_ = B;
        C_ = C;
        D_ = D;
 
        Q_ = state_stddevs.asDiagonal();
        Q_ = Q_ * Q_.transpose();
        R_ = measurement_stddevs.asDiagonal();
        R_ = R_ * R_.transpose();
 
        reset();
    }

    StateMatrix P() const { return P_; }

    void set_P(const StateMatrix &P) { P_ = P; }

    const StateVector &xhat() const { return xhat_; }

    double xhat(int i) const { return xhat_(i); }

    void set_xhat(const StateVector &xhat) { xhat_ = xhat; }

    void set_xhat(int i, double value) { xhat_(i) = value; }

    void reset() {
        xhat_.setZero();
        P_.setZero();
    }

    void predict(const InputVector &u, const double &dt) {
        // Q is discrete sqrt(process noise)
        EMat<STATES, STATES> Q = Q_ * dt;
        auto [A, B] = discretize_AB(A_, B_, dt);
 
        // Compute prior mean
        xhat_ = A * xhat_ + B * u;
 
        // Compute prior covariance
        P_ = A * P_ * A.transpose() + Q;
    }

    void correct(const OutputVector &y, const InputVector &u) { correct<OUTPUTS>(y, u, C_, D_, R_); }

    void correct(const OutputVector &y, const InputVector &u, const EMat<OUTPUTS, OUTPUTS> &R) {
        correct<OUTPUTS>(y, u, C_, D_, R);
    }

    template <int ROWS>
    void correct(
      const EVec<ROWS> &y, const InputVector &u, const EMat<ROWS, STATES> &C, const EMat<ROWS, INPUTS> &D,
      const EMat<ROWS, ROWS> &R
    ) {
        /*
         * Compute the innovation covariance
         *
         *   Py = CPCᵀ + R
         *
         */
        EMat<ROWS, ROWS> Py = C * P_ * C.transpose() + R;
        /* 
         * Compute the optimal Kalamn gain
         *
         *   K = (Py \ CPᵀ)ᵀ
         *
         */
        EMat<STATES, ROWS> K = Py.transpose().ldlt().solve(C * P_.transpose()).transpose();
        /* 
         * Compute the posterior mean
         *
         *   x̂ = x̂ + K(y - (Cx̂ + Du))
         *
         */
        xhat_ += K * (y - (C * xhat_ + D * u));
        /* 
         * Compute the posterior covariance using the Joseph form update equation
         *
         *   P = (I - KC)P(I - KC)ᵀ + KRKᵀ
         *
         */
        P_ = (EMat<STATES, STATES>::Identity() - K * C) * P_ * (EMat<STATES, STATES>::Identity() - K * C).transpose() +
             K * R * K.transpose();
    }
 
  private:
    StateVector xhat_;
    StateMatrix P_;
 
    StateMatrix Q_;
    EMat<OUTPUTS, OUTPUTS> R_;
 
    StateMatrix A_;
    EMat<STATES, INPUTS> B_;
    EMat<OUTPUTS, STATES> C_;
    EMat<OUTPUTS, INPUTS> D_;
 
    double dt_;
};
 
// allow using both names
template <int STATES, int INPUTS, int OUTPUTS> using KF = KalmanFilter<STATES, INPUTS, OUTPUTS>;