%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Program to perform data assimilation
%  on a damped oscillating system.
%
%
%  (c) 2002  Data Assimilation Research Centre
%
%  Original program by Matthew Martin
%
%  Change history:
%      Change to read error from file produced by randosc
%                       (Steven Fletcher)
%      Changes to make code more robust and add further options
%                       (Amos Lawless)
%      Added menu interface (Amos Lawless)
%      23/01/03 Change to position windows correctly for any screen
%               resolution (Amos Lawless)
%
%  Structure of program:
%    1. Inputs
%    2. `True' solution
%    3. Background solution
%    4. Observations
%    5. Covariance matrices
%    6. Assimilation
%    7. Plot results.
%
%  List of main variables
%    h:          Time step for Runge-Kutta scheme
%    l:          Damping coefficient
%    m:          Square of frequency
%
%    fmax:       Total time of assimilation + forecast
%    ob_f:       Frequency of observations in time steps
%    q:          x-axis data values for plots
%    sd:         Variance of observation error
%    tmax:       Total time of assimilation
%
%    xob(2,tmax): Matrix of observations
%                 Zero indicates no observation
%    R(2,2):      Observation error covariance matrix
%    B(2,2):      Background error covariance matrix
%    Qx(2,2):     Model error covariance matrix for Kalman Filter
%                 (default to zero)
%    Mt(2,2):     True model
%    Mb(2,2):     Background state evolution model
%
%    x(2,fmax):    True state vector
%                  x(:,i) is [x_1,x_2] at time i
%    xb(2,fmax):   Background state
%    x_sc(2,fmax): Analysis and forecast from SC scheme
%    x_ac(2,fmax): Analysis and forecast from AC scheme
%    x_oi(2,fmax): Analysis and forecast from OI scheme
%    x_kf(2,fmax): Analysis and forecast from Kalman Filter
%
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Close any open figure windows
g=gcf;
for i=1:g
  close(i)
end
% Determine screen size and set up figure positions
ss=get(0,'ScreenSize');
fig_width=0.5*(ss(3)*0.9);
fig_height=0.5*ss(4);
pos_1=[0.02*ss(3),0.45*ss(4),fig_width,fig_height];
pos_2=[fig_width+0.05*ss(3),0.45*ss(4),fig_width,fig_height];
p3l=0.6*ss(3);
p3b=0.05*ss(4);
p3w=0.3*ss(3);
p3h=0.3*ss(4);
pos_3=[p3l,p3b,p3w,p3h];
pos_4=[0.02*p3w,0.8*p3h,0.95*p3w,0.2*p3h];
pos_5=[0.02*p3w,0.02*p3h,0.95*p3w,0.75*p3h];
%------------------------------------------------------------------------
%%%%%%%%%%1.  INPUTS
%------------------------------------------------------------------------
l_assim=menu_asl('Please choose an assimilation scheme','Successive correction', ...
   'Optimal Interpolation', 'Analysis Correction','Kalman Filter');
l_assim=l_assim-1;
if (l_assim==0 | l_assim==2)
  n_its=menu_asl('How many iterations?','1','2','3','4','5');
end
if l_assim~=0
  l_weight=menu_asl('Use correct weighting matrices?','No','Yes');
  l_weight=l_weight-1;
end

fmax=500;	%time of assimilation + forecast
q=0:0.1:fmax*0.1-0.1;

%------------------------------------------------------------------------
%%%%%%%%%2.  TRUE SOLUTION 
%------------------------------------------------------------------------
%true solution of oscillating system using 2nd order Runge-Kutta method
x=zeros(2,fmax);  %true state vector
Mt=zeros(2,2);    %true model

%parameters
h=0.1;    % time step for Runge-Kutta scheme
l=0.1;    % damping coefficient
m=1.0;    % square of frequency
%initial conditions
x(1,1)=1.0;
x(2,1)=0.0;

%state evolution model
Mt(1,1)=(1-(((h^2)/2)*m));
Mt(1,2)=h*(1-((h/2)*l));
Mt(2,1)=(h*m*(((h/2)*l)-1));
Mt(2,2)=(1+(h*(-l+((h/2)*(l^2))-((h/2)*m))));

for i=1:fmax-1
  x(:,i+1)=Mt*x(:,i);  
end

disp(['*   Done true solution'])
%Plot true solution of x
h1=figure('Position',pos_1);
clf;	
plot(q,x(1,:),'k--');
legend('Truth')
xlabel('Time')
ylabel('x_1')
title('Plot of x_1 against time')
hold on
%----------------------------------------------------------------------------
%%%%%%%%3.  BACKGROUND SOLUTION 
%----------------------------------------------------------------------------
xb=zeros(2,fmax);

%Background solution initial conditions
xb(1,1)=1.5;
xb(2,1)=0.5;

%background state evolution model
Mb=Mt;

%start loop over time
for i=1:fmax-1
  xb(:,i+1)=Mb*xb(:,i);
end
%end loop over time

disp(['*   Done background solution'])
%Plot background solution
plot(q,xb(1,:),'b-.');
legend('Truth','Background')
%
%---------------------------------------------------------------
%%%%%%%%%4.  OBSERVATIONS 
%---------------------------------------------------------------
l_obs=menu_asl('How many time steps between observations?','1','2','5', ...
  '10','25');
  if l_obs==1
    ob_f=1;
  elseif l_obs==2
    ob_f=2;	
  elseif l_obs==3
    ob_f=5;	
  elseif l_obs==4
    ob_f=10;	
  elseif l_obs==5
    ob_f=25;	
  end
l_noise=menu_asl('Noise on observations?','No','Yes');
l_noise=l_noise-1;
if (l_noise~=0 & l_noise~=1)
  error('Invalid input for noise on observations')
end
tmax=250+ob_f;  %time of assimilation
%
xob=zeros(2,tmax);
vec=1:ob_f:tmax;
nobs=tmax/ob_f;	%number of obs
%
if l_noise==1
  l_read_noise=menu_asl('How read noise?','Generate in program',...
    'Read from file');
  l_read_noise=l_read_noise-1;
  if (l_read_noise==0)
    sc_x_noise=randn(nobs,1);
    sc_y_noise=randn(nobs,1);
    save ('noise.mat','sc_x_noise','sc_y_noise');
  elseif (l_read_noise==1)
    load('noise.mat','sc_x_noise','sc_y_noise');
  else
    error('Invalid input')
  end
  sd_str=inputdlg('Variance of observation error');
  sd=str2num(sd_str{1});
  var=sqrt(sd);
%
  NX=var*sc_x_noise;
  NY=var*sc_y_noise;
else
  for i=1:nobs;
    NX(i)=0.0;
    NY(i)=0.0;
  end
end

j=1;
for i=vec
  xob(1,i)=x(1,i)+NX(j);
  xob(2,i)=x(2,i)+NY(j);
  j=j+1;
end
disp(['*   Done observations'])
%Plot x observations
v=0.1*vec;
plot(v(2:nobs),xob(1,vec(2:nobs)),'om');
legend('Truth','Background','Observations')
hold on
l_test=menu_asl('Click to perform analysis','OK');
%---------------------------------------------------------------------
%%%%%%%%%%5.  ERROR COVARIANCE MATRICES 
%---------------------------------------------------------------------
if l_assim~=0 

BX=zeros(2,2);
BXi=zeros(2,2,tmax);
R=zeros(2,2);
Ri=zeros(2,2,nobs);

% Calculate observation error covariance at each observation time
for i=vec
  Ri(1,1,i)=(x(1,i)-xob(1,i))'*(x(1,i)-xob(1,i));
  Ri(1,2,i)=(x(1,i)-xob(1,i))'*(x(2,i)-xob(2,i));
  Ri(2,1,i)=(x(2,i)-xob(2,i))'*(x(1,i)-xob(1,i));
  Ri(2,2,i)=(x(2,i)-xob(2,i))'*(x(2,i)-xob(2,i));
end

if l_weight==1

% Calculate average observation error covariance matrix
  Rj=zeros(2,2,nobs);
  for j=1:2
    for k=1:2
      Rj(j,k,:)=Ri(j,k,vec);
      R(j,k)=(1/nobs)*sum(Rj(j,k,:));
    end
  end

% Calculate average background error covariance matrix
  for i=1:tmax
    BXi(1,1,i)=(x(1,i)-xb(1,i))'*(x(1,i)-xb(1,i));
    BXi(1,2,i)=(x(1,i)-xb(1,i))'*(x(2,i)-xb(2,i));
    BXi(2,1,i)=(x(2,i)-xb(2,i))'*(x(1,i)-xb(1,i));
    BXi(2,2,i)=(x(2,i)-xb(2,i))'*(x(2,i)-xb(2,i));
  end
  for j=1:2
    for k=1:2
      BX(j,k)=(1/tmax)*sum(BXi(j,k,:));
    end
  end

else

  R=eye(2);   % Set to identity matrix for l_weight==0
  BX=eye(2);  % Set to identity matrix for l_weight==0

end  % l_weight==1

end  % l_assim~=0
%-----------------------------------------------------------------------------------
%%%%%%%%%6.  ASSIMILATION
%-----------------------------------------------------------------------------------

if l_assim==0
%%%%%%%%%%%%%%%%%% SUCCESSIVE CORRECTIONS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp('Performing successive correction analysis')
assim_str='Successive correction';
WX=zeros(2,2);
x_sc=zeros(2,fmax);

%state weighting matrix
WX=0.5*eye(2,2);
x_sc(:,1)=xb(:,1);

%start loop over time
for i=1:fmax-1
  if i>=tmax
    xob(:,i+1)=0;
  end

%if there are observations at this time step
  if xob(:,i+1)~=0

%forecast on to time step i+1 
    x_sc(:,i+1)=Mb*x_sc(:,i);
  
%do analysis at time step i+1
    for j=1:n_its
      x_sc(:,i+1)=x_sc(:,i+1)+WX*(xob(:,i+1)-x_sc(:,i+1));
    end

%if there aren't observations 
  else

%forecast on to time step i+1
   x_sc(:,i+1)=Mb*x_sc(:,i);

  end
end
%end loop over time


elseif l_assim==1
%%%%%%%%%%%%%%%%%% OPTIMAL INTERPOLATION %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp('Performing OI analysis')
assim_str='Optimal interpolation';
WX=zeros(2,2);
x_oi=zeros(2,fmax);

if l_weight==1
%state weighting matrix
 WX=BX*pinv(BX+R);
else
 WX=eye(2,2);
end

x_oi(:,1)=xb(:,1);


%start loop over time
for i=1:fmax-1
  if i>=tmax
    xob(:,i+1)=0;
  end

%forecast on to time step i+1
    x_oi(:,i+1)=Mb*x_oi(:,i);

%if there are observations then do analysis at time step i+1
  if xob(:,i+1)~=0
     x_oi(:,i+1)=x_oi(:,i+1)+WX*(xob(:,i+1)-x_oi(:,i+1)); 
  end

%end loop over time
end


elseif l_assim==2
%%%%%%%%%%%%%%%%%%%%%%%%%% ANALYSIS CORRECTION %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp('Performing AC analysis')
assim_str='Analysis correction';
WX=zeros(2,2);
QX=zeros(2,2); %=(WX+I)^{-1}
x_ac=zeros(2,fmax);
y_ac=zeros(2,fmax);

%state weighting matrix
WX=BX*pinv(R);
QX=pinv(WX+eye(2,2));

x_ac(:,1)=xb(:,1);
y_ac(:,1)=xob(:,1);

%start loop over time
for i=1:fmax-1

  if i>=tmax
    xob(:,i+1)=0;
  end
%forecast on to time step i+1
    x_ac(:,i+1)=Mb*x_ac(:,i);
%if there are observations then do analysis at time step i+1
  if xob(:,i+1)~=0
    y_ac(:,i+1)=xob(:,i+1);
%
    for j=1:n_its
      x_ac(:,i+1)=x_ac(:,i+1)+WX*QX*(y_ac(:,i+1)-x_ac(:,i+1));
      y_ac(:,i+1)=y_ac(:,i+1)-QX*(y_ac(:,i+1)-x_ac(:,i+1));
    end
  end
end
%end loop over time

elseif l_assim==3
%%%%%%%%%%%%%%%%%%%%%%%%%% KALMAN FILTER %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp('Performing kalman filter analysis')
assim_str='Kalman filter';
KX=zeros(2,2,fmax);
Pfx=zeros(2,2,fmax);
Pax=zeros(2,2,fmax);
Qx=zeros(2,2);
x_kf=zeros(2,fmax);

%set model error covariance matrix
Qx(:,:)=0.0*eye(2,2);

%set forecast at i=1 to be the background guess there 
x_kf(:,1)=xb(:,1);

if l_weight==1
%forecast error covariance at i=1 is the background error 
%covariance there
  Pfx(:,:,1)=BXi(:,:,1);
else
%begin with some incorrect guess
  Pfx(:,:,1)=eye(2);
end

if xob(:,1)~=0
%calculate the gain matrix at i=1
KX(:,:,1)=Pfx(:,:,1)*pinv(Pfx(:,:,1)+Ri(:,:,1));

%calculate analysis error covariance matrix at i=1
Pax(:,:,1)=(eye(2,2)-KX(:,:,1))*Pfx(:,:,1);

else
  Pax(:,:,1)=Pfx(:,:,1);  %ASL
end

%start loop over time
for i=1:fmax-1
  if i>=tmax
    xob(:,i+1)=0;
  end

%forecast model
  x_kf(:,i+1)=Mb*x_kf(:,i);

%forecast error covariance matrix
  Pfx(:,:,i+1)=Mb(:,:)*Pax(:,:,i)*Mb(:,:)'+Qx(:,:);

%if there are observations
  if xob(:,i+1)~=0

%calculate the gain matrix 
    KX(:,:,i+1)=Pfx(:,:,i+1)*pinv(Pfx(:,:,i+1)+Ri(:,:,i+1));

%produce an analysis
    x_kf(:,i+1)=x_kf(:,i+1)+KX(:,:,i+1)*(xob(:,i+1)-x_kf(:,i+1));

%calculate analysis error covariance matrix
    Pax(:,:,i+1)=(eye(2,2)-KX(:,:,i+1))*Pfx(:,:,i+1);

%if there are no observations
  else

    Pax(:,:,i+1)=Pfx(:,:,i+1);

  end

%end loop over time
end

end
%End if on type of assimilation

disp(['***  FINISHED  ***'])
%-------------------------------------------------------------
%%%%%%%%%%7.  PLOT RESULTS
%-------------------------------------------------------------
%Plot analysed solution of x
if l_assim==0
  plot(q,x_sc(1,:),'r-');
elseif l_assim==1
  plot(q,x_oi(1,:),'r-');
elseif l_assim==2
  plot(q,x_ac(1,:),'r-');
elseif l_assim==3
  plot(q,x_kf(1,:),'r-');
end
legend('Truth','Background','Observations','Analysis')

yminmax=get(gca,'YLim');
yspace=(yminmax(2)-yminmax(1))*0.01;
yvals=yminmax(1):yspace:yminmax(2);
xval=(tmax-ob_f)*h;
line(xval,yvals,'LineStyle',':','Color','k')

%Plot errors in x
h2=figure('Position',pos_2);
clf;	
if l_assim==0
  plot(q,x(1,:)-x_sc(1,:),'k-')
elseif l_assim==1
  plot(q,x(1,:)-x_oi(1,:),'k-')
elseif l_assim==2
  plot(q,x(1,:)-x_ac(1,:),'k-')
elseif l_assim==3
  plot(q,x(1,:)-x_kf(1,:),'k-')
end
line(xval,yvals,'LineStyle','-','Color','k')
xlabel('Time')
ylabel('Error')
title('Plot of error (Truth-Analysis) against time')

%-----------------------------------------------------------------------
% Write out options chosen
%
h4=figure('Position',pos_3);
clf;
text1={'List of options chosen'};
text2={['Analysis scheme: ' assim_str]};
if (l_assim==0 | l_assim==2)
  text3={['Number of iterations: ' num2str(n_its)]};
else
  text3='';
end
if l_assim~0
  if l_weight==0
    text4={['Using incorrect weighting matrices']};
  else
    text4={['Using correct weighting matrices']};
  end
else
  text4='';
end
text5={['Time steps between observations: ' num2str(ob_f)]};
if l_noise==0
  text6={['No noise on observations']};
  text7='';
else
  text6={['Observations have random noise with variance ' num2str(sd)]};
  if l_read_noise==0
    text7={['Noise generated in program and saved to file']};
  else
    text7={['Noise read in from previously generated file']};
  end
end
%
str1=[text2;text3;text4;text5;text6;text7];
uicontrol('Style','text','Position',pos_4,'String',text1,...
 'FontSize',14,'FontWeight','bold')
uicontrol('Style','text','Position',pos_5,'String',str1,...
 'FontSize',12,'HorizontalAlignment','left')

%%%%%%%%%%%%%%% END OF PROGRAM %%%%%%%%%%%%%%%%%%%