final version of kf with reference solution

Author: adlarkin

KalmanFilter/kf.py

@@ -7,6 +7,7 @@
 from numpy.linalg import inv as mat_inv # matrix inverse
 import matplotlib.pyplot as plt
 from scipy.io import loadmat
+import pdb
 
 class Plotter():
     def __init__(self, times):
@@ -35,6 +36,28 @@ def __init__(self, times):
         self.k_v = []
         self.k_x = []
 
+        # plotting the covariance in position after each measurement and prediction
+        self.x_cov_times = []
+        self.x_cov_vals = []
+
+    def save_iteration_data(self, model, timestep, k_gains=None):
+        """
+        @type model: Model
+        """
+        self.v_sigma_pos.append(np.sqrt(model.sigma[0 , 0]) * 2)
+        self.v_sigma_neg.append(np.sqrt(model.sigma[0 , 0]) * 2 * -1)
+        self.v_error.append(model.vtr[0 , timestep] - model.mu[0 , 0])
+        self.x_sigma_pos.append(np.sqrt(model.sigma[1 , 1]) * 2)
+        self.x_sigma_neg.append(np.sqrt(model.sigma[1 , 1]) * 2 * -1)
+        self.x_error.append(model.xtr[0 , timestep] - model.mu[1 , 0])
+        self.v_true.append(model.vtr[0 , timestep])
+        self.x_true.append(model.xtr[0 , timestep])
+        self.v_pred.append(model.mu[0 , 0])
+        self.x_pred.append(model.mu[1 , 0])
+        if k_gains is not None:
+            self.k_v.append(k_gains[0 , 0])
+            self.k_x.append(k_gains[1 , 0])
+
     def plot(self):
         # increase vertical spacing between subplots
         # https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.subplots_adjust.html
@@ -45,54 +68,51 @@ def plot(self):
         # error covariance plots
         p1 = plt.figure(1)
         plt.subplot(211)
-        plt.plot(times, self.v_sigma_neg, 'r', label="Error Covariance")
-        plt.plot(times, self.v_sigma_pos, 'r')
-        plt.plot(times, self.v_error, 'b', alpha=.5, label="Velocity Estimation Error")
+        plt.plot(self.times, self.v_sigma_neg, 'r', label="Error Covariance")
+        plt.plot(self.times, self.v_sigma_pos, 'r')
+        plt.plot(self.times, self.v_error, 'b', label="Velocity Estimation Error")
         plt.title("Estimation error and error covariance vs time")
         plt.ylabel("Error (velocity)")
         plt.legend()
         plt.subplot(212)
-        plt.plot(times, self.x_sigma_neg, 'r', label="Error Covariance")
-        plt.plot(times, self.x_sigma_pos, 'r')
-        plt.plot(times, self.x_error, 'b', alpha=.5, label="Position Estimation Error")
+        plt.plot(self.times, self.x_sigma_neg, 'r', label="Error Covariance")
+        plt.plot(self.times, self.x_sigma_pos, 'r')
+        plt.plot(self.times, self.x_error, 'b', label="Position Estimation Error")
         plt.ylabel("Error (position)")
         plt.xlabel(x_label_str)
         plt.legend()
         p1.show()
 
         # state estimates vs ground truth plots
-        true_opacity = .6
-        pred_opacity = .75
         p2 = plt.figure(2)
-        plt.subplot(211)
-        plt.plot(times, self.v_true, alpha=true_opacity, label="True Velocity")
-        plt.plot(times, self.v_pred, alpha=pred_opacity, label="Predicted Velocity")
+        plt.plot(self.times, self.v_true, label="True Vel")
+        plt.plot(self.times, self.v_pred, label="Predicted Vel")
+        plt.plot(self.times, self.x_true, label="True Pos")
+        plt.plot(self.times, self.x_pred, label="Predicted Pos")
         plt.title("State estimates and true states vs time")
-        plt.ylabel("Velocity")
-        plt.legend()
-        plt.subplot(212)
-        plt.plot(times, self.x_true, alpha=true_opacity, label="True Position")
-        plt.plot(times, self.x_pred, alpha=pred_opacity, label="Predicted Prosition")
-        plt.ylabel("Postion")
         plt.xlabel(x_label_str)
+        plt.ylabel("Position (m) and Velocity (m / s^2)")
         plt.legend()
         p2.show()
 
         # kalman gain plots
-        gain = "Gain"
+        # no kalman gain at the first timestep since filtering starts at second timestep
         p3 = plt.figure(3)
-        plt.subplot(211)
-        plt.plot(times, self.k_v, label="Velocity")
+        plt.plot(self.times[1:], self.k_v, label="Velocity")
+        plt.plot(self.times[1:], self.k_x, label="Position")
         plt.title("Kalman gain vs time")
-        plt.ylabel(gain)
-        plt.legend()
-        plt.subplot(212)
-        plt.plot(times, self.k_x, label="Position")
-        plt.ylabel(gain)
+        plt.ylabel("Gain")
         plt.xlabel(x_label_str)
         plt.legend()
         p3.show()
 
+        p4 = plt.figure(4)
+        plt.plot(self.x_cov_times, self.x_cov_vals)
+        plt.title("Position covariance between prediction and measurement update")
+        plt.ylabel("Covariance")
+        plt.xlabel(x_label_str)
+        p4.show()
+
         # keep the plots open until user enters Ctrl+D to terminal (EOF)
         try:
             input()
@@ -117,6 +137,8 @@ def __init__(self, times, A, B, C, D, R, Q,
         """
         np.random.seed(seed)    # for debugging (removes randomness from noise)
 
+        self.times = times
+
         # state space
         self.A = A
         self.B = B
@@ -155,18 +177,23 @@ def load_matlab_data(self):
         # try on the matlab data (for comparison)
         self.mu = mu0
         self.sigma = Sig0
+        self.Q = Q
+        self.R = R
         self.control_inputs = u
         self.vtr = vtr
         self.xtr = xtr
         self.measurements = z
 
     def get_data(self):
+        # initial condition is v=0 and x=0
+        curr_state = np.array([[0] , [0]])
+        # input and measurement at start at timestep 1 since KF starts at timestep 1
         vtr = np.zeros((1,self.control_inputs.size))
         xtr = np.zeros((1,self.control_inputs.size))
         measurements = np.zeros((1,self.control_inputs.size))
 
-        curr_state = np.array(self.mu)
-        for i in range(self.control_inputs.size):
+        # updates begin at timestep 1 since timestep 0 is the initial state
+        for i in range(1, self.control_inputs.size):
             c_input = self.control_inputs[0 , i]
             c_input = np.reshape(c_input, (1,1))
             # get the next state from the control input
@@ -197,17 +224,22 @@ def make_measurement_noise(self):
 
     def kalman_filter(self):
         # for plotting
-        self.plotter = Plotter(times)
+        self.plotter = Plotter(self.times)
+
+        self.plotter.save_iteration_data(self, 0, None)
+        self.plotter.x_cov_times.append(0)
+        self.plotter.x_cov_vals.append(self.sigma[1 , 1])
+        for timestep in range(1,self.control_inputs.size):
 
-        for timestep in range(self.control_inputs.size):
             c_input = self.control_inputs[0 , timestep]
             c_input = np.reshape(c_input, (1,1))
             z = self.measurements[0 , timestep]
 
             # prediction
             mu_bar = mm(self.A, self.mu) + mm(self.B, c_input)
             sigma_bar = mm(self.A, mm(self.sigma, np.transpose(self.A))) + self.R
-            
+            self.plotter.x_cov_times.append(self.times[timestep])
+            self.plotter.x_cov_vals.append(sigma_bar[1 , 1])
             # correction
             c_transpose = np.transpose(self.C)
             matrix_one = mm(sigma_bar, c_transpose)
@@ -220,85 +252,74 @@ def kalman_filter(self):
             self.mu = mu
             self.sigma = sigma
 
-            # save info for plotting
-            # error covariance and estimation error plots
-            self.plotter.v_sigma_pos.append(np.sqrt(sigma[0 , 0]) * 2)
-            self.plotter.v_sigma_neg.append(np.sqrt(sigma[0 , 0]) * 2 * -1)
-            self.plotter.v_error.append(self.vtr[0 , timestep] - mu[0 , 0])
-            self.plotter.x_sigma_pos.append(np.sqrt(sigma[1 , 1]) * 2)
-            self.plotter.x_sigma_neg.append(np.sqrt(sigma[1 , 1]) * 2 * -1)
-            self.plotter.x_error.append(self.xtr[0 , timestep] - mu[1 , 0])
-            # state estimation plots
-            self.plotter.v_pred.append(mu[0 , 0])
-            self.plotter.v_true.append(self.vtr[0 , timestep])
-            self.plotter.x_pred.append(mu[1 , 0])
-            self.plotter.x_true.append(self.xtr[0 , timestep])
-            # kalman gain plots
-            self.plotter.k_v.append(k[0 , 0])
-            self.plotter.k_x.append(k[1 , 0])
+            self.plotter.save_iteration_data(self, timestep, k)
+            self.plotter.x_cov_times.append(self.times[timestep])
+            self.plotter.x_cov_vals.append(sigma[1 , 1])
 
     def plot_results(self):
         self.plotter.plot()
 
 
-########################################################################################
-############################## DEFINE PARAMETERS HERE ##################################
-########################################################################################
-# use the data from the matlab file, or make our own?
-use_file_data = False
-# eliminate randomness in noise generation? (For debugging)
-rand_seed = None
-# measurement covariance (default = .001)
-z_noise = .001
-# process noise associated with velocity state (default = .01)
-v_noise = .01
-# process noise associated with position state (default = .0001)
-x_noise = .0001
-# timestep (seconds)
-dt = .05
-#  initial belief (initial condition):
-# starting v and x, [[v] , [x]]
-mu = np.array([[0] , [0]])
-# starting covariance, [[v , 0] , [0 , x]]
-sigma = np.array([[1 , 0] , [0 , 1]])
-########################################################################################
-########################################################################################
-
-
-# robot model parameters
-b = 20      # drag coefficient
-m = 100     # mass
-
-# define continuous state space model
-A = np.array([[(-b/m) , 0] , [1 , 0]])
-B = np.array([[(1/m)] , [0]])
-C = np.array([0, 1])
-D = np.array([0])
-dt = .05
-
-# get a discrete state space model
-sys = control.ss(A, B, C, D)
-sys_d = control.c2d(sys, dt)
-sys_d.A = np.array(sys_d.A)
-sys_d.B = np.array(sys_d.B)
-sys_d.C = np.array(sys_d.C)
-sys_d.D = np.array(sys_d.D)
-
-# get the system/measurement noise
-R = np.array([[v_noise , 0] , [0 , x_noise]])
-Q = np.array([[z_noise]])
-
-# set the ground truth control inputs
-times = np.arange(0, 50+dt, dt)
-control_inputs = np.zeros((1, times.size))
-control_inputs[0 , 0:100] = 50
-control_inputs[0 , 100:500] = 0
-control_inputs[0 , 500:600] = -50
-control_inputs[0 , 600:] = 0
-
-# make the robot model
-uuv = Model(times, sys_d.A, sys_d.B, sys_d.C, sys_d.D, R, Q, 
-    mu, sigma, control_inputs, rand_seed, use_file_data)
-
-uuv.kalman_filter()
-uuv.plot_results()
+
+if __name__ == "__main__":
+    ########################################################################################
+    ############################## DEFINE PARAMETERS HERE ##################################
+    ########################################################################################
+    # use the data from the matlab file, or make our own?
+    use_file_data = True
+    # eliminate randomness in noise generation? (For debugging)
+    rand_seed = None
+    # measurement covariance (default = .001)
+    z_noise = .001
+    # process noise associated with velocity state (default = .01)
+    v_noise = .01
+    # process noise associated with position state (default = .0001)
+    x_noise = .0001
+    # timestep (seconds)
+    dt = .05
+    #  initial belief (initial condition):
+    # starting v and x: [[v] , [x]]
+    mu = np.array([[0] , [0]])
+    # starting covariance: [[v , 0] , [0 , x]]
+    sigma = np.array([[.75 , 0] , [0 , .05]])
+    ########################################################################################
+    ########################################################################################
+
+
+    # robot model parameters
+    b = 20      # drag coefficient
+    m = 100     # mass
+
+    # define continuous state space model
+    A = np.array([[(-b/m) , 0] , [1 , 0]])
+    B = np.array([[(1/m)] , [0]])
+    C = np.array([0, 1])
+    D = np.array([0])
+    dt = .05
+
+    # get a discrete state space model
+    sys = control.ss(A, B, C, D)
+    sys_d = control.c2d(sys, dt)
+    sys_d.A = np.array(sys_d.A)
+    sys_d.B = np.array(sys_d.B)
+    sys_d.C = np.array(sys_d.C)
+    sys_d.D = np.array(sys_d.D)
+
+    # get the system/measurement noise
+    R = np.array([[v_noise , 0] , [0 , x_noise]])
+    Q = np.array([[z_noise]])
+
+    # set the ground truth control inputs
+    input_times = np.arange(0, 50+dt, dt)
+    control_inputs = np.zeros((1, input_times.size))
+    control_inputs[0 , 0:100] = 50
+    control_inputs[0 , 100:500] = 0
+    control_inputs[0 , 500:600] = -50
+    control_inputs[0 , 600:] = 0
+
+    # make the robot model
+    uuv = Model(input_times, sys_d.A, sys_d.B, sys_d.C, sys_d.D, R, Q, 
+        mu, sigma, control_inputs, rand_seed, use_file_data)
+
+    uuv.kalman_filter()
+    uuv.plot_results()