fmax=1000;	%time of assimilation + forecast
q=0:0.01:fmax*0.01-0.01;

%-------------------------------------------------------------------------
%%%%%%%%%2.  TRUE SOLUTION 
%-------------------------------------------------------------------------
%true solution of Lorenz 1963 model equations using 2nd order Runge-Kutta method
x=zeros(fmax,1);
y=zeros(fmax,1);
z=zeros(fmax,1);

%parameters
h=0.01;    % time step for Runge-Kutta scheme
s=10.0;    % sigma
r=28;      % rho
b=8/3;     % beta

%initial conditions
x(1)=-5.4458d0;
y(1)=-5.4841d0;
z(1)=22.5606d0;

for i=1:fmax-1
k1x(i)=s*y(i)-s*x(i);
k1y(i)=r*x(i)-y(i)-x(i)*z(i);
k1z(i)=x(i)*y(i)-b*z(i);

k2x(i)=s*(y(i)+h*k1y(i))-s*(x(i)+h*k1x(i));
k2y(i)=r*(x(i)+h*k1x(i))-(y(i)+h*k1y(i))-(x(i)+h*k1x(i))*(z(i)+h*k1z(i));
k2z(i)=(x(i)+h*k1x(i))*(y(i)+h*k1y(i))-b*(z(i)+h*k1z(i));

x(i+1)=x(i)+0.5*h*(k1x(i)+k2x(i));
y(i+1)=y(i)+0.5*h*(k1y(i)+k2y(i));
z(i+1)=z(i)+0.5*h*(k1z(i)+k2z(i));	
  
end

%-----------------------------------------------------------------------
%%%%%%%%%3.  BACKGROUND SOLUTION 
%-----------------------------------------------------------------------
%Produce a background solution which starts from some perturbed initial
%conditions.

xb=zeros(fmax,1);
yb=zeros(fmax,1);
zb=zeros(fmax,1);

%Background solution initial conditions

xb(1)=-5.9d0;
yb(1)=-5.0d0;
zb(1)=24.0d0;

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

k1xb(i)=s*yb(i)-s*xb(i);
k1yb(i)=r*xb(i)-yb(i)-xb(i)*zb(i);
k1zb(i)=xb(i)*yb(i)-b*zb(i);

k2xb(i)=s*(yb(i)+h*k1yb(i))-s*(xb(i)+h*k1xb(i));
k2yb(i)=r*(xb(i)+h*k1xb(i))-(yb(i)+h*k1yb(i))-(xb(i)+h*k1xb(i))...
*(zb(i)+h*k1zb(i));
k2zb(i)=(xb(i)+h*k1xb(i))*(yb(i)+h*k1yb(i))-b*(zb(i)+h*k1zb(i));

xb(i+1)=xb(i)+0.5*h*(k1xb(i)+k2xb(i));
yb(i+1)=yb(i)+0.5*h*(k1yb(i)+k2yb(i));
zb(i+1)=zb(i)+0.5*h*(k1zb(i)+k2zb(i));	
  
end

disp(['*   done background solution'])
l_obs=input('Please select the number of time steps between observations: 1=25, 2=50, 3=100, 4=200, ');


  if l_obs==1
    ob_f=25;	
  elseif l_obs==2
    ob_f=50;	
  elseif l_obs==3
    ob_f=100;	
  elseif l_obs==4
    ob_f=200;	
  end
  l_noise=input('Is there noise on the observations:, 1=no, 2=yes, ');

l_noise=l_noise-1;
tmax=fmax;


%The observations are calculated from the true solution at certain 
%points with noise added. 

xob=zeros(tmax,1);
yob=zeros(tmax,1);
zob=zeros(tmax,1);
zobm=zeros(tmax,1);

nobs=tmax/ob_f;	%number of observations
vec=1:ob_f:tmax;

if l_noise==1
  l_read_noise=input('How to read noise?, 1=Generate in program, 2=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);
    sc_z_noise=randn(nobs,1);
    save ('noise.mat','sc_x_noise','sc_y_noise','sc_z_noise');
  elseif (l_read_noise==1)
    load('noise.mat','sc_x_noise','sc_y_noise','sc_z_noise');
  else
    error('Invalid input')
  end
  sd=input('Please enter the value of the observational error variance, ');
  var=sqrt(sd);
  RX=var*sc_x_noise;
  RY=var*sc_y_noise;
  RZ=var*sc_z_noise;
else
  for i=1:nobs;
    RX(i)=0.0;
    RY(i)=0.0;
    RZ(i)=0.0;
  end
end

j=1;
for i=vec
xob(i)=x(i)+RX(j);
yob(i)=y(i)+RY(j);
zob(i)=z(i)*exp(RZ(j)/40);
zobm(i)=log(zob(i));
j=j+1;
end
disp(['*   done observations'])

R=zeros(3,3);
Ri=zeros(3,3,nobs);

%Observation error covariance matrix, R.

  for i=vec
    Ri(1,1,i)=(x(i)-xob(i))'*(x(i)-xob(i));
    Ri(1,2,i)=(x(i)-xob(i))'*(y(i)-yob(i));
    Ri(1,3,i)=(x(i)-xob(i))'*(z(i)-zob(i));
    Ri(2,1,i)=(y(i)-yob(i))'*(x(i)-xob(i));
    Ri(2,2,i)=(y(i)-yob(i))'*(y(i)-yob(i));
    Ri(2,3,i)=(y(i)-yob(i))'*(z(i)-zob(i));
    Ri(3,1,i)=(z(i)-zob(i))'*(x(i)-xob(i));
    Ri(3,2,i)=(z(i)-zob(i))'*(y(i)-yob(i));
    Ri(3,3,i)=(z(i)-zob(i))'*(z(i)-zob(i));
  end
 
 for i=1:tmax
    BXi(1,1,i)=(x(i)-xb(i))'*(x(i)-xb(i));
    BXi(1,2,i)=(x(i)-xb(i))'*(y(i)-yb(i));
    BXi(1,3,i)=(x(i)-xb(i))'*(z(i)-zb(i));
    BXi(2,1,i)=(y(i)-yb(i))'*(x(i)-xb(i));
    BXi(2,2,i)=(y(i)-yb(i))'*(y(i)-yb(i));
    BXi(2,3,i)=(y(i)-yb(i))'*(z(i)-zb(i));
    BXi(3,1,i)=(z(i)-zb(i))'*(x(i)-xb(i));
    BXi(3,2,i)=(z(i)-zb(i))'*(y(i)-yb(i));
    BXi(3,3,i)=(z(i)-zb(i))'*(z(i)-zb(i));
  end
disp('Performing Gaussian kalman filter analysis')

KX=zeros(3,3,fmax);
Pfx=zeros(3,3,fmax);
Pax=zeros(3,3,fmax);
xfc=zeros(tmax);
yfc=zeros(tmax);
zfc=zeros(tmax);
x_kf=zeros(3,fmax);
x_fc=zeros(3,fmax);
x_ob=zeros(3,fmax);
Mx=zeros(3,3,fmax);
Qx=zeros(3,3,fmax);

%put obs into one array
x_ob=[xob';yob';zob'];

%set forecast at i=1 to be the background guess there 
xfc(1)=xb(1);
yfc(1)=yb(1);
zfc(1)=zb(1);

%put forecast values into one array
x_fc(:,1)=[xb(1);yb(1);zb(1)];

Pfx(:,:,1)=BXi(:,:,1);

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

x_kf(:,1)=x_fc(:,1);

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

for i=1:fmax-1
  if i>=tmax
    x_ob(:,i+1)=0;
  end

  xfc(i)=x_kf(1,i);
  yfc(i)=x_kf(2,i);
  zfc(i)=x_kf(3,i);

%tangent-linear model - linearised about the forecast state
  Mx(1,1,i)=1.0-(h*s)+(((h*h*s)/2.0)*(r+s-zfc(i)));
  Mx(1,2,i)=h*s-(((h*h*s)/2.0)*(1.0+s));
  Mx(1,3,i)=-(h*h*s*xfc(i))/2.0;
  Mx(2,1,i)=(h*(r-zfc(i)))+(((h*h)/2.0)*(zfc(i)-(r*s)-r+(b*zfc(i))...
  +(s*zfc(i))-(2.0*xfc(i)*yfc(i))))+(((h*h*h*s)/2.0)*((b*zfc(i))-...
  (yfc(i)*yfc(i))+(2.0*xfc(i)*yfc(i))));
  Mx(2,2,i)=1.0-h+(((h*h)/2.0)*((s*r)+1.0-(xfc(i)*xfc(i))-(s*zfc(i))))...
  +(((h*h*h*s)/2.0)*((b*zfc(i))-(2.0*xfc(i)*yfc(i))+(xfc(i)*xfc(i))));
  Mx(2,3,i)=(-h*xfc(i))+(((h*h)/2.0)*(xfc(i)+(b*xfc(i))-(s*yfc(i))+...
  (s*xfc(i))))+(((h*h*h*s*b)/2.0)*(yfc(i)+xfc(i)));
  Mx(3,1,i)=(h*yfc(i))+(((h*h)/2.0)*((2.0*r*xfc(i))-yfc(i)-(2.0*xfc(i)...
  *zfc(i))-(s*yfc(i))-(b*yfc(i))))+(((h*h*h*s)/2.0)*((r*yfc(i))-(yfc(i)...
  *zfc(i))-(2.0*r*xfc(i))+yfc(i)+(2.0*xfc(i)*zfc(i))));
  Mx(3,2,i)=(h*xfc(i))+(((h*h)/2.0)*((s*yfc(i))-xfc(i)-(s*xfc(i))-...
  (b*xfc(i))))+(((h*h*h*s)/2.0)*((r*xfc(i))-(2.0*yfc(i))-(xfc(i)*zfc(i))...
  +xfc(i)));
  Mx(3,3,i)=1.0-(h*b)+(((h*h)/2.0)*((b*b)-(xfc(i)*xfc(i))))+...
  (((h*h*h*s)/2.0)*((xfc(i)*xfc(i))-(xfc(i)*yfc(i))));

%Choose something for the model error covariance matrix.
%These are from Evensen, 1997.
  Qx(1,1,i)=0.1491;
  Qx(1,2,i)=0.1505;
  Qx(1,3,i)=0.0007;
  Qx(2,1,i)=0.1505;
  Qx(2,2,i)=0.9048;
  Qx(2,3,i)=0.0014;
  Qx(3,1,i)=0.0007;
  Qx(3,2,i)=0.0014;
  Qx(3,3,i)=0.9180;

%forecast the state
  k1xfc(i)=s*yfc(i)-s*xfc(i);
  k1yfc(i)=r*xfc(i)-yfc(i)-xfc(i)*zfc(i);
  k1zfc(i)=xfc(i)*yfc(i)-b*zfc(i);

  k2xfc(i)=s*(yfc(i)+h*k1yfc(i))-s*(xfc(i)+h*k1xfc(i));
  k2yfc(i)=r*(xfc(i)+h*k1xfc(i))-(yfc(i)+h*k1yfc(i))-(xfc(i)+h*k1xfc(i))...
  *(zfc(i)+h*k1zfc(i));
  k2zfc(i)=(xfc(i)+h*k1xfc(i))*(yfc(i)+h*k1yfc(i))-b*(zfc(i)+h*k1zfc(i));

  xfc(i+1)=xfc(i)+0.5*h*(k1xfc(i)+k2xfc(i));
  yfc(i+1)=yfc(i)+0.5*h*(k1yfc(i)+k2yfc(i));
  zfc(i+1)=zfc(i)+0.5*h*(k1zfc(i)+k2zfc(i));	

  x_fc(:,i+1)=[xfc(i+1);yfc(i+1);zfc(i+1)];

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

%If there is an observation at this time 
  if x_ob(:,i+1)~=0

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

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

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

  else

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

  end

end
disp(['***  FINISHED GAUSSIAN KALMAN FILTER  ***'])
x_kf_g=x_kf;

clear Ri;

for i=vec
    Ri(1,1,i)=(x(i)-xob(i))'*(x(i)-xob(i));
    Ri(1,2,i)=(x(i)-xob(i))'*(y(i)-yob(i));
    Ri(1,3,i)=(x(i)-xob(i))'*(log(z(i))-log(zob(i)));
    Ri(2,1,i)=(y(i)-yob(i))'*(x(i)-xob(i));
    Ri(2,2,i)=(y(i)-yob(i))'*(y(i)-yob(i));
    Ri(2,3,i)=(y(i)-yob(i))'*(log(z(i))-log(zob(i)));
    Ri(3,1,i)=(log(z(i))-log(zob(i)))'*(x(i)-xob(i));
    Ri(3,2,i)=(log(z(i))-log(zob(i)))'*(y(i)-yob(i));
    Ri(3,3,i)=(log(z(i))-log(zob(i)))'*(log(z(i))-log(zob(i)));
end
  
clear BXi

 for i=1:tmax
    BXi(1,1,i)=(x(i)-xb(i))'*(x(i)-xb(i));
    BXi(1,2,i)=(x(i)-xb(i))'*(y(i)-yb(i));
    BXi(1,3,i)=(x(i)-xb(i))'*(log(z(i))-log(zb(i)));
    BXi(2,1,i)=(y(i)-yb(i))'*(x(i)-xb(i));
    BXi(2,2,i)=(y(i)-yb(i))'*(y(i)-yb(i));
    BXi(2,3,i)=(y(i)-yb(i))'*(log(z(i))-log(zb(i)));
    BXi(3,1,i)=(log(z(i))-log(zb(i)))'*(x(i)-xb(i));
    BXi(3,2,i)=(log(z(i))-log(zb(i)))'*(y(i)-yb(i));
    BXi(3,3,i)=(log(z(i))-log(zb(i)))'*(log(z(i))-log(zb(i)));
 end
  
KX=zeros(3,3,fmax);
Pfx=zeros(3,3,fmax);
Pax=zeros(3,3,fmax);
xfc=zeros(tmax);
yfc=zeros(tmax);
zfc=zeros(tmax);
x_kf=zeros(3,fmax);
x_fc=zeros(3,fmax);
x_ob=zeros(3,fmax);
Mx=zeros(3,3,fmax);
Qx=zeros(3,3,fmax);

disp('Performing Mixed Gaussian-lognormal kalman filter analysis')

%put obs into one array
x_ob=[xob';yob';zob']; %Gaussian
x_obmx=[xob';yob';zobm'];  %Lognormal

%set forecast at i=1 to be the background guess there 
xfc(1)=xb(1);
yfc(1)=yb(1);
zfc(1)=zb(1);

x_fc(:,1)=[xb(1);yb(1);zb(1)];
Pfx(:,:,1)=BXi(:,:,1);

wb(:,:,1)=eye(3);
wb(3,3,1)=zfc(1);

wo(:,:,1)=eye(3);
wo(3,3,1)=1/zob(1);

KX(:,:,1)=Pfx(:,:,1)*wb(:,:,1)*wo(:,:,1)*pinv(wo(:,:,1)*wb(:,:,1)*Pfx(:,:,1)*wb(:,:,1)*wo(:,:,1)+Ri(:,:,1));
x_kf(:,1)=x_fc(:,1);

Pax(:,:,1)=(eye(3,3)-KX(:,:,1)*wo(:,:,1)*wb(:,:,1))*Pfx(:,:,1);

for i=1:fmax-1
  if i>=tmax
    x_ob(:,i+1)=0;
  end

  xfc(i)=x_kf(1,i);
  yfc(i)=x_kf(2,i);
  zfc(i)=x_kf(3,i);
  
    %Choose something for the model error covariance matrix.
%These are from Evensen, 1997.
  Qx(1,1,i)=0.1491;
  Qx(1,2,i)=0.1505;
  Qx(1,3,i)=0.0007;
  Qx(2,1,i)=0.1505;
  Qx(2,2,i)=0.9048;
  Qx(2,3,i)=0.0014;
  Qx(3,1,i)=0.0007;
  Qx(3,2,i)=0.0014;
  Qx(3,3,i)=0.9180;
  
  k1xfc(i)=s*yfc(i)-s*xfc(i);
  k1yfc(i)=r*xfc(i)-yfc(i)-xfc(i)*zfc(i);
  k1zfc(i)=xfc(i)*yfc(i)-b*zfc(i);

  k2xfc(i)=s*(yfc(i)+h*k1yfc(i))-s*(xfc(i)+h*k1xfc(i));
  k2yfc(i)=r*(xfc(i)+h*k1xfc(i))-(yfc(i)+h*k1yfc(i))-(xfc(i)+h*k1xfc(i))...
  *(zfc(i)+h*k1zfc(i));
  k2zfc(i)=(xfc(i)+h*k1xfc(i))*(yfc(i)+h*k1yfc(i))-b*(zfc(i)+h*k1zfc(i));

  xfc(i+1)=xfc(i)+0.5*h*(k1xfc(i)+k2xfc(i));
  yfc(i+1)=yfc(i)+0.5*h*(k1yfc(i)+k2yfc(i));
  zfc(i+1)=zfc(i)+0.5*h*(k1zfc(i)+k2zfc(i));	

  x_fc(:,i+1)=[xfc(i+1);yfc(i+1);zfc(i+1)];
  x_mx(:,i+1)=[xfc(i+1);yfc(i+1);log(zfc(i+1))];
  
    Pfx(1,1,i+1)=(xfc(i)-x(i))'*(xfc(i)-x(i))+Qx(1,1,i);
    Pfx(1,2,i+1)=(xfc(i)-x(i))'*(yfc(i)-y(i))+Qx(1,2,i);
    Pfx(1,3,i+1)=(xfc(i)-x(i))'*(log(zfc(i))-log(z(i)))+Qx(1,3,i);
    Pfx(2,1,i+1)=(yfc(i)-y(i))'*(xfc(i)-x(i))+Qx(2,1,i);
    Pfx(2,2,i+1)=(yfc(i)-y(i))'*(yfc(i)-y(i))+Qx(2,2,i);
    Pfx(2,3,i+1)=(yfc(i)-y(i))'*(log(zfc(i))-log(z(i)))+Qx(2,3,i);
    Pfx(3,1,i+1)=(log(zfc(i))-log(z(i)))'*(xfc(i)-x(i))+Qx(3,1,i);
    Pfx(3,2,i+1)=(log(zfc(i))-log(z(i)))'*(yfc(i)-y(i))+Qx(3,2,i);
    Pfx(3,3,i+1)=(log(zfc(i))-log(z(i)))'*(log(zfc(i))-log(z(i)))+Qx(3,3,i);
    
%If there is an observation at this time 
  if x_ob(:,i+1)~=0
wb(:,:,i+1)=eye(3);
wb(3,3,i+1)=zfc(i+1);

wo(:,:,i+1)=eye(3);
wo(3,3,i+1)=1/zob(i+1);
%calculate the gain matrices 

KX(:,:,i+1)=Pfx(:,:,i+1)*wb(:,:,i+1)*wo(:,:,i+1)*pinv(wo(:,:,i+1)*wb(:,:,i+1)*Pfx(:,:,i+1)*wb(:,:,i+1)*wo(:,:,i+1)+Ri(:,:,i+1));

%produce an analysis

x_kfi(:,i+1)=x_mx(:,i+1)+KX(:,:,i+1)*(x_obmx(:,i+1)-x_mx(:,i+1));

x_kf(1,i+1)=x_kfi(1,i+1);
x_kf(2,i+1)=x_kfi(2,i+1);
x_kf(3,i+1)=exp(x_kfi(3,i+1));

%calculate analysis error covariance matrices

Pax(:,:,i+1)=(eye(3,3)-KX(:,:,i+1)*wo(:,:,i+1)*wb(:,:,i+1))*Pfx(:,:,i+1);
  else

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

  end

end

disp(['***  FINISHED MIXED GAUSSIAN-LOGNORMAL KALMAN FILTER  ***'])

figure

subplot(2,1,1)
plot(vec(2:nobs),zob(vec(2:nobs)),'go','MarkerSize',6,'LineWidth',2)
hold on 
plot(z,'r','Linewidth',2)
plot(x_kf(3,:),'b','Linewidth',2)
plot(x_kf_g(3,:),'k','Linewidth',2)
xlabel 'time'
title('Z from KF and MXKF','FontSize',10,'FontWeight','bold')
set(gca,'FontWeight','bold','FontSize',10)
subplot(2,1,2)
plot(vec(2:nobs),xob(vec(2:nobs)),'go','MarkerSize',6,'LineWidth',2)
hold on 
plot(x,'r','Linewidth',2)
plot(x_kf(1,:),'b','Linewidth',2)
plot(x_kf_g(1,:),'k','Linewidth',2)
xlabel 'time'
title('X from KF and MXKF','FontSize',10,'FontWeight','bold')
legend('obs','truth','MXKF','KF','location','northwest','FontWeight','bold')
set(gca,'FontWeight','bold','FontSize',10)
figure

subplot(2,1,1)
plot(z'./x_kf(3,:),'k','Linewidth',2)
hold on 
plot(z'./x_kf_g(3,:),'k--','Linewidth',2)
title ('Z ANALYSIS ERROR','FontSize',10,'FontWeight','bold')
set(gca,'FontWeight','bold','FontSize',10)
ylabel 'z_a / z_t'
xlabel 'time'
subplot(2,1,2)
plot(x'-x_kf(1,:),'k','Linewidth',2)
hold on 
plot(x'-x_kf_g(1,:),'k--','Linewidth',2)
title('X ANALYSIS ERROR','FontSize',10,'FontWeight','bold')
legend('MXKF ERROR','KF ERROR','location','northwest')
ylabel 'x_a - x_t'
xlabel 'time'
set(gca,'FontWeight','bold','FontSize',10)