from numarray import *
from MLab import *
from pylab import *

NUM = 1000
STD_DEV_X = 0.001
STD_DEV_Y = 0.1

"""
Kalman Filtering
Carlos Rocha
MIT Media Lab

System:
    x = f(x) + Nx = Ax + Nx
    y = h(x) + Ny

Problem:

    theta = 0.1*t + 4* sin(0.01*t)

    y = sin(theta)
    theta(n+1) = 2*theta(n) + theta(n-1)

Extended Kalman Filter:
    x = [theta(n), theta(n-1)]
    A = [2, -1; 1, 0]
    h = sin(x)
    B = dh/dx = cos(x)

"""
def main() :

    yn = map(y_func, range(0,NUM))
    ym = zeros((NUM,1),Float64)
    ye = zeros((NUM,1),Float64)

    #################
    # INTIALIZE
    I = identity(2, typecode=Float32)
    A = array([[2.0, -1.0],[1.0, 0.0]])
    E = I

    x = array ([0.1 , 0.0])

    Nx = pow(STD_DEV_X,2) * I
    Ny = pow(STD_DEV_Y,2);
    
    for i in range(0,NUM):

        # PREDICT
        x = dot(A,x);
        E = dot(A,dot(E,transpose(A))) + Nx

        # MEASURE
        y = yn[i] + randn() * STD_DEV_Y
        B = array([cos(x[0]), 0.0])

        # CORRECT
        K = dot(E,B) / (dot(B,dot(E,B)) + Ny)
        x = x + dot(K,(y - sin(x[0])))
        E = dot((I - outerproduct(K,B)),E)

        # PLOT
        ym[i] = y
        ye[i] = sin(x[0])

    plot(ym,'b+')
    plot(ye,'g')
    title('Kalman Filter')
    show()

def y_func(t):
    return sin(0.1*t + 4* sin(0.01*t))

if __name__ == "__main__" : main()