kalman.py

import numpy as np


class KalmanFilter(object):
    def __init__(self,F,Q,H,R,G,u):
        """
        Initialize the dynamical system models
        
        Parameters
        ----------
        F : ndarray of shape (n,n)
            The state transition model
        Q : ndarray of shape (n,n)
            The covariance matrix for the state noise
        H : ndarray of shape (m,n)
            The observation model
        R : ndarray of shape (m,m)
            The covariance matrix for observation noise
        G : ndarray of shape (m,n)
            The control model
        u : ndarray of shape (n,)
            The control vector
        """
        # set attributes
        self.F = F
        self.Q = Q
        self.H = H
        self.R = R
        self.G = G
        self.u = u
    
    
    def evolve(self,x0,N):
        """
        Compute the first N states and observations generated by the Kalman system

        Parameters
        ----------
        x0 : ndarray of shape (n,)
            The initial state
        N : integer
            The number of time steps to evolve

        Returns
        -------
        states : ndarray of shape (n,N)
            The i-th column gives the i-th state
        obs : ndarray of shape (m,N)
            The i-th column gives the i-th observation
        """
        # initialize the sizes
        n = x0.shape[0]
        m = (self.H @ x0).shape[0]

        # initialize the output
        states = np.zeros((n,N))
        obs = np.zeros((m,N))
        states[:,0] = x0
        obs[:,0] = x0[:2]

        # iterate to compute the states and observations
        for i in range(1,N):
            states[:,i] = self.F @ states[:,i-1] + self.G @ self.u + np.random.multivariate_normal(np.zeros(n),self.Q)
            obs[:,i] = self.H @ states[:,i] + np.random.multivariate_normal(np.zeros(m),self.R)

        return states, obs


    def estimate(self,x0,P0,z, return_norms = False):
        """
        Compute the state estimates using the kalman filter

        Parameters
        ----------
        x0 : ndarray of shape (n,)
            The initial state estimate
        P0 : ndarray of shape (n,n)
            The initial error covariance matrix
        z : ndarray of shape(m,N)
            Sequence of N observations (each column is an observation)

        Returns
        -------
        out : ndarray of shape (n,N)
            Sequence of state estimates (each column is an estimate)
        """
        n = x0.shape[0]
        N = z.shape[1]

        # initialize the output
        output = np.zeros((n,N))
        output[:,0] = x0
        xk1 = x0
        pk1 = P0

        # iterate to compute the state estimates
        for i in range(1,N):
            # prediction step
            xk = self.F @ xk1 + self.G @ self.u
            pk = self.F @ pk1 @ self.F.T + self.Q

            # update step
            yh = z[:,i] - self.H @ xk
            testing = self.H @ pk @ self.H.T
            # print(testing.shape)

            Sk = self.H @ pk @ self.H.T + self.R
            Kk = pk @ self.H.T @ np.linalg.inv(Sk)
            xk = xk + Kk @ yh
            pk = (np.eye(n) - Kk @ self.H) @ pk

            # save the estimate
            output[:,i] = xk

            # update the state
            xk1 = xk
            pk1 = pk

        return output
    
    
    def predict(self,x,k):
        """
        Predict the next k states in the absence of observations

        Parameters
        ----------
        x : ndarray of shape (n,)
            The current state estimate
        k : integer
            The number of states to predict

        Returns
        -------
        out : ndarray of shape (n,k)
            The next k predicted states
        """
        out = np.empty((x.shape[0],k))

        # iterate to compute the state estimates
        for i in range(k):
            out_i = self.F @ x + self.G @ self.u
            out[:,i] = out_i
            x = out_i
        
        return out
    

    def rewind(self,x,k):
        """
        Predict the states from time 0 through k-1 in the absence of observations
    
        Parameters
        ----------
        x : ndarray of shape (n,)
            The state estimate at time k
        k : integer
            The current time step
    
        Returns
        -------
        out : ndarray of shape (n,k)
            The predicted states from time 0 up through k-1 (in that order)
        """
        out = np.empty((x.shape[0],k))

        # iterate to compute the state estimates
        for i in range(k):
            out_i = np.linalg.inv(self.F) @ (x - self.G @ self.u)
            out[:,i] = out_i
            x = out_i

        return out