Kalman filtering

In learning the principle of Kalman filtering, I hope to deepen the understanding of the formula and principle through python.

Take notes

import numpy as np
import math
import matplotlib.pyplot as plt

'''
    dynam_params：The dimensions of the state space.
    measure_params: the dimension of the measure value; 
    control_params: the dimension of the control vector, default is 0.
'''
class Kalman(object):
    '''
        INIT KALMAN
    '''
    def __init__(self, dynam_params, measure_params,control_params = 0,type = np.float32):
        self.dynam_params = dynam_params
        self.measure_params = measure_params
        self.control_params = control_params
        # self

        ### The following should all be able to determine the dimension based on the input dimension value
        if(control_params ! = 0):
            self.controlMatrix = np.mat(np.zeros((dynam_params, control_params)),type) # control matrix
        else:
            self.controlMatrix = None
        self.errorCovPost =  np.mat(np.zeros((dynam_params, dynam_params)),type) # P_K
        self.errorCovPre =  np.mat(np.zeros((dynam_params, dynam_params)),type) # P_k-1
        self.gain =  np.mat(np.zeros((dynam_params, measure_params)),type) # K

        self.measurementMatrix =  np.mat(np.zeros((measure_params, dynam_params)),type) 
        self.measurementNoiseCov =  np.mat(np.zeros((measure_params, measure_params)),type) 
        self.processNoiseCov =  np.mat(np.zeros((dynam_params, dynam_params)),type) 
        self.transitionMatrix =  np.mat(np.zeros((dynam_params, dynam_params)),type)

        self.statePost =  np.array(np.zeros((dynam_params, 1)),type)
        self.statePre =  np.array(np.zeros((dynam_params, 1)),type)
        
        # The diagonal is initialized to 1
        ## np.diag_indices returns the index of the main diagonal as a tuple
        ### F state transfer matrix diagonal initialized to 1
        row,col = np.diag_indices(self.transitionMatrix.shape[0])
        self.transitionMatrix[row,col] = np.array(np.ones(self.transitionMatrix.shape[0]))
        ### R measure noise Diagonal initialized to 1
        row,col = np.diag_indices(self.measurementNoiseCov.shape[0])
        self.measurementNoiseCov[row,col] = np.array(np.ones(self.measurementNoiseCov.shape[0]))
        ### Q Process noise Diagonal initialized to 1
        row,col = np.diag_indices(self.processNoiseCov.shape[0])
        self.processNoiseCov[row,col] = np.array(np.ones(self.processNoiseCov.shape[0]))
    
    def predict(self,control_vector = None):
        '''
            PREDICT
        '''
        # Predicted value
        F = self.transitionMatrix
        x_update = self.statePost
        B = self.controlMatrix

        if(self.control_params == 0):
            x_predict = F * x_update
        else:
            x_predict = F * x_update + B * control_vector
        self.statePre = x_predict

        # P_k
        P_k_minus = self.errorCovPost
        Q = self.processNoiseCov
        self.errorCovPre = F * P_k_minus * F.T + Q

        self.statePost = self.statePre
        self.errorCovPost = self.errorCovPre
        return x_predict

    def correct(self,mes):
        '''
            CORRECT
        '''
        # K update
        K = self.gain
        P_k = self.errorCovPost
        H = self.transitionMatrix
        R = self.measurementNoiseCov
        K = P_k * H.T * np.linalg.inv(H * P_k * H.T + R)
        self.gain = K

        # Calculate the estimated value of State
        x_predict = self.statePre
        x_update = x_predict +  K * (mes - H * x_predict)
        self.statePost = x_update

        # P_k update
        P_pre = self.errorCovPre
        P_k_post = P_pre - K * H * P_pre
        self.errorCovPost = P_k_post
        return x_update

if __name__ == '__main__':
    
    pos = np.array([
        [10,    50],
        [12,    49],    
        [11,    52],     
        [13,    52.2],     
        [12.9,  50]], np.float32)
    kalman = Kalman(2,2)
    kalman.measurementMatrix = np.mat([[1,0],[0,1]],np.float32)
    kalman.transitionMatrix = np.mat([[1,0],[0,1]], np.float32)
    kalman.processNoiseCov = np.mat([[1,0],[0,1]], np.float32) * 1e-4
    kalman.measurementNoiseCov = np.mat([[1,0],[0,1]], np.float32) * 1e-4 
    kalman.statePre =  np.mat([[6],[6]],np.float32)
   #kalman.statePre =  np.mat([[6],[6]],np.float32)
    for i in range(len(pos)):
        mes = np.reshape(pos[i,:],(2,1))
        y = kalman.predict()
        
        print("before correct mes",mes[0],mes[1])
        x = kalman.correct(mes)
        print (kalman.statePost[0],kalman.statePost[1])
        print (kalman.statePre[0],kalman.statePre[1])
        print ('measurement:\t',mes[0],mes[1])  
        print ('correct:\t',x[0],x[1])
        print ('predict:\t',y[0],y[1])     
        print ('='*30)