Extended Kalman filter sample

ekf.py

# -*- coding: utf-8 -*-
 
import numpy as np
import matplotlib.pyplot as plt
 
def main():
    # 初期化
    T = 30 # 観測数
    r = 10.0 # 半径
    w = 1.0*10/180 * np.pi # 角速度[rad/s]

    x = np.mat([[0],[-5]]) # 初期位置
    X = [x] # 実際の状態
    Y = [x] # 観測
    U = np.mat([[r],[w]]) # 操作量(一定)
 
    # state x = f(x_,u,v), v~N(0,Q)
    Q = np.mat([[0.5,0.0],[0.0,0.5]])

    # observation Y = x + w, w~N(0,R)
    R = np.mat([[0.5,0.0], [0.0,0.5]])
    
    def f(t,x,u):
        x0 = u[0,0]*u[1,0]*np.cos(u[1,0]*t)+x[0,0]
        x1 = u[0,0]*u[1,0]*np.sin(u[1,0]*t)+x[1,0]
        return np.mat([[x0],[x1]])

    def Jf(t,x,u):
        """
        解析的に求めるf(x)のヤコビ行列
        """
        return np.mat([[u[0,0]*u[1,0]*np.cos(u[1,0]*t),0],[0,u[0,0]*u[1,0]*np.sin(u[1,0]*t)]])

    # 観測データの生成
    for t in range(T):
        x = f(t,x,U)+np.random.multivariate_normal([0,0],Q,1).T
        X.append(x)
        y = x + np.random.multivariate_normal([0,0],R,1).T
        Y.append(y)
     
    # EKF
    _x = np.mat([[0],[-5]])
    Sigma = np.mat([[1,0],[0,1]])    
    _X = [_x] # 推定
    for t in range(T):
        # prediction
        A = Jf(t,_x,U)
        _x_ = f(t,_x,U)
        Sigma_ = Q + A * Sigma * A.T
        # update
        yi = Y[t+1] - _x_
        S = Sigma_ + R
        G = Sigma_ * S.I
        _x = _x_ + G * yi
        Sigma = Sigma_ - G * Sigma_
        _X.append(_x)

    # 描画
    a, b = np.array(np.concatenate(X,axis=1))
    plt.plot(a,b,'rs-', label="X_correct")
    a, b = np.array(np.concatenate(Y,axis=1))
    plt.plot(a,b,'g^-', label="Y (=X+N(0,R))")
    a, b = np.array(np.concatenate(_X,axis=1))
    plt.plot(a,b,'bo-', label="X_estimate")
    plt.legend(bbox_to_anchor=(0.80, 0.00), loc='lower left', borderaxespad=0)
    plt.axis('equal')
    plt.show()
 
if __name__ == '__main__':
    main()

Sample of extended Kalman filter (EKF)