clear;set(0,'defaultaxesfontsize',14);format long
%%% Kalman Filter in 2D, with the first component observed
%%% based on Ex 2.3 in: Law, Stuart, Zygalakis (2015), "Data Assimilation"

%% setup

N=1e3;% number of steps
dim=2;% dimension of state
I=eye(dim);% identity operator
gamma=1;% observational noise variance is gamma^2*I
sigma=0.5;% model noise covariance is sigma^2*I
C0=1.*eye(dim);% prior initial condition variance
m0=[1;1];% prior initial condition mean
sd=10;rng(sd);% choose random number seed
A=1.*[0 1;-1 1];% dynamics determined by A (change prefactor from 0.5 to 1 or 2)
%A=[0 1;-1 0];% dynamics determined by A (unity norm case)


m=zeros(dim,N);v=m;y=zeros(N,1);c=zeros(dim,dim,N);% pre-allocate space
x(:,1)=m0+sqrtm(C0)*randn(dim,1);% initial truth
m(:,1)=5+10*randn(dim,1);% initial mean/estimate
c(:,:,1)=100*C0;% initial covariance
H=[1,0];% observation operator

%% solution: assimilate data!

for j=1:N   
    x(:,j+1)=A*x(:,j) + sigma*randn(dim,1);% truth
    y(j)=H*x(:,j+1)+gamma*randn;% observation
    
    mhat=A*m(:,j);% mean prediction
    chat=A*c(:,:,j)*A'+sigma^2*I;% covariance prediction
    
    d=y(j)-H*mhat;% innovation
    K=(chat*H')/(H*chat*H'+gamma^2);% Kalman gain
    m(:,j+1)=mhat+K*d;% mean update
    c(:,:,j+1)=(I-K*H)*chat;% covariance update    
end

figure(1);ns=51;
% plot([0:ns-1],m(2,1:ns)+reshape(sqrt(c(2,2,1:ns)),1,ns),'r--');
% plot([0:ns-1],m(2,1:ns)-reshape(sqrt(c(2,2,1:ns)),1,ns),'r--');
cup=m(2,1:ns)+reshape(2.*sqrt(c(2,2,1:ns)),1,ns);
clo=m(2,1:ns)-reshape(2.*sqrt(c(2,2,1:ns)),1,ns);
tspan=[0:ns-1 ns-1:-1:0];
cspan=[cup clo(ns:-1:1)];
p=fill(tspan,cspan,'red');
p.FaceColor = [1 0.8 0.8];      
p.EdgeColor = 'none'; 
hold on
grid;xlabel('iteration');ylabel('estimate')
plot([0:ns-1],x(2,1:ns));plot([0:ns-1],m(2,1:ns),'*');
hold off
title('unobserved component');

figure(2);plot([0:N],reshape(c(1,1,:)+c(2,2,:),N+1,1));hold
plot([0:N],cumsum(reshape(c(1,1,:)+c(2,2,:),N+1,1))./[1:N+1]','m', ...
'Linewidth',2); grid; hold;xlabel('iteration');ylabel('total variance');axis([1 50 0 50]);
legend('covariance (trace)','running average')
title('Kalman filter (co)variance');

figure(3);plot([0:N],sum((x-m).^2));hold;
plot([0:N],cumsum(sum((x-m).^2))./[1:N+1],'m','Linewidth',2);grid
hold;xlabel('iteration');ylabel('error');axis([1 250 0 50]);
legend('L2 error','MSE')
title('Kalman filter error')









