% Jyvaskylä Summer School 2014
% Aku Seppanen
% University of Eastern Finland, Kuopio
% Department of Applied Physics

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Example 5.1: Prior, likelihood and posterior %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear

f_true = [0 1]';

% likelihood
K = [0 1; .02 1];
var_n = .002;
Gamma_n = var_n*eye(2);
iGamma_n = inv(Gamma_n);
n = sqrt(var_n)*randn(2,1);
g_obs = K*f_true + n;

% Try different priors
% In each case, run the results several times (different realizations of n)
%priorcase = 1 % Narrow prior distribution, complementary to likelihood
%priorcase = 2 % wide prior distribution
%priorcase = 3 % Narrow prior distribution, not a very good complement to likelihood
priorcase = 4 % False prior information

if priorcase == 1 
    eta_f = [0.05 2.5]';
    Gamma_f = [0.05 0; 0 6];
elseif priorcase == 2 
    eta_f = [.2 1.2]';
    Gamma_f = [4 0; 0 4];
elseif priorcase == 3
    eta_f = [.5 1.1]';
    Gamma_f = [10 .5; .5 .05];
elseif priorcase == 4
    eta_f = [2 2.5]';
    Gamma_f = [0.05 0; 0 6];
end

iGamma_f = inv(Gamma_f);

ff1 = linspace(-4,4,50*8);
ff2 = linspace(0,5,50*5);
[f1,f2] = meshgrid(ff1,ff2);
pi_pr = zeros(size(f1,1),size(f1,2));
pi_likel = zeros(size(f1,1),size(f1,2));

for ii = 1:size(f1,1)
    for jj = 1:size(f1,2) 
       f = [f1(ii,jj);f2(ii,jj)];
       pi_pr(ii,jj) = exp(-5*(f-eta_f)'*iGamma_f*(f-eta_f));
       pi_likel(ii,jj) = exp(-5*(g_obs-K*f)'*iGamma_n*(g_obs-K*f));
    end
end
pi_post = pi_pr.*pi_likel;
    
f_ML = inv(K'*iGamma_n*K)*(K'*iGamma_n*g_obs);
f_MAP = inv(K'*iGamma_n*K + iGamma_f)*(K'*iGamma_n*g_obs + iGamma_f*eta_f);

figure(1),clf
contour(f1,f2,pi_pr,20), hold on
plot(f_true(1),f_true(2),'k+')
%axis([min(ff1),max(ff1),min(ff2),max(ff2)]), 
%h=surf(f1,f2,pi_pr), axis([min(ff1),max(ff1),min(ff2),max(ff2)]), 
axis equal
title('Prior density')
xlabel('f_1'),ylabel('f_2')
legend('Prior density','True f')

figure(2),clf
contour(f1,f2,pi_likel,30), hold on
%surf(f1,f2,pi_likel,'EdgeColor','none'), view(2), hold on
plot(f_true(1),f_true(2),'k+')
plot(f_ML(1),f_ML(2),'ko')
%axis([min([min(ff1),f_ML(1)]),max([max(ff1),f_ML(1)]),min([min(ff2),f_ML(2)]),max([max(ff2),f_ML(2)])])
axis equal
title('Likelihood density')
xlabel('f_1'),ylabel('f_2')
legend('Likelihood density','True f','f_{ML}')

figure(3),clf
contour(f1,f2,pi_post,30), hold on
plot(f_true(1),f_true(2),'k+')
%axis([min(ff1),max(ff1),min(ff2),max(ff2)]), 
plot(f_MAP(1),f_MAP(2),'ko')
axis equal
title('Posterior density')
xlabel('f_1'),ylabel('f_2')
legend('Posterior density','True f','f_{MAP}')


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Example 5.2: Metropolis-Hastings MCMC %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear

ff1 = linspace(-2,2,400);
ff2 = linspace(-1.5,2.5,400);
[f1,f2] = meshgrid(ff1,ff2);

pi_post = exp(-10*(f1.^2-f2).^2 - (f2-1/4).^4);

figure(1),clf
contour(f1,f2,pi_post,20), hold on
xlabel('f_1'),ylabel('f_2')
axis equal

% number of samples
Ns = 1000; 
%Ns = 2000; 
%Ns = 10000; 

f_init = [2 -1]'; % initial value

% Select the variance of the proposal density gamma
%gamma = .02; % chain moves too slowly
gamma = .5; % better choice!

f_samples = zeros(2,Ns);
f_samples(:,1) = f_init;

figure(2),clf
contour(f1,f2,pi_post,20), hold on
xlabel('f_1'),ylabel('f_2')
axis equal

for k = 2:Ns
  f_old = f_samples(:,k-1);
  f_proposed = gamma*randn(2,1) + f_old;
  pi_post_old = exp(-10*(f_old(1).^2-f_old(2)).^2 - (f_old(2)-1/4).^4);
  pi_post_proposed = exp(-10*(f_proposed(1).^2-f_proposed(2)).^2 - (f_proposed(2)-1/4).^4);
  alpha = min([1,pi_post_proposed/pi_post_old]);
  t = rand(1);
  if alpha >= t
      f_samples(:,k) = f_proposed;
  else
      f_samples(:,k) = f_old;
  end
end

% Visualization of the chain...
% A bit too slow...
SHOW_THE_CHAIN = 0;
if SHOW_THE_CHAIN
  for k = 2:Ns
    plot(f_samples(1,k),f_samples(2,k),'k+'),drawnow
  end
end

% Plot all samples at once
figure(3),clf
contour(f1,f2,pi_post,20), hold on
xlabel('f_1'),ylabel('f_2')
axis equal
plot(f_samples(1,:),f_samples(2,:),'k+')

burn_in_period_length = 100; %just a guess; should be analyzed...
f_CM = 1/Ns*sum(f_samples(:,burn_in_period_length+1:end),2);

plot(f_CM(1),f_CM(2),'rx','MarkerSize',20,'LineWidth',5)
legend('\pi_{post}(f)','samples','f_{CM}')



