Kalman filter

# coding: utf-8
import sys
import matplotlib.pyplot
import numpy as np


# maybe rewrite this using iterators?
# also make it more general so a, b, psi, phi and u can depend on t
def kalman(data, qmat, rmat, a, b, psi, phi, mu, sigma0, u):
	# Initialization
	# number of time steps
	n = data.shape[0]+1
	p = mu.shape[0]
	q = psi.shape[0]
	r = u.shape[0]

	ident = np.eye(p)
	# forecasting output X_{t|t-1}
	xtt1 = np.empty((n, p))
	# filtering output X_{t|t}
	xtt = np.empty((n, p))
	xtt[0] = mu
	# forecasting autocov Σ_{t|t-1}
	sigmatt1 = np.empty((n, p, p))
	# filtering autocov Σ_{t|t}
	sigmatt = np.empty((n, p, p))
	sigmatt[0] = sigma0
	# Kalman gain
	k = np.empty((n, p, q))


	for t, yt in enumerate(data, start=1):
		# (4.24)
		xtt1[t] = phi.dot(xtt[t-1]) + a.dot(u)
		# (4.25)
		sigmatt1[t] = phi.dot(sigmatt[t]).dot(phi.T) + qmat
		# (4.26)
		aux_k = np.linalg.inv(psi.dot(sigmatt1[t]).dot(psi.T) + rmat)
		k[t] = sigmatt1[t].dot(psi.T).dot(aux_k)
		# (4.27)
		xtt[t] = xtt1[t] + k[t].dot(yt - psi.dot(xtt1[t]) - b.dot(u))
		# (4.28)
		sigmatt[t] = (ident - k[t].dot(psi)).dot(sigmatt1[t])
	
	return xtt



def main(datafile='../traj1.dat'):
	# data = np.loadtxt(datafile)
	data = np.array(
		[[ 12.799954,   5.0418512],
		[ 16.642359 ,  12.97475  ],
		[ 21.99176  ,  10.857849 ],
		[ 22.526264 ,  13.018157 ],
		[ 29.905157 ,  10.709719 ],
		[ 15.497875 ,  16.449986 ],
		[ 25.494703 ,  16.770262 ],
		[ 27.290452 ,  16.614255 ],
		[ 30.660548 ,  22.402779 ],
		[ 35.892427 ,  21.89261  ],
		[ 35.685455 ,  28.388866 ],
		[ 50.224023 ,  34.383331 ],
		[ 32.598139 ,  19.698292 ],
		[ 39.528487 ,  39.720937 ],
		[ 45.911273 ,  27.534317 ],
		[ 52.917502 ,  25.414518 ],
		[ 43.192391 ,  38.625689 ],
		[ 53.423811 ,  34.754563 ],
		[ 55.486396 ,  36.394745 ],
		[ 55.381929 ,  35.208579 ],
		[ 54.143932 ,  39.957321 ],
		[ 56.759931 ,  33.771814 ],
		[ 58.099124 ,  48.334456 ],
		[ 73.405186 ,  51.652688 ],
		[ 69.849073 ,  53.492712 ],
		[ 64.707123 ,  43.462027 ],
		[ 66.521287 ,  48.476394 ],
		[ 67.52618  ,  48.47316  ],
		[ 74.662846 ,  59.827654 ],
		[ 73.282193 ,  51.127846 ],
		[ 76.882351 ,  72.867909 ],
		[ 72.159922 ,  60.10723  ],
		[ 75.866495 ,  79.537762 ],
		[ 74.778813 ,  61.094879 ],
		[ 69.659365 ,  69.377823 ],
		[ 74.86766  ,  72.016416 ],
		[ 78.654064 ,  65.544144 ],
		[ 80.93329  ,  84.874477 ],
		[ 80.588094 ,  78.453651 ],
		[ 67.524694 ,  78.0018   ],
		[ 69.736706 ,  79.454077 ],
		[ 65.265348 ,  79.820764 ],
		[ 85.064374 ,  78.722904 ],
		[ 71.914052 ,  87.228363 ],
		[ 59.301647 ,  81.81093  ],
		[ 61.086335 ,  89.235195 ],
		[ 49.332856 ,  89.826734 ],
		[ 54.41791  ,  86.06546  ],
		[ 58.720166 ,  87.580531 ],
		[ 53.928438 ,  76.63464  ],
		[ 44.532845 ,  88.746796 ],
		[ 49.264044 ,  81.016947 ],
		[ 38.986997 ,  85.6208   ],
		[ 43.342518 ,  79.356393 ],
		[ 37.901723 ,  82.579754 ],
		[ 33.683886 ,  84.059568 ],
		[ 30.769241 ,  69.712877 ],
		[ 22.344014 ,  85.858518 ],
		[ 26.164395 ,  83.40624  ],
		[ 18.367652 ,  78.701893 ],
		[ 26.458433 ,  88.694971 ],
		[ 21.879814 ,  77.973257 ],
		[ 19.868844 ,  79.467536 ],
		[ 33.68683  ,  79.87646  ],
		[  9.9335449,  75.115428 ],
		[ 27.365321 ,  69.542269 ],
		[ 20.600756 ,  64.241611 ],
		[ 21.448112 ,  67.245199 ],
		[ 21.297358 ,  59.802801 ],
		[ 26.373251 ,  58.16045  ],
		[ 25.436362 ,  60.717019 ],
		[ 27.018141 ,  57.929736 ],
		[ 27.519767 ,  50.658256 ],
		[ 30.755161 ,  55.440872 ],
		[ 31.411575 ,  54.321715 ],
		[ 24.873565 ,  50.999102 ],
		[ 32.652864 ,  52.781879 ],
		[ 39.474335 ,  53.329268 ],
		[ 40.300628 ,  42.336927 ],
		[ 41.269163 ,  45.769114 ],
		[ 46.884142 ,  42.17084  ],
		[ 43.56104  ,  45.405261 ],
		[ 48.872118 ,  52.393591 ],
		[ 54.242416 ,  48.916691 ],
		[ 54.20074  ,  42.48142  ],
		[ 61.57182  ,  47.90673  ],
		[ 58.987328 ,  43.930994 ],
		[ 57.377832 ,  44.616558 ],
		[ 74.083425 ,  35.857887 ],
		[ 70.295675 ,  35.965528 ],
		[ 72.164301 ,  32.619386 ],
		[ 81.566121 ,  40.65701  ],
		[ 76.52148  ,  31.05232  ],
		[ 76.187694 ,  43.563853 ],
		[ 82.235187 ,  30.180752 ],
		[ 74.555223 ,  33.725647 ],
		[ 84.552837 ,  25.001934 ],
		[ 92.985648 ,  39.762553 ],
		[ 90.321526 ,  14.971742 ],
		[ 86.996321 ,  28.122595 ]])

	# https://www.cl.cam.ac.uk/~rmf25/papers/Understanding%20the%20Basis%20of%20the%20Kalman%20Filter.pdf
	
	# sample period
	delta = 1
	# state space dimension
	p = 4
	# observation space dimension
	q = 2
	# "exogenous input" dimension
	r = 3
	# Variance of the white noise A
	vara = 1
	# Variance of the white noise V
	varv = 1


	# Φ_t in the poly is Φ from the pdf
	# X_k = Φ X_{k-1} + Π A_{k-1}
	phi = np.eye(p)
	phi[0, 2] = delta
	phi[1, 3] = delta


	# Ψ_t in the poly
	psi = np.zeros((q, p))
	psi[0, 0] = 1
	psi[1, 1] = 1

	# Q = cov(W_k) = cov(Π A_{k-1})
	qmat = np.zeros((p, p))
	qmat[0, 0] = vara*(delta**4)/4.0
	qmat[1, 1] = vara*(delta**4)/4.0
	qmat[0, 2] = vara*(delta**3)/2.0
	qmat[1, 3] = vara*(delta**3)/2.0
	qmat[2, 0] = vara*(delta**3)/2.0
	qmat[3, 1] = vara*(delta**3)/2.0
	qmat[2, 3] = vara*(delta**2)
	qmat[3, 3] = vara*(delta**2)

	rmat = np.eye(q)*(varv*varv)

	# As per question 5
	sigma0 = np.eye(p)*100

	# As per question 5
	mu = np.zeros((p,))

	# we don't have any of that
	a = np.zeros((p, r))
	b = np.zeros((q, r))
	u = np.zeros((r,))
	# TODO: Plot this instead of printing
	# [:,:2] – print only the position estimates, not the velocity
	print kalman(data, qmat, rmat, a, b, psi, phi, mu, sigma0, u)[:,:2]
	

if __name__ == '__main__':
	# if len(sys.argv) != 2:
	# 	print("Pass data file as parameter")
	# 	sys.exit(2)
	main()