kalman.h

Go to the documentation of this file.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
#ifndef KALMAN_H
#define KALMAN_H

#include <functional>
namespace sp
{
#define FCN_XUW                                                                \
  [=](arma::mat x, arma::mat u,                                                \
      arma::mat w) // Lambda function f(x,u,w) ([capture] by copy)
using fcn_t = std::function<double(arma::mat, arma::mat, arma::mat)>;
using fcn_v = std::vector<fcn_t>;
using fcn_m = std::vector<fcn_v>;

arma_inline arma::mat eval_fcn(const fcn_v f, const arma::mat &x,
                               const arma::mat &u, const arma::mat &w)
{
  arma::mat y((arma::uword)(f.size()), 1, arma::fill::zeros);
  for (arma::uword n = 0; n < y.n_rows; n++)
    y(n, 0) = f[n](x, u, w);
  return y;
}

arma_inline arma::mat eval_fcn(const fcn_v f, const arma::mat &x,
                               const arma::mat &u)
{
  arma::mat w0(0, 0, arma::fill::zeros);
  return eval_fcn(f, x, u, w0);
}

arma_inline arma::mat eval_fcn(const fcn_v f, const arma::mat &x)
{
  arma::mat w0(0, 0, arma::fill::zeros);
  arma::mat u0(w0);
  return eval_fcn(f, x, u0, w0);
}

arma_inline void lti2discr(const arma::mat &F, const arma::mat &W,
                           const arma::mat &Qc, const double dT, arma::mat &A,
                           arma::mat &Q)
{
  arma::uword M = F.n_rows;

  // Solve A
  A = arma::expmat(F * dT);

  // Solve Q by using matrix fraction decomposition
  arma::mat AB                          = arma::zeros(2 * M, M);
  arma::mat CD                          = arma::zeros(2 * M, 2 * M);
  arma::mat EF                          = arma::zeros(2 * M, M);
  EF.submat(M, 0, 2 * M - 1, M - 1)     = arma::eye(M, M);
  CD.submat(0, 0, M - 1, M - 1)         = F;
  CD.submat(M, M, 2 * M - 1, 2 * M - 1) = -F.t();
  CD.submat(0, M, M - 1, 2 * M - 1)     = W * Qc * W.t();

  AB = arma::expmat(CD * dT) * EF;

  Q = AB.rows(0, M - 1) * arma::inv(AB.rows(M, 2 * M - 1));
}


class KF
{
protected:
  arma::uword N;   
  arma::uword M;   
  arma::uword L;   
  bool lin_proc;   
  bool lin_meas;   
  arma::mat x;     
  arma::mat z_err; 
  arma::mat A;     
  arma::mat B;     
  arma::mat H;     
  arma::mat P;     
  arma::mat Q;     
  arma::mat R;     
  arma::mat K;     
  fcn_v f;         
  fcn_v h;         
public:
  KF(arma::uword _N, arma::uword _M, arma::uword _L)
  {
    N        = _N; // Nr of states
    M        = _M; // Nr of measurements/observations
    L        = _L; // Nr of inputs
    lin_proc = true;
    lin_meas = true;
    x.set_size(N, 1);
    x.zeros();
    z_err.set_size(M, 1);
    z_err.zeros();
    A.set_size(N, N);
    A.eye();
    B.set_size(N, L);
    B.zeros();
    H.set_size(M, N);
    H.zeros();
    P.set_size(N, N);
    P.eye();
    Q.set_size(N, N);
    Q.eye();
    R.set_size(M, M);
    R.eye();
    K.set_size(N, M);
    K.zeros();
  }

  ~KF() {}

  void clear(void)
  {
    K.zeros();
    P.eye();
    x.zeros();
  }

  //  Set/get functions
  void set_state_vec(const arma::mat &_x) { x = _x; } // Set state vector.[Nx1]
  void set_trans_mat(const arma::mat &_A)
  {
    A = _A;
  } // Set state transition matrix.[NxN]
  void set_control_mat(const arma::mat &_B)
  {
    B = _B;
  } // Set input matrix.[NxL]
  void set_meas_mat(const arma::mat &_H)
  {
    H = _H;
  } // Set measurement matrix.[MxN]
  void set_err_cov(const arma::mat &_P)
  {
    P = _P;
  } // Set error covariance matrix.[NxN]
  void set_proc_noise(const arma::mat &_Q)
  {
    Q = _Q;
  } // Set process noise cov. matrix.
  void set_meas_noise(const arma::mat &_R)
  {
    R = _R;
  } // Set meas noise cov. matrix.
  void set_kalman_gain(const arma::mat &_K)
  {
    K = _K;
  }                            // Set Kalman gain matrix.[NxM]
  void set_trans_fcn(fcn_v _f) // Set state transition functions
  {
    f        = _f;
    lin_proc = false;
  }
  void set_meas_fcn(fcn_v _h) // Set measurement functions
  {
    h        = _h;
    lin_meas = false;
  }

  arma::mat get_state_vec(void) { return x; }   // Get states [Nx1]
  arma::mat get_err(void) { return z_err; }     // Get pred error [Mx1]
  arma::mat get_kalman_gain(void) { return K; } // Get Kalman gain [NxM]
  arma::mat get_err_cov(void) { return P; }     // Get error cov matrix [NxN]

  void predict(const arma::mat u)
  {
    x = A * x + B * u;     // New state
    P = A * P * A.t() + Q; // New error covariance
  }

  void predict(void)
  {
    arma::mat u0(L, 1, arma::fill::zeros);
    predict(u0);
  }

  void update(const arma::mat z)
  {
    // Compute the Kalman gain
    K = P * H.t() * inv(H * P * H.t() + R);

    // Update estimate with measurement z
    z_err = z - H * x;
    x     = x + K * z_err;

    // Joseph’s form covariance update
    arma::mat Jf = arma::eye<arma::mat>(N, N) - K * H;
    P            = Jf * P * Jf.t() + K * R * K.t();
  }

  void rts_smooth(const arma::mat &Xf, const arma::cube &Pf, arma::mat &Xs,
                  arma::cube &Ps)
  {
    arma::uword Nf = Xf.n_cols;
    arma::mat X_pred(N, 1);
    arma::mat P_pred(N, N);
    arma::mat C(N, N);

    // Copy forward data
    Xs = Xf;
    Ps = Pf;

    // Backward filter
    for (arma::uword n = Nf - 2; n > 0; n--)
    {
      // Project state and error covariance
      X_pred = A * Xf.col(n);
      P_pred = A * Pf.slice(n) * A.t() + Q;

      // Update
      C = Pf.slice(n) * A.t() * inv(P_pred);
      Xs.col(n) += C * (Xs.col(n + 1) - X_pred);
      Ps.slice(n) += C * (Ps.slice(n + 1) - P_pred) * C.t();
    }
  }
}; // End class KF

class EKF : public KF
{
protected:
  fcn_m f_jac; 
  fcn_m h_jac; 
  double dx;   
public:
  EKF(arma::uword _N, arma::uword _M, arma::uword _L) : KF(_N, _M, _L)
  {
    dx = 1e-7; // Default diff step size
  }

  //  Set/get functions
  void set_diff_step(double _dx) { dx = _dx; } // Set diff step size
  void set_state_jac(fcn_m _f)
  {
    f_jac = _f;
  } // Set EKF state transition jacobian
  void set_meas_jac(fcn_m _h)
  {
    h_jac = _h;
  } // Set EKF measurement transition jacobian

  void jacobian_diff(arma::mat &_F, fcn_v _f, const arma::mat &_x)
  {
    arma::uword nC = _F.n_cols;
    arma::uword nR = static_cast<arma::uword>(_f.size());
    arma::mat z0(
        nC, 1, arma::fill::zeros); // Zero matrix, assume dim u and w <= states

    if (nR == 0 || nR != _F.n_rows)
      err_handler("Function list is empty or wrong size");

    for (arma::uword c = 0; c < nC; c++)
    {
      arma::mat xp(_x);
      arma::mat xm(_x);
      xp(c, 0) += dx;
      xm(c, 0) -= dx;

      // Finite diff approx, evaluate at x,u=0,w=0
      for (arma::uword r = 0; r < nR; r++)
        _F(r, c) = (_f[r](xp, z0, z0) - _f[r](xm, z0, z0)) / (2 * dx);
    }
  }
  void jacobian_diff(arma::mat &_F, fcn_v _f)
  {
    jacobian_diff(_F, _f, x); // Use current state from object
  }

  void jacobian_analytical(arma::mat &_F, fcn_m _f_m, const arma::mat &_x)
  {
    arma::uword nC = _F.n_cols;
    arma::uword nR = static_cast<arma::uword>(_f_m.size());

    if (nR == 0 || nR != _F.n_rows)
      err_handler("Function list is empty or wrong size");

    arma::mat z0(
        nC, 1, arma::fill::zeros); // Zero matrix, assume dim u and w <= states
    for (arma::uword c = 0; c < nC; c++)
    {
      for (arma::uword r = 0; r < nR; r++)
        _F(r, c) = _f_m[r][c](_x, z0, z0);
    }
  }
  void jacobian_analytical(arma::mat &_F, fcn_m _f_m)
  {
    jacobian_analytical(_F, _f_m, x); // Use current state from object
  }

  void predict(const arma::mat u)
  {
    if (!lin_proc)
    {
      // Update A with jacobian approx or analytical if set
      if (f_jac.size() > 0)
        jacobian_analytical(A, f_jac);
      else
        jacobian_diff(A, f);

      // Predict state   x+ = f(x,u,0)
      x = eval_fcn(f, x, u);
    }
    else // Linear process
      x = A * x + B * u;

    // Project error covariance
    P = A * P * A.t() + Q;
  }

  void predict(void)
  {
    arma::mat u0(L, 1, arma::fill::zeros);
    predict(u0);
  }

  void update(const arma::mat z)
  {
    arma::mat z_hat(M, 1);

    if (!lin_meas) // Nonlinear measurement
    {
      // Update H with jacobian approx or analytical if set
      if (h_jac.size() > 0)
        jacobian_analytical(H, h_jac);
      else
        jacobian_diff(H, h);

      // Update measurement
      z_hat = eval_fcn(h, x);
    }
    else // Linear meas
      z_hat = H * x;

    // Calc residual
    z_err = z - z_hat;

    // Compute the Kalman gain
    K = P * H.t() * inv(H * P * H.t() + R);

    // Update estimate with measurement residual
    x = x + K * z_err;

    // Joseph’s form covariance update
    arma::mat Jf = arma::eye<arma::mat>(N, N) - K * H;
    P            = Jf * P * Jf.t() + K * R * K.t();
  }

  void rts_smooth(const arma::mat &Xf, const arma::cube &Pf, arma::mat &Xs,
                  arma::cube &Ps)
  {
    arma::uword Nf = Xf.n_cols;
    arma::mat X_pred(N, 1);
    arma::mat P_pred(N, N);
    arma::mat C(N, N);

    // Copy forward data
    Xs = Xf;
    Ps = Pf;

    // Backward filter
    for (arma::uword n = Nf - 2; n > 0; n--)
    {
      if (!lin_proc)
      {
        // Update A with jacobian approx or analytical if set
        if (f_jac.size() > 0)
          jacobian_analytical(A, f_jac, Xf.col(n));
        else
          jacobian_diff(A, f, Xf.col(n));

        // Project state
        X_pred = eval_fcn(f, Xf.col(n));
      }
      else // Linear process
      {
        // Project state
        X_pred = A * Xf.col(n);
      }

      // Project error covariance
      P_pred = A * Pf.slice(n) * A.t() + Q;

      // Update
      C = Pf.slice(n) * A.t() * inv(P_pred);
      Xs.col(n) += C * (Xs.col(n + 1) - X_pred);
      Ps.slice(n) += C * (Ps.slice(n + 1) - P_pred) * C.t();
    }
  }

}; // End class EKF

class UKF : public KF
{
protected:
  double alpha; 
  double beta;  
  double kappa; 
  double lambda;

  arma::mat X;  
  arma::mat S;  
  arma::mat C;  
  arma::vec Wx; 
  arma::vec Wp; 
public:
  UKF(arma::uword _N, arma::uword _M, arma::uword _L) : KF(_N, _M, _L)
  {
    alpha  = 1e-3;
    beta   = 2.0;
    kappa  = 0;
    lambda = alpha * alpha * (_N + kappa) - _N;
    X.set_size(_N, 2 * _N + 1);
    X.zeros();
    Wx.set_size(2 * _N + 1);
    Wx.zeros();
    Wp.set_size(2 * _N + 1);
    Wp.zeros();
  }

  //  Set/get functions
  void set_alpha(double _a) { alpha = _a; }   // Set alpha
  void set_beta(double _b) { beta = _b; }     // Set beta
  void set_kappa(double _k) { kappa = _k; }   // Set kappa
  void set_lambda(double _l) { lambda = _l; } // Set lambda

  void update_weights(void)
  {
    // Update lambda
    lambda = alpha * alpha * (N + kappa) - N;

    // Update weights
    Wx(0) = lambda / (N + lambda);
    Wp(0) = lambda / (N + lambda) + (1 - alpha * alpha + beta);

    for (arma::uword n = 1; n <= 2 * N; n++)
    {
      Wx(n) = 1 / (2 * (N + lambda));
      Wp(n) = Wx(n);
    }
  }

  void update_sigma(const arma::mat &_x, const arma::mat &_P)
  {
    // Update sigma points using Cholesky decomposition
    arma::mat _A = sqrt(N + lambda) * arma::chol(_P, "lower");

    X = arma::repmat(_x, 1, 2 * N + 1);
    X.cols(1, N) += _A;
    X.cols(N + 1, 2 * N) -= _A;
  }

  arma::mat ut(const arma::mat &_x, const arma::mat &_P, const fcn_v _f)
  {
    arma::uword Ny = static_cast<arma::uword>(_f.size());
    arma::mat y(Ny, 1);
    S.set_size(Ny, Ny);
    C.set_size(N, Ny);

    update_weights();
    update_sigma(_x, _P);

    // Propagate sigma points through nonlinear function
    arma::mat Xy(Ny, 2 * N + 1);
    for (arma::uword n = 0; n < 2 * N + 1; n++)
      Xy.col(n) = eval_fcn(_f, X.col(n));

    // New mean
    y = Xy * Wx;

    // New cov
    S = (Xy - arma::repmat(y, 1, 2 * N + 1)) * arma::diagmat(Wp) *
        (Xy - arma::repmat(y, 1, 2 * N + 1)).t();

    // New crosscov
    C = (X - arma::repmat(_x, 1, 2 * N + 1)) * arma::diagmat(Wp) *
        (Xy - arma::repmat(y, 1, 2 * N + 1)).t();

    return y;
  }

  void predict(const arma::mat u)
  {
    if (!lin_proc) // Nonlinear process
    {
      // Do the Unscented Transform
      x = ut(x, P, f) + B * u;

      // Add process noise cov
      P = S + Q;
    }
    else // Linear process
    {
      // Project state
      x = A * x + B * u;

      // Project error covariance
      P = A * P * A.t() + Q;
    }
  }

  void predict(void)
  {
    arma::mat u0(L, 1, arma::fill::zeros);
    predict(u0);
  }

  void update(const arma::mat z)
  {
    arma::mat z_hat(M, 1);
    if (!lin_meas) // Nonlinear measurement
    {
      // Do the Unscented Transform
      z_hat = ut(x, P, h);

      // Add measurement noise cov
      S = S + R;

      // Compute the Kalman gain
      K = C * arma::inv(S);

      // Update estimate with measurement residual
      z_err = z - z_hat;
      x     = x + K * z_err;

      // Update covariance, TODO: improve numerical perf. Josephs form?
      P = P - K * S * K.t();
    }
    else // Linear measurement
    {
      // Calc residual
      z_err = z - H * x;

      // Compute the Kalman gain
      K = P * H.t() * inv(H * P * H.t() + R);

      // Update estimate with measurement residual
      x = x + K * z_err;

      // Joseph’s form covariance update
      arma::mat Jf = arma::eye<arma::mat>(N, N) - K * H;
      P            = Jf * P * Jf.t() + K * R * K.t();
    }
  }

  void rts_smooth(const arma::mat &Xf, const arma::cube &Pf, arma::mat &Xs,
                  arma::cube &Ps)
  {
    arma::uword Nf = Xf.n_cols;
    arma::mat X_pred(N, 1, arma::fill::zeros);
    arma::mat P_pred(N, N, arma::fill::zeros);
    arma::mat D_pred(N, N, arma::fill::zeros);

    // Copy forward data
    Xs = Xf;
    Ps = Pf;

    // Backward filter
    for (arma::uword k = Nf - 2; k > 0; k--)
    {
      if (!lin_proc)
      {
        // Do the unscented transform
        X_pred = ut(Xf.col(k), Pf.slice(k), f);
        P_pred = S + Q;

        // Update
        D_pred = C * inv(P_pred);
        Xs.col(k) += D_pred * (Xs.col(k + 1) - X_pred);
        Ps.slice(k) += D_pred * (Ps.slice(k + 1) - P_pred) * D_pred.t();
      }
      else // Linear process
      {
        // Project state
        X_pred = A * Xf.col(k);

        // Project error covariance
        P_pred = A * Pf.slice(k) * A.t() + Q;

        // Update
        D_pred = Pf.slice(k) * A.t() * inv(P_pred);
        Xs.col(k) += D_pred * (Xs.col(k + 1) - X_pred);
        Ps.slice(k) += D_pred * (Ps.slice(k + 1) - P_pred) * D_pred.t();
      }
    }
  }
}; // End class UKF
} // namespace sp
#endif // KALMAN_H