%
% Identify the initial condition in a 1D advection
% diffusion equation from masurements of the solution
% using a regularized least squares approach.
% The PDE is
%
%    y_t(x,t) - diff  y_xx(x,t) + adv y_x(x,t)  = 0  x in (0,1), t in (0,T)
%
%                                        y(0,t) = y(1,t),    t in (0,T)
%                                      y_x(0,t) = y_x(1,t),  t in (0,T)
%
%                                        y(x,0) = y0(x),     x \in (0,1)
%
% The advection diffusion equation is discretized using
% upwind finite differences.
%
%
% Matthias Heinkenschloss
% Jan 9, 2017 (revised Jan 17, 2020)


set(0, 'defaultaxesfontsize',24,'defaultaxeslinewidth',2.2,...
       'defaultlinelinewidth',2.2,'defaultpatchlinewidth',2.2,...
       'defaulttextfontsize',24);
   
clear all



%% The discretizeed differential equation
nx    = 100;     % number of spatial subintervals
h     = 1/nx;    % step size in space
x     = (0:h:1);


diff = 0.001;  % diffusion
adv  = 1;     % advection
% set subdiagonal 
c       = (diff/h^2 + adv/h)*ones(nx,1);
% set diagonal
d       = (-2*diff/h^2 - adv/h)*ones(nx,1);
% set superdiagonal
e       = (diff/h^2)*ones(nx,1);
% generate system matrix
K       = spdiags([c d e], -1:1, nx, nx);
K(1,nx) = (diff/h^2 + adv/h);
K(nx,1) = (diff/h^2);


%% generate observations
I  = eye(nx);
mx = 5;         % nx/mx must be integer 
H  = I((nx/mx):(nx/mx):nx,:);

%% regularization matrix
R  = eye(nx);

%% time discretization parameters
T  = 0.5;   % final time (initial time = 0)
nt = 50;    % number of time steps
dt = T/nt;  % time step size
mt = 25;    % nt/mt must be integer


%% plot the solution of the y'(t) = K y(t);  y(0) = y0_true
fprintf( '\n\n Solve advection difusion equation for true initial data  \n' )
y0_true = exp( -100*(x-0.5).^2)';


yhist           = zeros(nx+1,nt+1);  % record solution at times for plotting

y               = y0_true(2:nx+1);
yhist(2:nx+1,1) = y; yhist(1,1)= yhist(nx+1,1);
I               = speye(size(K,1));

for it = 1:nt
    y  = (I-(0.5*dt)*K) \ ((I+(0.5*dt)*K)*y); 
    yhist(2:nx+1,it+1) = y; yhist(1,it+1)= yhist(nx+1,it+1);
end

figure(1)
mesh( (0:h:1), (0:dt:T), yhist' )
%title(' Solution of y''(t) = K y(t);  y(0) = Gaussian)
axis('tight')
xlabel('x')
ylabel('t')
%print -depsc data-assim-ode1-sol.eps

fprintf('hit return to continue ...\n')
pause


%% Test matrix-vector products and transposes
fprintf( '\n\n Test matrix-vector products and transposes  \n' )

v      = rand(nx,1);
z      = rand(size(H,1)*mt+nx,1);
rho  = 0;
[ z1 ] = Av(v, K, H, R, dt, nt, mt, rho );
[ v1 ] = Atz(z, K, H, R, dt, nt, mt, rho );

fprintf( ' rho                    = %12.6e  \n', rho )
fprintf( ' z''*(A*v)               = %12.6e  \n', z'*z1 )
fprintf( ' (A''*z)''*v              = %12.6e  \n', v1'*v )
fprintf( ' |z''*(A*v) - (A''*z)''*v| = %12.6e  \n', abs(z'*z1-v1'*v) )


v      = rand(nx,1);
z      = rand(size(H,1)*mt+nx,1);
rho  = 1.0;
[ z1 ] = Av(v, K, H, R, dt, nt, mt, rho );
[ v1 ] = Atz(z, K, H, R, dt, nt, mt, rho );

fprintf( ' rho                    = %12.6e  \n', rho )
fprintf( ' z''*(A*v)               = %12.6e  \n', z'*z1 )
fprintf( ' (A''*z)''*v              = %12.6e  \n', v1'*v )
fprintf( ' |z''*(A*v) - (A''*z)''*v| = %12.6e  \n', abs(z'*z1-v1'*v) )

fprintf('hit return to continue ...\n')
pause

%% Generate Observations using   y(0) = y0_true
fprintf( '\n\n Generate Observations \n' )

y0_true = exp( -100*(x-0.5).^2)';


mh = size(H,1);
  
I = speye(size(K,1));
y = y0_true(2:nx+1);
b = zeros(mt*mh,1);
for j = 1:mt
    for i = 1:nt/mt
        y  = (I-(0.5*dt)*K) \ ((I+(0.5*dt)*K)*y); 
    end
    b((j-1)*mh+1:j*mh) = H*y;
end
% add 1% error to data
b   = b .* (1 + 0.01*randn(size(b,1),1));


rho = 0;
fprintf( ' Solve the least squares problem with rho = %12.6e  \n', rho )

maxit = nx;
tol   = 1.e-5;

% We solve the least squares problem with weighting matrix W = I
% compute A'*b
b   = [ b; zeros(nx,1)];
Atb = Atz(b, K, H, R, dt, nt, mt, rho );

afun  = @(x)AtAv(x, K, H, R, dt, nt, mt, rho );

[y0_comp,flag,relres,iter,resvec] = pcg(afun,Atb,tol,maxit);
y0_comp = [y0_comp(nx); y0_comp];
fprintf( ' pcg returned after %d iterations with flag %d  \n', iter, flag)

figure(2)
semilogy((0:iter), resvec, '*-' ); 
xlabel('iteration')
ylabel('CG residual')

figure(1)
plot( (0:h:1), y0_true, 'r--' ); hold on
plot( (0:h:1), y0_comp, 'b-' ); hold off
axis([0,1,-0.2,1.5])
xlabel('x')
legend('true y_0','computed y_0')
%print -depsc data-assim-ode1-comp-rho0.eps

fprintf('hit return to continue ...\n')
pause

rho = 1.e-2;
fprintf( ' Solve the least squares problem with rho = %12.6e  \n', rho )

maxit = nx;
tol   = 1.e-5;
% We solve the least squares problem with weighting matrix W = I
% A'*b as before

afun  = @(x)AtAv(x, K, H, R, dt, nt, mt, rho );

[y0_comp,flag,relres,iter,resvec] = pcg(afun,Atb,tol,maxit);
y0_comp = [y0_comp(nx); y0_comp];
fprintf( ' pcg returned after %d iterations with flag %d  \n', iter, flag)

figure(2)
semilogy((0:iter), resvec, '*-' ); 
xlabel('iteration')
ylabel('CG residual')

figure(1)
plot( (0:h:1), y0_true, 'r--' ); hold on
plot( (0:h:1), y0_comp, 'b-' ); hold off
axis([0,1,-0.2,1.5])
xlabel('x')
legend('true y_0','computed y_0')
%print -depsc data-assim-ode1-comp-rho1-2.eps
