%% Kalman filter demonstration for tracking the heat equation
clear all; close all;
set(0,'DefaultLineLineWidth',2,'DefaultAxesFontSize',[15],...
    'DefaultAxesFontWeight','bold');

gamma = .0001; % observational noise
sigma = .01; % innovation uncertainty
T = .1; % terminal time
ntime = 100;
dt = T/ntime;

%% Simulate measurement data using a dense spatial grid
n0 = 150;
h0 = 1/n0;
x0 = (h0:h0:1-h0)';
L0 = toeplitz([-2 1 zeros(1,n0-3)]);
L0(1,1) = -1; L0(end,end) = -1;
L0 = L0/h0^2;
init0 = x0.^2 .* (1-x0).^2;
y = [];
for ii = 1:ntime
    tmp = expm(dt*ii*L0)*init0;
    y = [y;tmp(1),tmp(end)];
end
y = y+gamma * randn(size(y)); % Contaminate measurements with noise

%% The actual inversion takes place on a coarser computational grid
n = 100;
h = 1/n;
x = (h:h:1-h)';
L = toeplitz([-2 1 zeros(1,n-3)]);
L(1,1) = -1; L(end,end) = -1;
L = L/h^2;
init = x.^2 .* (1-x).^2;

%% Kalman filter
xfilt = zeros(n-1,1);
pr = toeplitz([2 -1 zeros(1,n-3)]);
pr(1,1) = -1; pr(end,end)=-1;
D = .1*inv(pr'*pr);
M = inv(eye(n-1) - dt*L);
H = eye(n-1);
H = H([1,n-1],:);
for ii = 1:ntime
    % Prediction
    xpred = M*xfilt;
    Dpred = M*D*M'+sigma^2*eye(n-1);
    % Correction
    K = Dpred * H'*inv(H*Dpred*H'+gamma^2*eye(2));
    xfilt = xpred + K*(y(ii,:)'-H*xpred);
    D = Dpred - K*H*Dpred;
    % Display results
    if mod(ii,ntime/10) == 0
        figure(ii);
        plot(x(2:end-1),xfilt(2:end-1),'b',x0,expm(dt*ii*L0)*init0,'r--');
        title(['Time $t$ = ',num2str(dt*ii)],'interpreter','latex');
        legend('reconstruction','ground truth');
        drawnow;
        pause(1);
    end
end
