% Code to test the Unscented Kalman filter for a linear state-space model.
% Though the UKF block is very powerful in it's ability to track signals
% generated by a nonlinear model, this test is intended to be a sanity
% check.
% 
% Author: Saurabh Prasad
% Last modified: November 2, 2005.
%
% Dependent files - 
%                 |- sigmapt.m - compute sigma points
%                 |- sigmawt.m - compute 'predicted' observation, y from
%                                the sigma point theory.
%                 |- sigmawt_process.m - compute 'predicted' state vectors
%                                from the sigma point theory

clear all;
N = 100;
L = 1; 
num_sigma = 2*L+1;
lambda = 3 - L;
sigma_n = .2605;  % -----> Define State and Process noise std. deviations, 
sigma_v = .2;      % -----> sigma_v and sigma_n respectively


% Defining the weights:
W_m = zeros(1,(2*L+1));
W_c = zeros(1,(2*L+1));
W_m(1)=lambda/(L+lambda);
W_c(1)=(lambda/(L+lambda));
for j = 2 : 2*L+1
    W_m(j)=1/(2*(L+lambda));
    W_c(j)=1/(2*(L+lambda));
end;

% Store the 'clean' series generated by the process model 

P = [0.03 0.06 0.70 0.05];
x(1) = 0.1;    z(1)= x(1) + normrnd(0,sigma_n^2);
x(2) = 0.2;    z(2)= x(2) + normrnd(0,sigma_n^2);
x(3) = 0.5;    z(3)= x(3) + normrnd(0,sigma_n^2);
x(4) = 0.8;    z(4)= x(4) + normrnd(0,sigma_n^2);
% The state space model is a 4'th order AR process with an AWGN (0 mean, 0.02 variance)
% driving noise. The measurement model is simple with variable noise
% variance.
for i = 5 : 100
    x(i)=P(1)*x(i-1)+P(2)*x(i-2)+P(3)*x(i-3)+P(4)*x(i-4) + normrnd(0,sigma_v^2);
    z(i)= x(i) + normrnd(0,sigma_n^2); 
end;
state = x;
y_corrupted = z;


% Augment with process and measurement noise
X_a = [y_corrupted(1) 0  0]'; 

P_0 = var(state);

P_a = [P_0, 0, 0; 0, sigma_v, 0; 0, 0, sigma_n];
mu_a = [0, 0, 0]';

state_estimate(1) = y_corrupted(1);

for data_index = 1 : 3
    y_filtered(data_index) = y_corrupted(data_index);
end;

for data_index = 4 : length(y_corrupted)    
    % Generate the sigma points about the current estimate of the state-vector 
    [sigma] = sigmapt(X_a,mu_a',P_a,L);   
    for j = 1 : num_sigma
        sigma_i = sigma(j,:);
        sigma_state = sigma_i(1);  % in a multi-dimensional case, replace (1) with length of state vector.
        state_sigma(j,:) = sigma_state ;%+ normrnd(0,sigma_v);
        clear sigma_state;
    end;
    
    % Compute the function Y = g(X) on the sigma points:
	% Let the sample nonlinearity be y = x;
	for i = 1 : L
        for j = 1 : 2*L+1
            Y(j,i) = sigma(j,i);
        end;
	end;
    
    [mu_x, sigma_x] = sigmawt_process(sigma,L);
    x_predicted = mu_x; 
    P_predicted = sigma_x; 
    
    [mu_y, sigma_y] = sigmawt(sigma,L);
    
    y_predicted = mu_y;
    
    % The following updates the innovations covariance matrix
    P_yy = 0.00005*eye(L,L);  %Initializing the diagonals to a very small non-zero value, to prevent an 
                              %ill-conditioned matrix.
	for i = 1 : 2*L+1
        P_yy = P_yy + W_c(i)*((Y(i,:)-y_predicted)')*((Y(i,:)-y_predicted));
	end;
    
    % The following updates the Pxy
    P_xy = zeros(L,L);
	for i = 1 : 2*L+1
        P_xy = P_xy + W_c(i)*((state_sigma(i,:)-x_predicted)')*((Y(i,:)-y_predicted));
	end;
        
    % Updating the Kalman Gain
    K = P_xy*inv(P_yy);

    % Updating the filtered state-estimate
    x_estimate = x_predicted + K*(y_corrupted(data_index)-y_predicted);
    state_estimate(data_index) = x_estimate;
    
    P_estimate = P_predicted - K*P_yy*K';
    y_filtered(data_index) = x_estimate;
    
    P_a = [P_estimate, 0, 0; 0, sigma_v, 0; 0, 0, sigma_n]; 
    X_a = [x_predicted, 0, 0]; 

end;


figure;
plot(state,'b+');
hold on;
plot(y_corrupted,'rh'); 
plot(y_filtered,'g'); 
legend('Clean Series','Corrupted Series', 'Filtered Series');
title('noise variance = 0.6480');

