Milk and Butter price inference

kalman.py

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Ellipse
import math

# Kalman

def lkf_predict(x, P, A, B, u, Q):
    xpred = A.dot(x) + B.dot(u)
    Ppred = A.dot(P).dot(A.T) + Q
    return xpred, Ppred

def lkf_correct(x, P, z, H, R):
    y = z - H.dot(x)
    S = H.dot(P).dot(H.T) + R
    K = P.dot(H.T).dot(np.linalg.inv(S))
    xcorr = x + K.dot(y)
    Pcorr = (np.eye(K.shape[0]) - K.dot(H)).dot(P)
    return xcorr, Pcorr

# 95% confidence ellipse params

def conf_ellipse(x, P, chi_square=5.991):
    w, v = np.linalg.eig(P)
    w = np.abs(w)
    major = np.argmax(w)
    minor = (major + 1) % 2
    angle = math.atan2(v[1, major], v[0, major])
    if angle < 0.:
        angle += 2 * np.pi
    
    width = 2 * math.sqrt(w[major] * chi_square)
    height = 2 * math.sqrt(w[minor] * chi_square)

    return width, height, angle

# Purchase

def go_shopping(n, milk_price, butter_price, sigma_total):
    amounts = np.random.uniform(0., 1.5, size=(n, 2))
    # amounts = np.repeat([[1, 1]], n, axis=0)
    pricepaid = amounts[:, 0] * milk_price + amounts[:, 1] * butter_price    
    pricepaid += np.random.normal(0, sigma_total, size=amounts.shape[0])
    return np.hstack((amounts, pricepaid.reshape(-1,1)))

sigma_total = 1
purchases = go_shopping(100, milk_price=2.5, butter_price=3., sigma_total=sigma_total)

b = purchases[:, 2]
A = purchases[:, :2]
print(np.linalg.solve(A.T.dot(A), A.T.dot(b)))

# Kalman filter setup
x = np.array([[1.], [1.]])
P = np.eye(2) * 3.
A = np.array([
    [1, 0],
    [0, 1],
])
Q = np.eye(2) * 0.001
R = np.eye(1) * sigma_total**2
B = np.zeros((2,2))
u = np.zeros((2,1))

# Draw related

fig, ax = plt.subplots()
ax.set_title('Milk and Butter')
ax.set_xlim(0, 5)
ax.set_ylim(0, 5)
ax.set_xlabel('Milk')
ax.set_ylabel('Butter')

w, h, a = conf_ellipse(x, P)
ellipse = Ellipse(xy=(x[0,0],x[1,0]), width=w, height=h, angle=np.degrees(a), animated=True, alpha=0.8)
xy, = plt.plot([], [], marker='+', color='r', ls='', alpha=0.8)

# Animation related
def update(i):
    global ax, x, P, ellipse
   
    if i == 0:        
        ax.add_patch(ellipse)
        return ellipse,
    else:
        # Update state estimate
        x, P = lkf_predict(x, P, A, B, u, Q)
        
        H = np.array([purchases[i, 0], purchases[i, 1]]).reshape(1, 2)


        z = np.array([purchases[i, 2]]).reshape(1,1)
        x, P = lkf_correct(x, P, z, H, R)
        print(P)
        
        w, h, a = conf_ellipse(x, P)
        ellipse.angle = np.degrees(a)
        ellipse.width = w
        ellipse.height = h
        ellipse.center = (x[0,0], x[1,0])
        xy.set_xdata(x[0,0])
        xy.set_ydata(x[1,0])
        return ellipse, xy
    
anim = FuncAnimation(fig, update, interval=50, frames=purchases.shape[0], blit=True)
plt.grid()
plt.show()