%
% cycle the ensemble Kalman filter on the Henon attractor.
%
% Square root methods appear unstable.
%
% Originator: G.J. Hakim, University of Washington
%             hakim@atmos.washington.edu
%
% released under GNU General Public License version 3. http://www.gnu.org/licenses/gpl.html
%
% version control:
% $Date: 2009-02-26 11:59:58 -0800 (Thu, 26 Feb 2009) $
% $Revision: 48 $
% $Author: hakim $
% $Id: henon_ensda.m 48 2009-02-26 19:59:58Z hakim $

clear all
randn('state',0); % insure the same sequence of random numbers

update = 1;     % DA method (1= perturbed obs; 2= potter, 3=EnKF)
q = 1e-4*eye(2);% model error covariance
A = eye(2);     % initial analysis error covariance
R = eye(2)*1e-2;% observation error covariance
nassim = 50;   % number of assim cycles
nsteps = 1;     % do assimilation every nsteps steps
climo = 0;      % option for climo covariance (0 = flow dependent)
Nens = 10;      % number of ensemble members
inflate = 0.8; % covariance inflation 
H = eye(2);     % obs operator (eye(3) gives identity obs)
Nobs = 2;       % number of observations

par = [1.4 0.3]; % [a b] % Henon parameter values 

if update == 1
  disp('running perturbed observations...')
  inflate = 0;
elseif update == 2
  disp('running the potter algorithm...')
elseif update == 3
  disp('running the enkf algorithm...')
else
  error('wrong update code specified')
end
  
% get a state on the attractor
[x, y] = henon_get_ic(1,1,par);

% IC for truth
xt = x; yt = y;

% populate initial ensemble analysis by perturbing true state
pert = 1e-1*randn(1,Nens); 
xa = xt + pert; ya = yt + pert;

% check to see if we're using a fixed B
if climo == 1
  load henon_climo_B; Bc = B; 
end

disp(['Cycling ON the attractor...']);

for k=1:nassim

  % step to the next assimilation time
  xb = xa; yb = ya;
  for t = 1:1:nsteps
    % advance truth with the full nonlinear model
    [xt yt] = henon(xt,yt,par);

    % advance background ensemble with the full nonlinear model
	 for n = 1:1:Nens
      [xb(n), yb(n)] = henon(xb(n),yb(n),par);
    end
  end  

  % new observations from noise about truth
  Y = [xt; yt] + diag(diag(randn(2,1))*sqrtm(R));
  
  % new background error covariance from ensemble estimate
  if climo == 0
	 xp = xb - mean(xb); yp = yb - mean(yb);
	 if inflate ~= 0 % square-root filter
		Xbp = inflate*[xp; yp];
	 else % perturbed obs
		Xbp = [xp; yp];
		Xb = [xb; yb];
	 end
	 B = (Xbp*Xbp')/(Nens - 1) + q*eye(2);
  else
	 B = Bc;
  end
  
  xbm = [mean(xb) mean(yb)]';  % ensemble mean
   
  % update step for ensemble mean (xam), perturbations from the mean (Xap), 
  % and the full state Xa = Xam + Xap
  if update == 1
	 Xa = perturbed_obs(Xb,B,Y,R); 
	 xam = mean(Xa,2);
	 [tmp Xam] = meshgrid(ones(Nens,1),xam);
	 Xap = Xa - Xam; 
  elseif update == 2
	 [xam,Xap] = potter(xbm,Xbp,Y,H,R); 
	 [tmp Xam] = meshgrid(ones(Nens,1),xam);
	 Xa = Xam + Xap;
  elseif update == 3
	 loc = ones(Nobs,2); % this does no localization; size: Nobs x dof
	 [xam,Xap] = enkf_update(xbm,Xbp,Y,H,R,loc);	 
	 [tmp Xam] = meshgrid(ones(Nens,1),xam);
	 Xa = Xam + Xap;
  end

  A = (Xap*Xap')/(Nens - 1);

  disp(['trace of B and A: ' num2str(trace(B),2) '    ' num2str(trace(A),2)])
  
  xa = Xa(1,:); ya = Xa(2,:);

  % update step
  if 1 == 0
	 for n = 1:1:Nens
		Xb = [xb(n); yb(n)];
		if inflate == 0 % perturbed obs for each member
		  Yn = Y + diag(diag(randn(2,1))*sqrtm(R));
		else % square root filter
		  Yn = Y;
		end
		K = B*inv(B + R);
		Xa = Xb + K*(Yn - Xb);
		xa(n) = Xa(1); ya(n) = Xa(2);
	 end
	 A = (eye(2) - K)*B;
  end
  
  xbm = mean(xb); ybm = mean(yb);
  xam = mean(xa); yam = mean(ya);
  
  %error statistics for the ensemble mean
  xbe(k) = xbm - xt; ybe(k) = ybm - yt;
  xae(k) = xam - xt; yae(k) = yam - yt;
  xye(k) = Y(1) - xt; yye(k) = Y(2) - yt;
  xaev(k) = A(1,1); yaev(k) = A(2,2);
  xbev(k) = B(1,1); ybev(k) = B(2,2);

  % check for filter divergence & blowup
  if abs(xae(k)) > 1 & A(1,1) <1e-5 
	 error('filter divergence!')
  end
  if abs(xae(k)) > 10
	 error('filter blow up!')
  end

end % end of assimilation loop

xbe_var = var(xbe);
ybe_var = var(ybe);
xae_var = var(xae);
yae_var = var(yae);

xbe_rms = sqrt(xbe_var/nassim);

if 1 == 1
  
% plot error statistics
figure(1); clf
plot(abs(xbe),'b-'); hold on
plot(abs(xae),'r-')
c = ['mean background: ' num2str(mean(abs(xbe(2:end))),3) ' +/- ' num2str(std(abs(xbe(2:end))),3) '    mean analysis:   ' num2str(mean(abs(xae(2:end))),3) ' +/- ' num2str(std(abs(xae(2:end))),3)];
yl = get(gca,'ylim');
yoff = yl(1) - yl(2)*.12;
text(0,yoff,c);
xlabel('Assimilation Step','fontweight','bold','fontsize',12)
ylabel('Absolute Error','fontweight','bold','fontsize',12)
title('Ensemble Kalman Filter Absolute Error (x)','fontweight','bold','fontsize',12)
set(gca,'linewidth',2,'fontweight','bold','fontsize',10)
legend('background','analysis')
if climo == 0
%  print -dpng -r80 henon_enkf_abs_error.png
else
%  print -dpng -r80 henon_enkf_abs_error_climo_B.png
end

figure(2); clf
plot(xbev(2:end),'b-'); hold on
plot(xaev(2:end),'r-')
c = ['mean background: ' num2str(mean(xbev(2:end)),3) ' +/- ' num2str(std(xbev(2:end)),3) '    mean analysis:   ' num2str(mean(xaev(2:end)),3) ' +/- ' num2str(std(xaev(2:end)),3)];
yl = get(gca,'ylim');
yoff = yl(1) - yl(2)*.12;
text(0,yoff,c);
xlabel('Assimilation Step','fontweight','bold','fontsize',12)
ylabel('Error Variance','fontweight','bold','fontsize',12)
title('Ensemble Kalman Filter Error Variance (x)','fontweight','bold','fontsize',12)
set(gca,'linewidth',2,'fontweight','bold','fontsize',10)
legend('background','analysis')
if climo == 0
%  print -dpng -r80 henon_enkf_error_var.png
else
%  print -dpng -r80 henon_enkf_error_var_climo_B.png  
end


end
